feat: swap embedding backend to model2vec and fix search arg-swap

Author: nicosammito

pyproject.toml

@@ -14,6 +14,6 @@ dependencies = [
     "litellm==1.83.7",
     "pydantic==2.12.5",
     "qdrant-client>=1.16.1",
-    "sentence-transformers>=5.1.2",
+    "model2vec>=0.3.0",
     "tucana==0.0.72",
 ]

src/model.py

@@ -1,7 +1,8 @@
-from sentence_transformers import SentenceTransformer
+from model2vec import StaticModel
 
-MODEL_NAME = 'all-MiniLM-L6-v2'
+MODEL_NAME = 'minishlab/potion-base-8M'
+EMBEDDING_DIM = 256
 
 
-def load_vector_model() -> SentenceTransformer:
-    return SentenceTransformer(MODEL_NAME)
+def load_vector_model() -> StaticModel:
+    return StaticModel.from_pretrained(MODEL_NAME)

src/store/few_shots_store.py

@@ -2,8 +2,8 @@
 from pathlib import Path
 from typing import Any, List
 
+from model2vec import StaticModel
 from qdrant_client import QdrantClient
-from sentence_transformers import SentenceTransformer
 
 from src.logger import get_logger
 from src.schema.few_shot_schema import FewShot
@@ -13,7 +13,7 @@
 
 
 class FewShotsStore(Store):
-    def __init__(self, memory_client: QdrantClient, vector_model: SentenceTransformer):
+    def __init__(self, memory_client: QdrantClient, vector_model: StaticModel):
         super().__init__(
             memory_client,
             vector_model,
@@ -50,7 +50,7 @@ def validate(self, payload: Any) -> FewShot:
         return FewShot.model_validate(payload)
 
     def search(self, group_identifier: str, prompt: str, limit=5) -> List[FewShot]:
-        return super().search(prompt, group_identifier, limit)
+        return super().search(group_identifier, prompt, limit)
 
     def find(self, group_identifier: str, identifier: str) -> FewShot | None:
         return super().find(group_identifier, identifier)

src/store/flow_type_store.py

@@ -1,8 +1,8 @@
 import re
 from typing import List, Any
 
+from model2vec import StaticModel
 from qdrant_client import QdrantClient
-from sentence_transformers import SentenceTransformer
 
 from src.schema.data_type_schema import DataType
 from src.schema.flow_type_schema import FlowType
@@ -81,7 +81,7 @@ def replacement(match):
 
 
 class FlowTypeStore(Store):
-    def __init__(self, memory_client: QdrantClient, vector_model: SentenceTransformer):
+    def __init__(self, memory_client: QdrantClient, vector_model: StaticModel):
         super().__init__(
             memory_client,
             vector_model,
@@ -105,7 +105,7 @@ def validate(self, payload: Any) -> FlowType:
         return FlowType.model_validate(payload)
 
     def search(self, group_identifier: str, prompt: str, limit=5) -> List[FlowType]:
-        return super().search(prompt, group_identifier, limit)
+        return super().search(group_identifier, prompt, limit)
 
     def find(self, group_identifier: str, identifier: str) -> FlowType | None:
         return super().find(group_identifier, identifier)

src/store/function_store.py

@@ -1,8 +1,8 @@
 import re
 from typing import List, Any
 
+from model2vec import StaticModel
 from qdrant_client import QdrantClient
-from sentence_transformers import SentenceTransformer
 
 from src.schema.data_type_schema import DataType
 from src.schema.function_schema import FunctionDefinition
@@ -81,7 +81,7 @@ def replacement(match):
 
 
 class FunctionStore(Store):
-    def __init__(self, memory_client: QdrantClient, vector_model: SentenceTransformer):
+    def __init__(self, memory_client: QdrantClient, vector_model: StaticModel):
         super().__init__(
             memory_client,
             vector_model,
@@ -105,7 +105,7 @@ def validate(self, payload: Any) -> FunctionDefinition:
         return FunctionDefinition.model_validate(payload)
 
     def search(self, group_identifier: str, prompt: str, limit=5) -> List[FunctionDefinition]:
-        return super().search(prompt, group_identifier, limit)
+        return super().search(group_identifier, prompt, limit)
 
     def find(self, group_identifier: str, identifier: str) -> FunctionDefinition | None:
         return super().find(group_identifier, identifier)

src/store/vector_store.py

@@ -5,7 +5,9 @@
 from apscheduler.schedulers.background import BackgroundScheduler
 from qdrant_client import QdrantClient
 from qdrant_client.models import Distance, VectorParams, PointStruct, MatchAny, Filter, FieldCondition, MatchValue
-from sentence_transformers import SentenceTransformer
+from model2vec import StaticModel
+
+from src.model import EMBEDDING_DIM
 
 from src.logger import get_logger
 
@@ -16,7 +18,7 @@ class Store(ABC):
     def __init__(
             self,
             client: QdrantClient,
-            model: SentenceTransformer,
+            model: StaticModel,
             collection_name: str,
             payload_identifier: str,
             group_identifier: str,
@@ -39,7 +41,7 @@ def _setup_collection(self):
             self.client.create_collection(
                 collection_name=self.collection_name,
                 vectors_config=VectorParams(
-                    size=384,
+                    size=EMBEDDING_DIM,
                     distance=Distance.COSINE
                 ),
             )