github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-07-03 11:47:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

NIT: Instead of hardcoding k in each definition, define it as a param above. (#2675)

Browse Source 
Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
Co-authored-by: Davis Chase <130488702+dev2049@users.noreply.github.com>
This commit is contained in:
escafati 2023-05-18 21:35:31 -03:00 committed by GitHub
parent 3df2d831f9
commit e027a38f33
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
langchain/vectorstores/chroma.py
View File

@ -17,7 +17,8 @@ if TYPE_CHECKING:
import chromadb import chromadb
import chromadb.config 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]: def _results_to_docs(results: Any) -> List[Document]:
@ -164,7 +165,7 @@ class Chroma(VectorStore):
def similarity_search( def similarity_search(
self, self,
query: str, query: str,
k: int = 4, k: int = DEFAULT_K,
filter: Optional[Dict[str, str]] = None, filter: Optional[Dict[str, str]] = None,
**kwargs: Any, **kwargs: Any,
) -> List[Document]: ) -> List[Document]:
@ -184,7 +185,7 @@ class Chroma(VectorStore):
def similarity_search_by_vector( def similarity_search_by_vector(
self, self,
embedding: List[float], embedding: List[float],
k: int = 4, k: int = DEFAULT_K,
filter: Optional[Dict[str, str]] = None, filter: Optional[Dict[str, str]] = None,
**kwargs: Any, **kwargs: Any,
) -> List[Document]: ) -> List[Document]:
@ -204,7 +205,7 @@ class Chroma(VectorStore):
def similarity_search_with_score( def similarity_search_with_score(
self, self,
query: str, query: str,
k: int = 4, k: int = DEFAULT_K,
filter: Optional[Dict[str, str]] = None, filter: Optional[Dict[str, str]] = None,
**kwargs: Any, **kwargs: Any,
) -> List[Tuple[Document, float]]: ) -> List[Tuple[Document, float]]:
@ -234,7 +235,7 @@ class Chroma(VectorStore):
def max_marginal_relevance_search_by_vector( def max_marginal_relevance_search_by_vector(
self, self,
embedding: List[float], embedding: List[float],
k: int = 4, k: int = DEFAULT_K,
fetch_k: int = 20, fetch_k: int = 20,
lambda_mult: float = 0.5, lambda_mult: float = 0.5,
filter: Optional[Dict[str, str]] = None, filter: Optional[Dict[str, str]] = None,
@ -277,7 +278,7 @@ class Chroma(VectorStore):
def max_marginal_relevance_search( def max_marginal_relevance_search(
self, self,
query: str, query: str,
k: int = 4, k: int = DEFAULT_K,
fetch_k: int = 20, fetch_k: int = 20,
lambda_mult: float = 0.5, lambda_mult: float = 0.5,
filter: Optional[Dict[str, str]] = None, filter: Optional[Dict[str, str]] = None,