chore: review update with score_threshold

This commit is contained in:
Anush008 2024-10-23 00:54:47 +05:30
parent d1bef44e2b
commit bc27046f36
No known key found for this signature in database
6 changed files with 24 additions and 13 deletions

View file

@ -38,7 +38,9 @@ class ChromaIndex(EmbeddingIndex):
ids=[f"{c.document_id}:chunk-{i}" for i, c in enumerate(chunks)],
)
async def query(self, embedding: NDArray, k: int) -> QueryDocumentsResponse:
async def query(
self, embedding: NDArray, k: int, score_threshold: float
) -> QueryDocumentsResponse:
results = await self.collection.query(
query_embeddings=[embedding.tolist()],
n_results=k,