Merge pull request #8 from code0-tech/feat/#2

nicosammito · web-flow · commit 77b1ea7b4252 · 2026-06-21T18:47:31.000+02:00
Alpine/musl compatibility
diff --git a/pyproject.toml b/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",
 ]
diff --git a/src/endpoint/generation/generate_endpoint.py b/src/endpoint/generation/generate_endpoint.py
@@ -71,28 +71,6 @@ def Prompt(self, request: pb2.PromptRequest, context) -> pb2.FlowResponse:
                 details=f"The specified model_identifier '{request.model_identifier}' does not exist. Please provide a valid model_identifier for flow generation."
             )
 
-        if len(
-                self.function_store.get_all(
-                    group_identifier=str(request.project_id)
-                )
-        ) <= 0 and (len(request.functions) <= 0):
-            log.warning(f"[Prompt] Rejected — no functions for project={request.project_id}")
-            context.abort(
-                code=grpc.StatusCode.ABORTED,
-                details="No functions found for the given project_id. Please add functions before requesting a prompt generation."
-            )
-
-        if len(
-                self.flow_type_store.get_all(
-                    group_identifier=str(request.project_id)
-                )
-        ) <= 0 and (len(request.flow_types) <= 0):
-            log.warning(f"[Prompt] Rejected — no flow types for project={request.project_id}")
-            context.abort(
-                code=grpc.StatusCode.ABORTED,
-                details="No flow types found for the given project_id. Please add flow_types before requesting a prompt generation."
-            )
-
         functions = [
             map_to_function_schema(fn)
             for fn in request.functions
@@ -122,6 +100,30 @@ def Prompt(self, request: pb2.PromptRequest, context) -> pb2.FlowResponse:
                 data_types=data_types
             )
 
+        stored_functions_count = len(self.function_store.get_all(group_identifier=str(request.project_id)))
+        stored_flow_types_count = len(self.flow_type_store.get_all(group_identifier=str(request.project_id)))
+        stored_few_shots_count = len(self.few_shots_store.get_all(group_identifier="global"))
+        log.info(
+            f"[Prompt] Store state — project={request.project_id} "
+            f"functions={stored_functions_count} "
+            f"flow_types={stored_flow_types_count} "
+            f"few_shots={stored_few_shots_count}"
+        )
+
+        if stored_functions_count <= 0:
+            log.warning(f"[Prompt] Rejected — no functions for project={request.project_id}")
+            context.abort(
+                code=grpc.StatusCode.ABORTED,
+                details="No functions found for the given project_id. Please add functions before requesting a prompt generation."
+            )
+
+        if stored_flow_types_count <= 0:
+            log.warning(f"[Prompt] Rejected — no flow types for project={request.project_id}")
+            context.abort(
+                code=grpc.StatusCode.ABORTED,
+                details="No flow types found for the given project_id. Please add flow_types before requesting a prompt generation."
+            )
+
         prompt_functions = self.function_store.search(
             group_identifier=str(request.project_id),
             prompt=request.prompt,
@@ -272,28 +274,6 @@ def Flow(self, request: pb2.FlowRequest, context) -> pb2.FlowResponse:
                 details=f"The specified model_identifier '{request.model_identifier}' does not exist. Please provide a valid model_identifier for flow generation."
             )
 
-        if len(
-                self.function_store.get_all(
-                    group_identifier=str(request.project_id)
-                )
-        ) <= 0 and (len(request.functions) <= 0):
-            log.warning(f"[Flow] Rejected — no functions for project={request.project_id}")
-            context.abort(
-                code=grpc.StatusCode.ABORTED,
-                details="No functions found for the given project_id. Please add functions before requesting a prompt generation."
-            )
-
-        if len(
-                self.flow_type_store.get_all(
-                    group_identifier=str(request.project_id)
-                )
-        ) <= 0 and (len(request.flow_types) <= 0):
-            log.warning(f"[Flow] Rejected — no flow types for project={request.project_id}")
-            context.abort(
-                code=grpc.StatusCode.ABORTED,
-                details="No flow types found for the given project_id. Please add flow_types before requesting a prompt generation."
-            )
-
         functions = [
             map_to_function_schema(fn)
             for fn in request.functions
@@ -323,6 +303,30 @@ def Flow(self, request: pb2.FlowRequest, context) -> pb2.FlowResponse:
                 data_types=data_types
             )
 
