Weaviate: make retrievers return dicts (#491)

deepset-ai · Feb 27, 2024 · 9361465 · 9361465
1 parent 6119e5c
commit 9361465
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/...tions/weaviate/src/haystack_integrations/components/retrievers/weaviate/bm25_retriever.py b/...tions/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/.../weaviate/src/haystack_integrations/components/retrievers/weaviate/embedding_retriever.py b/.../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}