Attended PR comments

Author: aponcedeleonch

src/codegate/config.py

@@ -57,7 +57,9 @@ class Config:
     force_certs: bool = False
 
     max_fim_hash_lifetime: int = 60 * 5  # Time in seconds. Default is 5 minutes.
-    persona_threshold = 0.75  # Min value is 0 (max similarity), max value is 2 (orthogonal)
+    # Min value is 0 (max similarity), max value is 2 (orthogonal)
+    # The value 0.75 was found through experimentation. See /tests/muxing/test_semantic_router.py
+    persona_threshold = 0.75
 
     # Provider URLs with defaults
     provider_urls: Dict[str, str] = field(default_factory=lambda: DEFAULT_PROVIDER_URLS.copy())

src/codegate/db/connection.py

@@ -560,7 +560,7 @@ async def add_persona(self, persona: PersonaEmbedding) -> None:
         )
 
         try:
-            # For Pydantic we conver the numpy array to a string when serializing.
+            # For Pydantic we convert the numpy array to string when serializing with .model_dumpy()
             # We need to convert it back to a numpy array before inserting it into the DB.
             persona_dict = persona.model_dump()
             persona_dict["description_embedding"] = persona.description_embedding
@@ -615,17 +615,19 @@ async def _exec_select_conditions_to_pydantic(
                     raise e
                 return None
 
-    async def _exec_vec_db_query(
-        self, sql_command: str, conditions: dict
-    ) -> Optional[CursorResult]:
+    async def _exec_vec_db_query_to_pydantic(
+        self, sql_command: str, conditions: dict, model_type: Type[BaseModel]
+    ) -> List[BaseModel]:
         """
         Execute a query on the vector database. This is a separate connection to the SQLite
         database that has the vector extension loaded.
         """
         conn = self._get_vec_db_connection()
+        conn.row_factory = sqlite3.Row
         cursor = conn.cursor()
-        cursor.execute(sql_command, conditions)
-        return cursor
+        results = [model_type(**row) for row in cursor.execute(sql_command, conditions)]
+        conn.close()
+        return results
 
     async def get_prompts_with_output(self, workpace_id: str) -> List[GetPromptWithOutputsRow]:
         sql = text(
@@ -985,14 +987,10 @@ async def get_distance_to_persona(
             WHERE id = :id
         """
         conditions = {"id": persona_id, "query_embedding": query_embedding}
-        persona_distance_cursor = await self._exec_vec_db_query(sql, conditions)
-        persona_distance_raw = persona_distance_cursor.fetchone()
-        return PersonaDistance(
-            id=persona_distance_raw[0],
-            name=persona_distance_raw[1],
-            description=persona_distance_raw[2],
-            distance=persona_distance_raw[3],
+        persona_distance = await self._exec_vec_db_query_to_pydantic(
+            sql, conditions, PersonaDistance
         )
+        return persona_distance[0]
 
 
 def init_db_sync(db_path: Optional[str] = None):

src/codegate/db/models.py

@@ -243,8 +243,6 @@ class MuxRule(BaseModel):
     updated_at: Optional[datetime.datetime] = None
 
 
-# Pydantic doesn't support numpy arrays out of the box. Defining a custom type
-# Reference: https://github.com/pydantic/pydantic/issues/7017
 def nd_array_custom_before_validator(x):
     # custome before validation logic
     return x
@@ -255,6 +253,11 @@ def nd_array_custom_serializer(x):
     return str(x)
 
 
+# Pydantic doesn't support numpy arrays out of the box hence we need to construct a custom type.
+# There are 2 things necessary for a Pydantic custom type: Validator and Serializer
+# The lines below build our custom type
+# Docs: https://docs.pydantic.dev/latest/concepts/types/#adding-validation-and-serialization
+# Open Pydantic issue for npy support: https://github.com/pydantic/pydantic/issues/7017
 NdArray = Annotated[
     np.ndarray,
     BeforeValidator(nd_array_custom_before_validator),
@@ -263,17 +266,33 @@ def nd_array_custom_serializer(x):
 
 
 class Persona(BaseModel):
+    """
+    Represents a persona object.
+    """
+
     id: str
     name: str
     description: str
 
 
 class PersonaEmbedding(Persona):
-    description_embedding: NdArray  # sqlite-vec will handle numpy arrays directly
+    """
+    Represents a persona object with an embedding.
+    """
+
+    description_embedding: NdArray
 
     # Part of the workaround to allow numpy arrays in pydantic models
     model_config = ConfigDict(arbitrary_types_allowed=True)
 
 
 class PersonaDistance(Persona):
+    """
+    Result of an SQL query to get the distance between the query and the persona description.
+
+    A vector similarity search is performed to get the distance. Distance values ranges [0, 2].
+    0 means the vectors are identical, 2 means they are orthogonal.
+    See [sqlite docs](https://alexgarcia.xyz/sqlite-vec/api-reference.html#vec_distance_cosine)
+    """
+
     distance: float

src/codegate/muxing/semantic_router.py

@@ -13,6 +13,17 @@
 logger = structlog.get_logger("codegate")
 
 
+REMOVE_URLS = re.compile(r"https?://\S+|www\.\S+")
+REMOVE_EMAILS = re.compile(r"\S+@\S+")
+REMOVE_CODE_BLOCKS = re.compile(r"```[\s\S]*?```")
+REMOVE_INLINE_CODE = re.compile(r"`[^`]*`")
+REMOVE_HTML_TAGS = re.compile(r"<[^>]+>")
+REMOVE_PUNCTUATION = re.compile(r"[^\w\s\']")
+NORMALIZE_WHITESPACE = re.compile(r"\s+")
+NORMALIZE_DECIMAL_NUMBERS = re.compile(r"\b\d+\.\d+\b")
+NORMALIZE_INTEGER_NUMBERS = re.compile(r"\b\d+\b")
+
+
 class PersonaDoesNotExistError(Exception):
     pass
 
@@ -54,27 +65,27 @@ def _clean_text_for_embedding(self, text: str) -> str:
         text = "".join([c for c in text if not unicodedata.combining(c)])
 
         # Remove URLs
-        text = re.sub(r"https?://\S+|www\.\S+", " ", text)
+        text = REMOVE_URLS.sub(" ", text)
 
         # Remove email addresses
-        text = re.sub(r"\S+@\S+", " ", text)
+        text = REMOVE_EMAILS.sub(" ", text)
 
         # Remove code block markers and other markdown/code syntax
-        text = re.sub(r"```[\s\S]*?```", " ", text)  # Code blocks
-        text = re.sub(r"`[^`]*`", " ", text)  # Inline code
+        text = REMOVE_CODE_BLOCKS.sub(" ", text)
+        text = REMOVE_INLINE_CODE.sub(" ", text)
 
         # Remove HTML/XML tags
-        text = re.sub(r"<[^>]+>", " ", text)
+        text = REMOVE_HTML_TAGS.sub(" ", text)
 
         # Normalize numbers (replace with placeholder)
-        text = re.sub(r"\b\d+\.\d+\b", " NUM ", text)  # Decimal numbers
-        text = re.sub(r"\b\d+\b", " NUM ", text)  # Integer numbers
+        text = NORMALIZE_DECIMAL_NUMBERS.sub(" NUM ", text)  # Decimal numbers
+        text = NORMALIZE_INTEGER_NUMBERS.sub(" NUM ", text)  # Integer numbers
 
         # Replace punctuation with spaces (keeping apostrophes for contractions)
-        text = re.sub(r"[^\w\s\']", " ", text)
+        text = REMOVE_PUNCTUATION.sub(" ", text)
 
         # Normalize whitespace (replace multiple spaces with a single space)
-        text = re.sub(r"\s+", " ", text)
+        text = NORMALIZE_WHITESPACE.sub(" ", text)
 
         # Convert to lowercase and strip
         text = text.strip()
@@ -91,7 +102,7 @@ async def _embed_text(self, text: str) -> np.ndarray:
             self._embeddings_model, [cleaned_text], n_gpu_layers=self._n_gpu
         )
         # Use only the first entry in the list and make sure we have the appropriate type
-        logger.debug("Text embedded in semantic routing", text=cleaned_text[:100])
+        logger.debug("Text embedded in semantic routing", text=cleaned_text[:50])
         return np.array(embed_list[0], dtype=np.float32)
 
     async def add_persona(self, persona_name: str, persona_desc: str) -> None: