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( def query_doc(
collection_name: str, collection_name: str,
query: str, query_embedding: list[float],
embedding_function,
k: int, k: int,
): ):
try: try:
result = VECTOR_DB_CLIENT.search( result = VECTOR_DB_CLIENT.search(
collection_name=collection_name, collection_name=collection_name,
vectors=[embedding_function(query)], vectors=[query_embedding],
limit=k, limit=k,
) )
@ -182,15 +181,17 @@ def query_collection(
embedding_function, embedding_function,
k: int, k: int,
) -> dict: ) -> dict:
results = [] results = []
query_embedding = embedding_function(query)
for collection_name in collection_names: for collection_name in collection_names:
if collection_name: if collection_name:
try: try:
result = query_doc( result = query_doc(
collection_name=collection_name, collection_name=collection_name,
query=query,
k=k, k=k,
embedding_function=embedding_function, query_embedding=query_embedding,
) )
results.append(result.model_dump()) results.append(result.model_dump())
except Exception as e: except Exception as e: