refactor: use pgvector ORM methods instead of raw SQL

Author: Igboke

app/repositories/cached_query_repository.py

@@ -16,36 +16,20 @@ def __init__(self, session: SyncSession):
     def find_similar_query(self, query_embedding: List[float]) -> CachedQuery | None:
         """
         Finds a cached query if its embedding is within the similarity threshold.
-        This is done in a single, efficient query.
+        This is done in a single, efficient query using the ORM.
         """
 
         distance_threshold = 1 - SIMILARITY_THRESHOLD
 
-        stmt = text("""
-            SELECT id
-            FROM cached_queries
-            WHERE (question_embedding <=> CAST(:query_embedding AS vector)) < :distance_threshold
-            ORDER BY question_embedding <=> CAST(:query_embedding AS vector)
-            LIMIT 1
-        """)
+        stmt = select(CachedQuery).options(
+            selectinload(CachedQuery.source_chunks)
+        ).where(
+            CachedQuery.question_embedding.cosine_distance(query_embedding) < distance_threshold
+        ).order_by(
+            CachedQuery.question_embedding.cosine_distance(query_embedding)
+        ).limit(1)
 
-        result = self.session.execute(
-            stmt,
-            {
-                "query_embedding": query_embedding,
-                "distance_threshold": distance_threshold
-            }
-        ).scalar_one_or_none()
-
-        if result:
-            similar_query_id = result
-            final_stmt = select(CachedQuery).options(
-                selectinload(CachedQuery.source_chunks)
-            ).where(CachedQuery.id == similar_query_id)
-            
-            return self.session.execute(final_stmt).scalar_one_or_none()
-
-        return None
+        return self.session.execute(stmt).scalar_one_or_none()
 
     def save_query(self, question_text: str, query_embedding: list[float], answer: str, source_chunks: list[Chunk]):
         """Saves a new query and links it to its source chunks."""

app/repositories/chunk_repository.py

@@ -18,32 +18,10 @@ def delete_by_document_id_sync(self, document_id: str):
     def find_relevant_chunks_sync(self, query_embedding: list[float], top_k: int = 5) -> list[Chunk]:
         """
         Finds the most relevant chunks using vector similarity.
-
-        Strategy: Use raw SQL for fast vector search to get IDs only,
-        then fetch via ORM to get properly-attached objects that can be used in relationships.
         """
- 
-        stmt = text("""
-            SELECT id
-            FROM chunks
-            ORDER BY embedding <=> CAST(:query_embedding AS vector)
-            LIMIT :top_k
-        """)
-
-        result = self.session.execute(
-            stmt,
-            {"query_embedding": query_embedding.tolist() if hasattr(query_embedding, 'tolist') else query_embedding, "top_k": top_k}
-        )
-
-        chunk_ids = [row[0] for row in result]
-
-        if not chunk_ids:
-            return []
-
-        chunks = self.session.execute(
-            select(Chunk).where(Chunk.id.in_(chunk_ids))
-        ).scalars().all()
+        stmt = select(Chunk).order_by(
+            Chunk.embedding.cosine_distance(query_embedding)
+        ).limit(top_k)
 
-        chunks_dict = {chunk.id: chunk for chunk in chunks}
-        return [chunks_dict[chunk_id] for chunk_id in chunk_ids if chunk_id in chunks_dict]
+        return self.session.execute(stmt).scalars().all()