refactor: Improve query performance by adding filtered IDs (#137)

* feat: Improved query performance by adding filters

* added atlas mongo function

Author: shirsho-dutt

app/routes/document_routes.py

@@ -83,10 +83,10 @@ async def health_check():
 async def get_documents_by_ids(ids: list[str] = Query(...)):
     try:
         if isinstance(vector_store, AsyncPgVector):
-            existing_ids = await vector_store.get_all_ids()
+            existing_ids = await vector_store.get_filtered_ids(ids)
             documents = await vector_store.get_documents_by_ids(ids)
         else:
-            existing_ids = vector_store.get_all_ids()
+            existing_ids = vector_store.get_filtered_ids(ids)
             documents = vector_store.get_documents_by_ids(ids)
 
         # Ensure all requested ids exist
@@ -121,10 +121,10 @@ async def get_documents_by_ids(ids: list[str] = Query(...)):
 async def delete_documents(document_ids: List[str] = Body(...)):
     try:
         if isinstance(vector_store, AsyncPgVector):
-            existing_ids = await vector_store.get_all_ids()
+            existing_ids = await vector_store.get_filtered_ids(document_ids)
             await vector_store.delete(ids=document_ids)
         else:
-            existing_ids = vector_store.get_all_ids()
+            existing_ids = vector_store.get_filtered_ids(document_ids)
             vector_store.delete(ids=document_ids)
 
         if not all(id in existing_ids for id in document_ids):
@@ -456,10 +456,10 @@ async def load_document_context(id: str):
     ids = [id]
     try:
         if isinstance(vector_store, AsyncPgVector):
-            existing_ids = await vector_store.get_all_ids()
+            existing_ids = await vector_store.get_filtered_ids(ids)
             documents = await vector_store.get_documents_by_ids(ids)
         else:
-            existing_ids = vector_store.get_all_ids()
+            existing_ids = vector_store.get_filtered_ids(ids)
             documents = vector_store.get_documents_by_ids(ids)
 
         # Ensure the requested id exists

app/services/vector_store/async_pg_vector.py

@@ -6,6 +6,9 @@
 class AsyncPgVector(ExtendedPgVector):
     async def get_all_ids(self) -> list[str]:
         return await run_in_executor(None, super().get_all_ids)
+    
+    async def get_filtered_ids(self, ids: list[str]) -> list[str]:
+        return await run_in_executor(None, super().get_filtered_ids, ids)
 
     async def get_documents_by_ids(self, ids: list[str]) -> list[Document]:
         return await run_in_executor(None, super().get_documents_by_ids, ids)

app/services/vector_store/atlas_mongo_vector.py

@@ -44,6 +44,10 @@ def similarity_search_with_score_by_vector(
     def get_all_ids(self) -> list[str]:
         # Return unique file_id fields in self._collection
         return self._collection.distinct("file_id")
+    
+    def get_filtered_ids(self, ids: list[str]) -> list[str]:
+        # Return unique file_id fields filtered by the provided ids
+        return self._collection.distinct("file_id", {"file_id": {"$in": ids}})
 
     def get_documents_by_ids(self, ids: list[str]) -> list[Document]:
         # Return documents filtered by file_id

app/services/vector_store/extended_pg_vector.py

@@ -10,6 +10,12 @@ def get_all_ids(self) -> list[str]:
         with Session(self._bind) as session:
             results = session.query(self.EmbeddingStore.custom_id).all()
             return [result[0] for result in results if result[0] is not None]
+        
+    def get_filtered_ids(self, ids: list[str]) -> list[str]:
+        with Session(self._bind) as session:
+            query = session.query(self.EmbeddingStore.custom_id).filter(self.EmbeddingStore.custom_id.in_(ids))
+            results = query.all()
+            return [result[0] for result in results if result[0] is not None]
 
     def get_documents_by_ids(self, ids: list[str]) -> list[Document]:
         with Session(self._bind) as session:

tests/conftest.py

@@ -26,6 +26,10 @@ def dummy_post_init(self):
 class DummyVectorStore:
     def get_all_ids(self) -> list[str]:
         return ["testid1", "testid2"]
+    
+    def get_filtered_ids(self, ids) -> list[str]:
+        dummy_ids = ["testid1", "testid2"]
+        return [id for id in dummy_ids if id in ids]
 
     async def get_documents_by_ids(self, ids: list[str]) -> list[Document]:
         return [

tests/test_main.py

@@ -29,6 +29,12 @@ async def dummy_get_all_ids():
         return ["testid1", "testid2"]
     monkeypatch.setattr(vector_store, "get_all_ids", dummy_get_all_ids)
 
+    # Override get_filtered_ids as an async function.
+    async def dummy_get_filtered_ids(ids):
+        dummy_ids = ["testid1", "testid2"]
+        return [id for id in dummy_ids if id in ids]
+    monkeypatch.setattr(vector_store, "get_filtered_ids", dummy_get_filtered_ids)
+
     # Override get_documents_by_ids as an async function.
     async def dummy_get_documents_by_ids(ids):
         return [