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

Co-authored-by: Dev 2049 <dev.dev2049@gmail.com>
Co-authored-by: Davis Chase <130488702+dev2049@users.noreply.github.com>
searx_updates
escafati 1 year ago committed by GitHub
parent 3df2d831f9
commit e027a38f33
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

@ -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,

Loading…
Cancel
Save