gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: Fix batch search function

Browse source 
Rewrite batch search to work as async gather

Fix
This commit is contained in:
Igor Ilic 2024-12-03 11:58:34 +01:00
parent 764c0895df
commit c301498da0
1 changed files with 8 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
cognee/infrastructure/databases/vector/milvus/MilvusAdapter.py
View file

@ -216,10 +216,15 @@ class MilvusAdapter(VectorDBInterface):
raise e raise e
async def batch_search(self, collection_name: str, query_texts: List[str], limit: int, with_vectors: bool = False): async def batch_search(self, collection_name: str, query_texts: List[str], limit: int, with_vectors: bool = False):
def query_search(query_vector): query_vectors = await self.embed_data(query_texts)
return self.search(collection_name, query_vector=query_vector, limit=limit, with_vector=with_vectors)
return [await query_search(query_vector) for query_vector in await self.embed_data(query_texts)] return await asyncio.gather(
*[self.search(collection_name=collection_name,
query_vector=query_vector,
limit=limit,
with_vector=with_vectors,
) for query_vector in query_vectors]
)
async def delete_data_points(self, collection_name: str, data_point_ids: list[str]): async def delete_data_points(self, collection_name: str, data_point_ids: list[str]):
from pymilvus import MilvusException from pymilvus import MilvusException