From 6a4ee07e4f12f1d2c7d4103d83f79cbb5d27d51a Mon Sep 17 00:00:00 2001 From: Kentaro Tanaka Date: Sat, 4 Mar 2023 17:20:18 +0900 Subject: [PATCH] Fix type hint of 'vectorstore_cls' arg in `SemanticSimilarityExampleSelector` (#1427) Hello! Thank you for the amazing library you've created! While following the tutorial at [the link(`Using an example selector`)](https://langchain.readthedocs.io/en/latest/modules/prompts/examples/few_shot_examples.html#using-an-example-selector), I noticed that passing Chroma as an argument to from_examples results in a type hint error. Error message(mypy): ``` Argument 3 to "from_examples" of "SemanticSimilarityExampleSelector" has incompatible type "Type[Chroma]"; expected "VectorStore" [arg-type]mypy(error) ``` This pull request fixes the type hint and allows the VectorStore class to be specified as an argument. --- langchain/prompts/example_selector/semantic_similarity.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/langchain/prompts/example_selector/semantic_similarity.py b/langchain/prompts/example_selector/semantic_similarity.py index e4c8e553..604a04e6 100644 --- a/langchain/prompts/example_selector/semantic_similarity.py +++ b/langchain/prompts/example_selector/semantic_similarity.py @@ -1,7 +1,7 @@ """Example selector that selects examples based on SemanticSimilarity.""" from __future__ import annotations -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Type from pydantic import BaseModel, Extra @@ -65,7 +65,7 @@ class SemanticSimilarityExampleSelector(BaseExampleSelector, BaseModel): cls, examples: List[dict], embeddings: Embeddings, - vectorstore_cls: VectorStore, + vectorstore_cls: Type[VectorStore], k: int = 4, input_keys: Optional[List[str]] = None, **vectorstore_cls_kwargs: Any, @@ -131,7 +131,7 @@ class MaxMarginalRelevanceExampleSelector(SemanticSimilarityExampleSelector, Bas cls, examples: List[dict], embeddings: Embeddings, - vectorstore_cls: VectorStore, + vectorstore_cls: Type[VectorStore], k: int = 4, input_keys: Optional[List[str]] = None, fetch_k: int = 20,