From e027a38f33049b837a33415322ab8083e72995e4 Mon Sep 17 00:00:00 2001 From: escafati Date: Thu, 18 May 2023 21:35:31 -0300 Subject: [PATCH] NIT: Instead of hardcoding k in each definition, define it as a param above. (#2675) Co-authored-by: Dev 2049 Co-authored-by: Davis Chase <130488702+dev2049@users.noreply.github.com> --- langchain/vectorstores/chroma.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/langchain/vectorstores/chroma.py b/langchain/vectorstores/chroma.py index e102abe4..f83c8c22 100644 --- a/langchain/vectorstores/chroma.py +++ b/langchain/vectorstores/chroma.py @@ -17,7 +17,8 @@ if TYPE_CHECKING: import chromadb import chromadb.config -logger = logging.getLogger(__name__) +logger = logging.getLogger() +DEFAULT_K = 4 # Number of Documents to return. def _results_to_docs(results: Any) -> List[Document]: @@ -164,7 +165,7 @@ class Chroma(VectorStore): def similarity_search( self, query: str, - k: int = 4, + k: int = DEFAULT_K, filter: Optional[Dict[str, str]] = None, **kwargs: Any, ) -> List[Document]: @@ -184,7 +185,7 @@ class Chroma(VectorStore): def similarity_search_by_vector( self, embedding: List[float], - k: int = 4, + k: int = DEFAULT_K, filter: Optional[Dict[str, str]] = None, **kwargs: Any, ) -> List[Document]: @@ -204,7 +205,7 @@ class Chroma(VectorStore): def similarity_search_with_score( self, query: str, - k: int = 4, + k: int = DEFAULT_K, filter: Optional[Dict[str, str]] = None, **kwargs: Any, ) -> List[Tuple[Document, float]]: @@ -234,7 +235,7 @@ class Chroma(VectorStore): def max_marginal_relevance_search_by_vector( self, embedding: List[float], - k: int = 4, + k: int = DEFAULT_K, fetch_k: int = 20, lambda_mult: float = 0.5, filter: Optional[Dict[str, str]] = None, @@ -277,7 +278,7 @@ class Chroma(VectorStore): def max_marginal_relevance_search( self, query: str, - k: int = 4, + k: int = DEFAULT_K, fetch_k: int = 20, lambda_mult: float = 0.5, filter: Optional[Dict[str, str]] = None,