vector service update

service/vector_store_service.py

@@ -1,11 +1,34 @@
 import numpy as np
 
 class VectorStoreService:
-    def __init__(self, embeddings, documents):
-        self.embeddings = np.array(embeddings)
+    """
+    Simple in-memory vector store.
+    Stores embeddings and corresponding documents.
+    """
+
+    def __init__(self, embeddings: list[list[float]], documents: list[str]):
+        assert len(embeddings) == len(documents), "Embeddings and documents must match in length."
+        self.embeddings = np.array(embeddings, dtype=np.float32)
         self.documents = documents
 
-    def search(self, query_embedding, k=3):
-        scores = query_embedding @ self.embeddings.T
-        top_k = scores.argsort()[-k:][::-1]
-        return [self.documents[i] for i in top_k]
+    def search_with_scores(self, query_vec: list[float], k: int = 3):
+        """
+        Returns top-k documents along with cosine similarity scores.
+        :param query_vec: single query vector
+        :param k: number of top results
+        :return: list of tuples (document, similarity_score)
+        """
+        query_vec = np.array(query_vec, dtype=np.float32)
+
+        # Normalize vectors for cosine similarity
+        query_norm = query_vec / np.linalg.norm(query_vec)
+        doc_norms = self.embeddings / np.linalg.norm(self.embeddings, axis=1, keepdims=True)
+
+        # Cosine similarity
+        similarities = doc_norms.dot(query_norm)
+
+        # Get top-k indices
+        top_indices = np.argsort(-similarities)[:k]
+
+        # Return documents with similarity scores
+        return [(self.documents[i], float(similarities[i])) for i in top_indices]