adding vectorstore_kwarg attribute to search_similarity function (#14604)

- **Description:** the ability to add all extra parameter of vectorstore
and using them SemanticSimilarityExampleSelector.
  - **Issue:** #14583
  - **Dependencies:** no dependensies
  - **Tag maintainer:** 
  - **Twitter handle:** @AmirMalekiz

---------

Co-authored-by: Amir Maleki <amaleki@fb.com>
Co-authored-by: Harrison Chase <hw.chase.17@gmail.com>
This commit is contained in:
amaleki2 2024-01-02 17:18:33 -08:00 committed by GitHub
parent e93be14c11
commit 413a56b8f1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -28,6 +28,8 @@ class SemanticSimilarityExampleSelector(BaseExampleSelector, BaseModel):
input_keys: Optional[List[str]] = None input_keys: Optional[List[str]] = None
"""Optional keys to filter input to. If provided, the search is based on """Optional keys to filter input to. If provided, the search is based on
the input variables instead of all variables.""" the input variables instead of all variables."""
vectorstore_kwargs: Optional[Dict[str, Any]] = None
"""Extra arguments passed to similarity_search function of the vectorstore."""
class Config: class Config:
"""Configuration for this pydantic object.""" """Configuration for this pydantic object."""
@ -51,8 +53,11 @@ class SemanticSimilarityExampleSelector(BaseExampleSelector, BaseModel):
# Get the docs with the highest similarity. # Get the docs with the highest similarity.
if self.input_keys: if self.input_keys:
input_variables = {key: input_variables[key] for key in self.input_keys} input_variables = {key: input_variables[key] for key in self.input_keys}
vectorstore_kwargs = self.vectorstore_kwargs or {}
query = " ".join(sorted_values(input_variables)) query = " ".join(sorted_values(input_variables))
example_docs = self.vectorstore.similarity_search(query, k=self.k) example_docs = self.vectorstore.similarity_search(
query, k=self.k, **vectorstore_kwargs
)
# Get the examples from the metadata. # Get the examples from the metadata.
# This assumes that examples are stored in metadata. # This assumes that examples are stored in metadata.
examples = [dict(e.metadata) for e in example_docs] examples = [dict(e.metadata) for e in example_docs]