fix(open_embedding): use embedding for openai in batches as well (#53)

* fix(open_embedding): use embedding for openai in batches as well

* fix: test cases

Author: ArslanSaleem

backend/app/processing/file_preprocessing.py

@@ -43,7 +43,8 @@ def process_segmentation(project_id: int, asset_id: int, asset_file_name: str):
 
         vectorstore.add_docs(
             docs=docs,
-            metadatas=metadatas
+            metadatas=metadatas,
+            batch_size=100
         )
 
         project_repository.update_asset_content_status(
@@ -67,6 +68,7 @@ def preprocess_file(asset_id: int):
         # Get asset details from the database first
         with SessionLocal() as db:
             asset = project_repository.get_asset(db=db, asset_id=asset_id)
+
             if asset is None:
                 logger.error(f"Asset with id {asset_id} not found in the database")
                 return

backend/app/processing/process_queue.py

@@ -406,4 +406,4 @@ def vectorize_extraction_process_step(project_id: int, process_step_id: int, fil
     ]
 
     # Add documents to vectorstore
-    vectorstore.add_docs(docs=docs, metadatas=metadatas)
+    vectorstore.add_docs(docs=docs, metadatas=metadatas, batch_size=100)

backend/app/vectorstore/chroma.py

@@ -101,25 +101,16 @@ def add_docs(
         filename = metadatas[0].get('filename', 'unknown')
         logger.info(f"Adding {len(docs)} sentences to the vector store for file {filename}")
 
-        # If using OpenAI embeddings, add all documents at once
-        if self.settings.use_openai_embeddings and self.settings.openai_api_key:
-            logger.info("Using OpenAI embeddings")
+        # Batching the document processing
+        batch_size = batch_size or self._batch_size
+
+        for i in range(0, len(docs), batch_size):
+            logger.info(f"Processing batch {i} to {i + batch_size}")
             self._docs_collection.add(
-                documents=list(docs),
-                metadatas=metadatas,
-                ids=ids,
+                documents=docs[i : i + batch_size],
+                metadatas=metadatas[i : i + batch_size],
+                ids=ids[i : i + batch_size],
             )
-        else:
-            logger.info("Using default embedding function")
-            batch_size = batch_size or self._batch_size
-
-            for i in range(0, len(docs), batch_size):
-                logger.info(f"Processing batch {i} to {i + batch_size}")
-                self._docs_collection.add(
-                    documents=docs[i : i + batch_size],
-                    metadatas=metadatas[i : i + batch_size],
-                    ids=ids[i : i + batch_size],
-                )
 
         return list(ids)
 

backend/tests/processing/test_process_queue.py

@@ -215,7 +215,8 @@ def test_vectorize_extraction_process_step_single_reference(mock_chroma_db):
     # Assertions
     mock_vectorstore.add_docs.assert_called_once_with(
         docs=expected_docs,
-        metadatas=expected_metadatas
+        metadatas=expected_metadatas,
+        batch_size=100
     )
 
 @patch('app.processing.process_queue.ChromaDB')
@@ -261,7 +262,8 @@ def test_vectorize_extraction_process_step_multiple_references_concatenation(moc
     # Assertions
     mock_vectorstore.add_docs.assert_called_once_with(
         docs=expected_docs,
-        metadatas=expected_metadatas
+        metadatas=expected_metadatas,
+        batch_size=100
     )
 
 @patch('app.processing.process_queue.ChromaDB')  # Replace with the correct module path