+        stored_functions_count = len(self.function_store.get_all(group_identifier=str(request.project_id)))
+        stored_flow_types_count = len(self.flow_type_store.get_all(group_identifier=str(request.project_id)))
+        stored_few_shots_count = len(self.few_shots_store.get_all(group_identifier="global"))
+        log.info(
+            f"[Flow] Store state — project={request.project_id} "
+            f"functions={stored_functions_count} "
+            f"flow_types={stored_flow_types_count} "
+            f"few_shots={stored_few_shots_count}"
+        )
+
+        if stored_functions_count <= 0:
+            log.warning(f"[Flow] Rejected — no functions for project={request.project_id}")
+            context.abort(
+                code=grpc.StatusCode.ABORTED,
+                details="No functions found for the given project_id. Please add functions before requesting a prompt generation."
+            )
+
+        if stored_flow_types_count <= 0:
+            log.warning(f"[Flow] Rejected — no flow types for project={request.project_id}")
+            context.abort(
+                code=grpc.StatusCode.ABORTED,
+                details="No flow types found for the given project_id. Please add flow_types before requesting a prompt generation."
+            )
+
         prompt_functions = self.function_store.search(
             group_identifier=str(request.project_id),
             prompt=request.prompt,
diff --git a/src/model.py b/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)
diff --git a/src/store/few_shots_store.py b/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,13 +13,14 @@
 
 
 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,
             'few_shots',
             "identifier",
-            "identifier"
+            "identifier",
+            time_to_live=None,
         )
         self._load_models_from_json()
 
@@ -50,7 +51,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)
diff --git a/src/store/flow_type_store.py b/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)
diff --git a/src/store/function_store.py b/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)
diff --git a/src/store/vector_store.py b/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,11 +18,11 @@ class Store(ABC):
     def __init__(
             self,
             client: QdrantClient,
-            model: SentenceTransformer,
+            model: StaticModel,
             collection_name: str,
             payload_identifier: str,
             group_identifier: str,
-            time_to_live: int = 60 * 5
+            time_to_live: int | None = 60 * 5
 
     ):
         self.client = client
@@ -32,14 +34,15 @@ def __init__(
         self._id_counter = 0
         self._scheduler = None
         self._setup_collection()
-        self._start_garbage_collector()
+        if self.time_to_live is not None:
+            self._start_garbage_collector()
 
     def _setup_collection(self):
         if not self.client.collection_exists(collection_name=self.collection_name):
             self.client.create_collection(
                 collection_name=self.collection_name,
                 vectors_config=VectorParams(
-                    size=384,
+                    size=EMBEDDING_DIM,
                     distance=Distance.COSINE
                 ),
             )
@@ -81,13 +84,21 @@ def _garbage_collect_all(self):
             log.error(f"[GC] Error in collection '{self.collection_name}': {e}")
 
     def _reset_time_to_live(self, group_identifier: str):
-        items = self.get_all(group_identifier)
-        current_time = time.time()
+        if self.time_to_live is None:
+            return
 
-        for item in items:
-            item_dict = item.model_dump() if hasattr(item, 'model_dump') else item
-            item_dict["created_at"] = current_time
-            self.insert(group_identifier, item_dict)
+        self.client.set_payload(
+            collection_name=self.collection_name,
+            payload={"created_at": time.time()},
+            points=Filter(
+                must=[
+                    FieldCondition(
+                        key=self.group_identifier,
+                        match=MatchValue(value=group_identifier)
+                    )
+                ]
+            ),
+        )
 
     @abstractmethod
     def validate(self, payload: Any) -> Any:
@@ -101,8 +112,11 @@ def insert(self, group_identifier: str, payload: Any) -> None:
         if "created_at" not in payload:
             payload["created_at"] = time.time()
 
-        payload_str = str(payload)
-        vector = self.model.encode(payload_str).tolist()
+        semantic_payload = {
+            k: v for k, v in payload.items()
+            if k not in ("created_at", self.group_identifier)
+        }
+        vector = self.model.encode(str(semantic_payload)).tolist()
 
         payload[self.group_identifier] = group_identifier
 
diff --git a/uv.lock b/uv.lock

Original file line number	Diff line number	Diff line change
`@@ -14,6 +14,6 @@ dependencies = [`
`14`	`14`	`"litellm==1.83.7",`
`15`	`15`	`"pydantic==2.12.5",`
`16`	`16`	`"qdrant-client>=1.16.1",`
`17`		`- "sentence-transformers>=5.1.2",`
	`17`	`+ "model2vec>=0.3.0",`
`18`	`18`	`"tucana==0.0.72",`
`19`	`19`	`]`