Validate persona description is sufficiently different (#1225)

Closes: #1218

Check if the description for a new persona is different enough from
the existing personas descriptions. This is done to correctly
differentiate between personas

Author: aponcedeleonch

src/codegate/config.py

@@ -57,9 +57,14 @@ class Config:
     force_certs: bool = False
 
     max_fim_hash_lifetime: int = 60 * 5  # Time in seconds. Default is 5 minutes.
+
     # 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
+    # It's the threshold value to determine if a query matches a persona.
     persona_threshold = 0.75
+    # The value 0.3 was found through experimentation. See /tests/muxing/test_semantic_router.py
+    # It's the threshold value to determine if a persona description is similar to existing personas
+    persona_diff_desc_threshold = 0.3
 
     # Provider URLs with defaults
     provider_urls: Dict[str, str] = field(default_factory=lambda: DEFAULT_PROVIDER_URLS.copy())

src/codegate/db/connection.py

@@ -1004,6 +1004,26 @@ async def get_persona_by_name(self, persona_name: str) -> Optional[Persona]:
         )
         return personas[0] if personas else None
 
+    async def get_distance_to_existing_personas(
+        self, query_embedding: np.ndarray
+    ) -> List[PersonaDistance]:
+        """
+        Get the distance between a persona and a query embedding.
+        """
+        sql = """
+            SELECT
+                id,
+                name,
+                description,
+                vec_distance_cosine(description_embedding, :query_embedding) as distance
+            FROM personas
+        """
+        conditions = {"query_embedding": query_embedding}
+        persona_distances = await self._exec_vec_db_query_to_pydantic(
+            sql, conditions, PersonaDistance
+        )
+        return persona_distances
+
     async def get_distance_to_persona(
         self, persona_id: str, query_embedding: np.ndarray
     ) -> PersonaDistance:

src/codegate/db/models.py

@@ -245,6 +245,8 @@ class MuxRule(BaseModel):
 
 def nd_array_custom_before_validator(x):
     # custome before validation logic
+    if isinstance(x, bytes):
+        return np.frombuffer(x, dtype=np.float32)
     return x
 
 

src/codegate/muxing/semantic_router.py

@@ -28,6 +28,10 @@ class PersonaDoesNotExistError(Exception):
     pass
 
 
+class PersonaSimilarDescriptionError(Exception):
+    pass
+
+
 class SemanticRouter:
 
     def __init__(self):
@@ -36,6 +40,7 @@ def __init__(self):
         self._embeddings_model = f"{conf.model_base_path}/{conf.embedding_model}"
         self._n_gpu = conf.chat_model_n_gpu_layers
         self._persona_threshold = conf.persona_threshold
+        self._persona_diff_desc_threshold = conf.persona_diff_desc_threshold
         self._db_recorder = DbRecorder()
         self._db_reader = DbReader()
 
@@ -105,12 +110,38 @@ async def _embed_text(self, text: str) -> np.ndarray:
         logger.debug("Text embedded in semantic routing", text=cleaned_text[:50])
         return np.array(embed_list[0], dtype=np.float32)
 
+    async def _is_persona_description_diff(self, emb_persona_desc: np.ndarray) -> bool:
+        """
+        Check if the persona description is different enough from existing personas.
+        """
+        # The distance calculation is done in the database
+        persona_distances = await self._db_reader.get_distance_to_existing_personas(
+            emb_persona_desc
+        )
+        if not persona_distances:
+            return True
+
+        for persona_distance in persona_distances:
+            logger.info(
+                f"Persona description distance to {persona_distance.name}",
+                distance=persona_distance.distance,
+            )
+            # If the distance is less than the threshold, the persona description is too similar
+            if persona_distance.distance < self._persona_diff_desc_threshold:
+                return False
+        return True
+
     async def add_persona(self, persona_name: str, persona_desc: str) -> None:
         """
         Add a new persona to the database. The persona description is embedded
         and stored in the database.
         """
         emb_persona_desc = await self._embed_text(persona_desc)
+        if not await self._is_persona_description_diff(emb_persona_desc):
+            raise PersonaSimilarDescriptionError(
+                "The persona description is too similar to existing personas."
+            )
+
         new_persona = db_models.PersonaEmbedding(
             id=str(uuid.uuid4()),
             name=persona_name,