From ca446f3f11fd97069974a008d1e3a0062151bf10 Mon Sep 17 00:00:00 2001 From: "VERBI-BLN\\ignacio.llorca" Date: Mon, 5 Feb 2024 13:22:38 +0100 Subject: [PATCH] Allow persitent docstore and filters --- .../components/retrievers/chroma/retriever.py | 4 ++-- .../document_stores/chroma/document_store.py | 17 +++++++++++++---- 2 files changed, 15 insertions(+), 6 deletions(-) diff --git a/integrations/chroma/src/haystack_integrations/components/retrievers/chroma/retriever.py b/integrations/chroma/src/haystack_integrations/components/retrievers/chroma/retriever.py index 9388171f4..7ae3b26f4 100644 --- a/integrations/chroma/src/haystack_integrations/components/retrievers/chroma/retriever.py +++ b/integrations/chroma/src/haystack_integrations/components/retrievers/chroma/retriever.py @@ -66,7 +66,7 @@ class ChromaEmbeddingRetriever(ChromaQueryRetriever): def run( self, query_embedding: List[float], - _: Optional[Dict[str, Any]] = None, # filters not yet supported + filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None, ): """ @@ -80,4 +80,4 @@ def run( top_k = top_k or self.top_k query_embeddings = [query_embedding] - return {"documents": self.document_store.search_embeddings(query_embeddings, top_k)[0]} + return {"documents": self.document_store.search_embeddings(query_embeddings, top_k, filters)[0]} \ No newline at end of file diff --git a/integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py b/integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py index 1b11d047c..522de8bfe 100644 --- a/integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py +++ b/integrations/chroma/src/haystack_integrations/document_stores/chroma/document_store.py @@ -24,7 +24,7 @@ class ChromaDocumentStore: """ def __init__( - self, collection_name: str = "documents", embedding_function: str = "default", **embedding_function_params + self, collection_name: str = "documents", embedding_function: str = "default", persist_path: str = None, **embedding_function_params ): """ Initializes the store. The __init__ constructor is not part of the Store Protocol @@ -40,7 +40,7 @@ def __init__( self._embedding_function = embedding_function self._embedding_function_params = embedding_function_params # Create the client instance - self._chroma_client = chromadb.Client() + self._chroma_client = chromadb.Client() if persist_path == None else chromadb.PersistentClient(path=persist_path) self._collection = self._chroma_client.get_or_create_collection( name=collection_name, embedding_function=get_embedding_function(embedding_function, **embedding_function_params), @@ -185,14 +185,23 @@ def search(self, queries: List[str], top_k: int) -> List[List[Document]]: ) return self._query_result_to_documents(results) - def search_embeddings(self, query_embeddings: List[List[float]], top_k: int) -> List[List[Document]]: + def search_embeddings(self, query_embeddings: List[List[float]], top_k: int, filters: Optional[Dict[str, Any]] = None) -> List[List[Document]]: """ Perform vector search on the stored document, pass the embeddings of the queries - instead of their text + instead of their text. + + Accepts filters in haystack format. """ + if filters is not None: + chroma_filters = self._normalize_filters(filters=filters) + else: + chroma_filters = (None, None, None) + results = self._collection.query( query_embeddings=query_embeddings, n_results=top_k, + where=chroma_filters[1], + where_document=chroma_filters[2], include=["embeddings", "documents", "metadatas", "distances"], ) return self._query_result_to_documents(results)