fix: do not embed the query every single time

This commit is contained in:
Jannik Streidl 2024-09-30 16:18:02 +02:00
parent 37e0d47082
commit 6bff5a4d09

View File

@ -65,14 +65,13 @@ class VectorSearchRetriever(BaseRetriever):
def query_doc(
collection_name: str,
query: str,
embedding_function,
query_embedding: list[float],
k: int,
):
try:
result = VECTOR_DB_CLIENT.search(
collection_name=collection_name,
vectors=[embedding_function(query)],
vectors=[query_embedding],
limit=k,
)
@ -182,15 +181,17 @@ def query_collection(
embedding_function,
k: int,
) -> dict:
results = []
query_embedding = embedding_function(query)
for collection_name in collection_names:
if collection_name:
try:
result = query_doc(
collection_name=collection_name,
query=query,
k=k,
embedding_function=embedding_function,
query_embedding=query_embedding,
)
results.append(result.model_dump())
except Exception as e: