Refactoring to allow both list of lists, and single lists, for embeddings, due to provider differences

Author: tomsmoker

backend/src/app/services/vector_db/base.py

@@ -82,11 +82,16 @@ async def ensure_collection_exists(self) -> None:
     async def get_embeddings(
         self, texts: Union[str, List[str]]
     ) -> List[List[float]]:
-        """Get embeddings for the given text(s) using the LLM service."""
+        """Get embeddings for the given text(s) using the embedding service."""
         if isinstance(texts, str):
             texts = [texts]
         return await self.embedding_service.get_embeddings(texts)
 
+    async def get_single_embedding(self, text: str) -> List[float]:
+        """Get a single embedding for the given text."""
+        embeddings = await self.get_embeddings(text)
+        return embeddings[0]
+
     async def prepare_chunks(
         self, document_id: str, chunks: List[Document]
     ) -> List[Dict[str, Any]]:

backend/src/app/services/vector_db/milvus_service.py

@@ -155,9 +155,8 @@ async def vector_search(
         # Search for each query
         for query in queries:
             logger.info("Generating embedding.")
-
-            # Embed the query
-            embedded_query = await self.get_embeddings(query)
+            # Use get_single_embedding but wrap result in list for Milvus
+            embedded_query = [await self.get_single_embedding(query)]
 
             logger.info("Searching...")
 
@@ -323,7 +322,7 @@ def count_keywords(text: str, keywords: List[str]) -> int:
             )
 
         # Embed the query
-        embedded_query = await self.get_embeddings(query)
+        embedded_query = [await self.get_single_embedding(query)]
 
         try:
             # First, let's check if there are any vectors for this document_id

backend/src/app/services/vector_db/qdrant_service.py

@@ -75,7 +75,7 @@ async def vector_search(
 
         for query in queries:
             logger.info("Generating embedding.")
-            embedded_query = await self.get_embeddings(query)
+            embedded_query = await self.get_single_embedding(query)
             logger.info("Searching...")
 
             query_response = self.client.query_points(
@@ -162,7 +162,7 @@ def count_keywords(text: str, keywords: List[str]) -> int:
                 reverse=True,
             )
 
-        embedded_query = await self.get_embeddings(query)
+        embedded_query = await self.get_single_embedding(query)
         logger.info("Running semantic similarity search.")
 
         semantic_response = self.client.query_points(
@@ -194,8 +194,6 @@ def count_keywords(text: str, keywords: List[str]) -> int:
             combined_chunks, key=lambda chunk: chunk["chunk_number"]
         )
 
-        # Optionally, for each chunk, retrieve neighbouring chunks to ensure full context is retrieved
-
         # Eliminate duplicate chunks
         seen_chunks = set()
         formatted_output = []

backend/tests/test_service_vector_db_milvus.py

@@ -13,7 +13,7 @@ def __init__(self, embedding_service, llm_service, settings):
         self.embedding_service = embedding_service
         self.llm_service = llm_service
         self.settings = settings
-        self.client = Mock()  # Use regular Mock instead of AsyncMock
+        self.client = Mock()
 
         # Set up synchronous return values
         self.client.has_collection.return_value = True
@@ -36,12 +36,17 @@ async def upsert_vectors(self, vectors):
         }
 
     async def vector_search(self, queries, document_id):
+        # Mock using get_single_embedding
+        for query in queries:
+            _ = await self.get_single_embedding(query)
         return VectorResponseSchema(message="success", chunks=[])
 
     async def keyword_search(self, query, document_id, keywords):
         return VectorResponseSchema(message="success", chunks=[])
 
     async def hybrid_search(self, query, document_id, rules):
+        # Mock using get_single_embedding
+        _ = await self.get_single_embedding(query)
         return VectorResponseSchema(
             message="Query processed successfully.", chunks=[]
         )
@@ -107,3 +112,27 @@ async def test_delete_document(vector_db_service):
 
     assert result["status"] == "success"
     assert result["message"] == "Document deleted successfully."
+
+
+@pytest.mark.asyncio
+async def test_get_single_embedding(vector_db_service):
+    # Reset the mock before the test
+    vector_db_service.embedding_service.get_embeddings.reset_mock()
+
+    # Mock the embedding service to return a known value
+    vector_db_service.embedding_service.get_embeddings.return_value = [
+        [0.1, 0.2, 0.3]
+    ]
+
+    # Test getting a single embedding
+    result = await vector_db_service.get_single_embedding("test text")
+
+    # Verify the result
+    assert isinstance(result, list)
+    assert len(result) == 3  # Length of our mock embedding
+    assert result == [0.1, 0.2, 0.3]
+
+    # Verify the embedding service was called correctly
+    vector_db_service.embedding_service.get_embeddings.assert_called_once_with(
+        ["test text"]
+    )

backend/tests/test_service_vector_db_qdrant.py

@@ -77,7 +77,7 @@ async def test_upsert_vectors(qdrant_service):
 
 @pytest.mark.asyncio
 async def test_vector_search(qdrant_service, mock_embeddings_service):
-    mock_embeddings_service.get_embeddings.return_value = [0.1, 0.2]
+    mock_embeddings_service.get_embeddings.return_value = [[0.1, 0.2]]
 
     result = await qdrant_service.vector_search(["test query"], "test_doc")
 
@@ -88,11 +88,8 @@ async def test_vector_search(qdrant_service, mock_embeddings_service):
 
 @pytest.mark.asyncio
 async def test_hybrid_search(qdrant_service, mock_embeddings_service):
-    # Mock the embedding service response
-    mock_embeddings_service.get_embeddings.return_value = [0.1, 0.2]
+    mock_embeddings_service.get_embeddings.return_value = [[0.1, 0.2]]
 
-    # Mock the extract_keywords method directly on the qdrant_service
-    # since it's a method of QdrantService, not CompletionService
     with patch.object(
         qdrant_service,
         "extract_keywords",
@@ -134,3 +131,27 @@ async def test_delete_document(qdrant_service):
 async def test_keyword_search_not_implemented(qdrant_service):
     with pytest.raises(NotImplementedError):
         await qdrant_service.keyword_search("query", "doc_id", ["keyword"])
+
+
+@pytest.mark.asyncio
+async def test_get_single_embedding(qdrant_service):
+    # Reset the mock before the test
+    qdrant_service.embedding_service.get_embeddings.reset_mock()
+
+    # Mock the embedding service to return a known value
+    qdrant_service.embedding_service.get_embeddings.return_value = [
+        [0.1, 0.2, 0.3]
+    ]
+
+    # Test getting a single embedding
+    result = await qdrant_service.get_single_embedding("test text")
+
+    # Verify the result
+    assert isinstance(result, list)
+    assert len(result) == 3
+    assert result == [0.1, 0.2, 0.3]
+
+    # Verify the embedding service was called correctly
+    qdrant_service.embedding_service.get_embeddings.assert_called_once_with(
+        ["test text"]
+    )