diff --git a/cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py b/cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py index a8df0b2e0..9289bb6c8 100644 --- a/cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py +++ b/cognee/infrastructure/databases/hybrid/neptune_analytics/NeptuneAnalyticsAdapter.py @@ -320,7 +320,12 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface): self._na_exception_handler(e, query_string) async def batch_search( - self, collection_name: str, query_texts: List[str], limit: int, with_vectors: bool = False + self, + collection_name: str, + query_texts: List[str], + limit: int, + with_vectors: bool = False, + include_payload: bool = False, ): """ Perform a batch search using multiple text queries against a collection. @@ -343,7 +348,14 @@ class NeptuneAnalyticsAdapter(NeptuneGraphDB, VectorDBInterface): data_vectors = await self.embedding_engine.embed_text(query_texts) return await asyncio.gather( *[ - self.search(collection_name, None, vector, limit, with_vectors) + self.search( + collection_name, + None, + vector, + limit, + with_vectors, + include_payload=include_payload, + ) for vector in data_vectors ] ) diff --git a/cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py b/cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py index 5e2d3975a..bec97ca94 100644 --- a/cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py +++ b/cognee/infrastructure/databases/vector/chromadb/ChromaDBAdapter.py @@ -442,6 +442,7 @@ class ChromaDBAdapter(VectorDBInterface): query_texts: List[str], limit: int = 5, with_vectors: bool = False, + include_payload: bool = False, ): """ Perform multiple searches in a single request for efficiency, returning results for each diff --git a/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py b/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py index 49168ffcd..baef75d9e 100644 --- a/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py +++ b/cognee/infrastructure/databases/vector/lancedb/LanceDBAdapter.py @@ -260,6 +260,7 @@ class LanceDBAdapter(VectorDBInterface): .limit(limit) .to_list() ) + if not result_values: return [] normalized_values = normalize_distances(result_values) @@ -279,6 +280,7 @@ class LanceDBAdapter(VectorDBInterface): query_texts: List[str], limit: Optional[int] = None, with_vectors: bool = False, + include_payload: bool = False, ): query_vectors = await self.embedding_engine.embed_text(query_texts) @@ -289,6 +291,7 @@ class LanceDBAdapter(VectorDBInterface): query_vector=query_vector, limit=limit, with_vector=with_vectors, + include_payload=include_payload, ) for query_vector in query_vectors ] diff --git a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py index 6c73ad475..5e2c356ee 100644 --- a/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py +++ b/cognee/infrastructure/databases/vector/pgvector/PGVectorAdapter.py @@ -380,6 +380,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface): query_texts: List[str], limit: int = None, with_vectors: bool = False, + include_payload: bool = False, ): query_vectors = await self.embedding_engine.embed_text(query_texts) @@ -390,6 +391,7 @@ class PGVectorAdapter(SQLAlchemyAdapter, VectorDBInterface): query_vector=query_vector, limit=limit, with_vector=with_vectors, + include_payload=include_payload, ) for query_vector in query_vectors ] diff --git a/cognee/infrastructure/databases/vector/vector_db_interface.py b/cognee/infrastructure/databases/vector/vector_db_interface.py index e38de19e2..4376d8713 100644 --- a/cognee/infrastructure/databases/vector/vector_db_interface.py +++ b/cognee/infrastructure/databases/vector/vector_db_interface.py @@ -117,6 +117,7 @@ class VectorDBInterface(Protocol): query_texts: List[str], limit: Optional[int], with_vectors: bool = False, + include_payload: bool = False, ): """ Perform a batch search using multiple text queries against a collection. @@ -129,6 +130,9 @@ class VectorDBInterface(Protocol): - limit (Optional[int]): The maximum number of results to return for each query. - with_vectors (bool): Whether to include vector representations with search results. (default False) + - include_payload (bool): Whether to include the payload data with search. Search is faster when set to False. + Payload contains metadata about the data point, useful for searches that are only based on embedding distances + like the RAG_COMPLETION search type, but not needed when search also contains graph data. """ raise NotImplementedError