From cc0f2bc5f6158f5ca74cfbfcca7aa7665845343d Mon Sep 17 00:00:00 2001 From: Massimiliano Pippi Date: Thu, 29 Feb 2024 16:27:05 +0100 Subject: [PATCH] make write_documents compatible with the DocumentStore protocol (#505) --- .../document_stores/chroma/document_store.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) 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 0da89c87e..8bb90aa2e 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 @@ -154,14 +154,20 @@ def filter_documents(self, filters: Optional[Dict[str, Any]] = None) -> List[Doc return self._get_result_to_documents(result) - def write_documents(self, documents: List[Document], policy: DuplicatePolicy = DuplicatePolicy.FAIL) -> None: + def write_documents(self, documents: List[Document], policy: DuplicatePolicy = DuplicatePolicy.FAIL) -> int: """ Writes (or overwrites) documents into the store. - :param documents: a list of documents. - :param policy: not supported at the moment - :raises DuplicateDocumentError: Exception trigger on duplicate document if `policy=DuplicatePolicy.FAIL` - :returns: None + :param documents: + A list of documents to write into the document store. + :param policy: + Not supported at the moment. + + :raises ValueError: + When input is not valid. + + :returns: + The number of documents written """ for doc in documents: if not isinstance(doc, Document): @@ -183,6 +189,8 @@ def write_documents(self, documents: List[Document], policy: DuplicatePolicy = D self._collection.add(**data) + return len(documents) + def delete_documents(self, document_ids: List[str]) -> None: """ Deletes all documents with a matching document_ids from the document store.