core: add kwargs support to VectorStore (#25934)

has been missing the passthrough until now

---------

Co-authored-by: Erick Friis <erick@langchain.dev>
This commit is contained in:
Harrison Chase
2024-12-16 10:57:57 -08:00
committed by GitHub
parent 87c50f99e5
commit de7996c2ca
5 changed files with 50 additions and 38 deletions

View File

@@ -1057,7 +1057,9 @@ class VectorStoreRetriever(BaseRetriever):
def _get_ls_params(self, **kwargs: Any) -> LangSmithRetrieverParams:
"""Get standard params for tracing."""
ls_params = super()._get_ls_params(**kwargs)
_kwargs = self.search_kwargs | kwargs
ls_params = super()._get_ls_params(**_kwargs)
ls_params["ls_vector_store_provider"] = self.vectorstore.__class__.__name__
if self.vectorstore.embeddings:
@@ -1074,43 +1076,45 @@ class VectorStoreRetriever(BaseRetriever):
return ls_params
def _get_relevant_documents(
self, query: str, *, run_manager: CallbackManagerForRetrieverRun
self, query: str, *, run_manager: CallbackManagerForRetrieverRun, **kwargs: Any
) -> list[Document]:
_kwargs = self.search_kwargs | kwargs
if self.search_type == "similarity":
docs = self.vectorstore.similarity_search(query, **self.search_kwargs)
docs = self.vectorstore.similarity_search(query, **_kwargs)
elif self.search_type == "similarity_score_threshold":
docs_and_similarities = (
self.vectorstore.similarity_search_with_relevance_scores(
query, **self.search_kwargs
query, **_kwargs
)
)
docs = [doc for doc, _ in docs_and_similarities]
elif self.search_type == "mmr":
docs = self.vectorstore.max_marginal_relevance_search(
query, **self.search_kwargs
)
docs = self.vectorstore.max_marginal_relevance_search(query, **_kwargs)
else:
msg = f"search_type of {self.search_type} not allowed."
raise ValueError(msg)
return docs
async def _aget_relevant_documents(
self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
self,
query: str,
*,
run_manager: AsyncCallbackManagerForRetrieverRun,
**kwargs: Any,
) -> list[Document]:
_kwargs = self.search_kwargs | kwargs
if self.search_type == "similarity":
docs = await self.vectorstore.asimilarity_search(
query, **self.search_kwargs
)
docs = await self.vectorstore.asimilarity_search(query, **_kwargs)
elif self.search_type == "similarity_score_threshold":
docs_and_similarities = (
await self.vectorstore.asimilarity_search_with_relevance_scores(
query, **self.search_kwargs
query, **_kwargs
)
)
docs = [doc for doc, _ in docs_and_similarities]
elif self.search_type == "mmr":
docs = await self.vectorstore.amax_marginal_relevance_search(
query, **self.search_kwargs
query, **_kwargs
)
else:
msg = f"search_type of {self.search_type} not allowed."