From 792e528fc4be0abfd6e6e9519617a2b94a363705 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Micha=C5=82=20Pstr=C4=85g?= Date: Mon, 28 Oct 2024 03:03:25 +0100 Subject: [PATCH] small fix --- .../src/ragbits/document_search/_main.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/packages/ragbits-document-search/src/ragbits/document_search/_main.py b/packages/ragbits-document-search/src/ragbits/document_search/_main.py index ff37e9ead..40d3f0101 100644 --- a/packages/ragbits-document-search/src/ragbits/document_search/_main.py +++ b/packages/ragbits-document-search/src/ragbits/document_search/_main.py @@ -5,6 +5,7 @@ from ragbits.core.embeddings import Embeddings, get_embeddings from ragbits.core.vector_store import VectorStore, get_vector_store +from ragbits.core.vector_store.base import VectorStoreOptions from ragbits.document_search.documents.document import Document, DocumentMeta from ragbits.document_search.documents.element import Element from ragbits.document_search.documents.sources import GCSSource, LocalFileSource, Source @@ -98,7 +99,10 @@ async def search(self, query: str, config: SearchConfig | None = None) -> list[E elements = [] for rephrased_query in queries: search_vector = await self.embedder.embed_text([rephrased_query]) - entries = await self.vector_store.retrieve(search_vector[0]) + entries = await self.vector_store.retrieve( + vector=search_vector[0], + options=VectorStoreOptions(**config.vector_store_kwargs), + ) elements.extend([Element.from_vector_db_entry(entry) for entry in entries]) return self.reranker.rerank(elements)