UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-03-22 22:07:15 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

Merge pull request #5829 from jannikstdl/query-embedding-perf-fix

Browse Source 
fix:  performance issues on large collections
This commit is contained in:
Timothy Jaeryang Baek 2024-10-04 10:01:17 +02:00 committed by GitHub
commit 0876c9b5ef
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
backend/open_webui/apps/retrieval/utils.py
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: