From 93614651628a48ff9eaf31c4569f8190377946ed Mon Sep 17 00:00:00 2001 From: Stefano Fiorucci Date: Tue, 27 Feb 2024 11:11:35 +0100 Subject: [PATCH] Weaviate: make retrievers return dicts (#491) --- .../components/retrievers/weaviate/bm25_retriever.py | 3 ++- .../components/retrievers/weaviate/embedding_retriever.py | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/bm25_retriever.py b/integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/bm25_retriever.py index 6c27378cf..c37312604 100644 --- a/integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/bm25_retriever.py +++ b/integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/bm25_retriever.py @@ -47,4 +47,5 @@ def from_dict(cls, data: Dict[str, Any]) -> "WeaviateBM25Retriever": def run(self, query: str, filters: Optional[Dict[str, Any]] = None, top_k: Optional[int] = None): filters = filters or self._filters top_k = top_k or self._top_k - return self._document_store._bm25_retrieval(query=query, filters=filters, top_k=top_k) + documents = self._document_store._bm25_retrieval(query=query, filters=filters, top_k=top_k) + return {"documents": documents} diff --git a/integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/embedding_retriever.py b/integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/embedding_retriever.py index b8a163b56..38f7cd85f 100644 --- a/integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/embedding_retriever.py +++ b/integrations/weaviate/src/haystack_integrations/components/retrievers/weaviate/embedding_retriever.py @@ -71,10 +71,11 @@ def run( top_k = top_k or self._top_k distance = distance or self._distance certainty = certainty or self._certainty - return self._document_store._embedding_retrieval( + documents = self._document_store._embedding_retrieval( query_embedding=query_embedding, filters=filters, top_k=top_k, distance=distance, certainty=certainty, ) + return {"documents": documents}