Skip to content

Commit 79d0d52

Browse files
feat(Qdrant): SparseEmbedding instead of Dict
1 parent 2683a74 commit 79d0d52

File tree

3 files changed

+8
-10
lines changed

3 files changed

+8
-10
lines changed

integrations/qdrant/src/haystack_integrations/components/retrievers/qdrant/retriever.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -132,6 +132,7 @@ class QdrantSparseRetriever:
132132
```python
133133
from haystack_integrations.components.retrievers.qdrant import QdrantSparseRetriever
134134
from haystack_integrations.document_stores.qdrant import QdrantDocumentStore
135+
from haystack.dataclasses.sparse_embedding import SparseEmbedding
135136
136137
document_store = QdrantDocumentStore(
137138
":memory:",
@@ -140,8 +141,8 @@ class QdrantSparseRetriever:
140141
wait_result_from_api=True,
141142
)
142143
retriever = QdrantSparseRetriever(document_store=document_store)
143-
144-
retriever.run(query_sparse_embedding={"indices":[0, 1, 2, 3], "values":[0.1, 0.8, 0.05, 0.33]})
144+
sparse_embedding = SparseEmbedding(indices=[0, 1, 2, 3], values=[0.1, 0.8, 0.05, 0.33])
145+
retriever.run(query_sparse_embedding=sparse_embedding)
145146
```
146147
"""
147148

integrations/qdrant/src/haystack_integrations/document_stores/qdrant/document_store.py

+2-3
Original file line numberDiff line numberDiff line change
@@ -310,9 +310,8 @@ def query_by_sparse(
310310
return_embedding: bool = False, # noqa: FBT001, FBT002
311311
) -> List[Document]:
312312
qdrant_filters = self.qdrant_filter_converter.convert(filters)
313-
314-
query_indices = query_sparse_embedding["indices"]
315-
query_values = query_sparse_embedding["values"]
313+
query_indices = query_sparse_embedding.indices
314+
query_values = query_sparse_embedding.values
316315

317316
points = self.client.search(
318317
collection_name=self.index,

integrations/qdrant/tests/test_retriever.py

+3-5
Original file line numberDiff line numberDiff line change
@@ -227,14 +227,12 @@ def test_run(self, filterable_docs: List[Document]):
227227
document_store.write_documents(filterable_docs)
228228

229229
retriever = QdrantSparseRetriever(document_store=document_store)
230-
231-
results: List[Document] = retriever.run(query_sparse_embedding=self._generate_mocked_sparse_embedding(1)[0])
230+
sparse_embedding = SparseEmbedding(indices=[0, 1, 2, 3], values=[0.1, 0.8, 0.05, 0.33])
231+
results: List[Document] = retriever.run(query_sparse_embedding=sparse_embedding)
232232

233233
assert len(results["documents"]) == 10 # type: ignore
234234

235-
results = retriever.run(
236-
query_sparse_embedding=self._generate_mocked_sparse_embedding(1)[0], top_k=5, return_embedding=False
237-
)
235+
results = retriever.run(query_sparse_embedding=sparse_embedding, top_k=5, return_embedding=False)
238236

239237
assert len(results["documents"]) == 5 # type: ignore
240238

0 commit comments

Comments
 (0)