Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: 50x performance improvement for external embeddings (#19296)

Browse Source 
* Update utils.py (#77)

Co-authored-by: Claude <noreply@anthropic.com>

* refactor: address code review feedback for embedding performance improvements (#92)

Co-authored-by: Claude <noreply@anthropic.com>

* fix: prevent sentence transformers from blocking async event loop (#95)

Co-authored-by: Claude <noreply@anthropic.com>

---------

Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Classic298
2025-11-22 20:54:59 -05:00
committed by GitHub
co-authored by Claude
parent bb3e222e09
commit 902c6cfbea
4 changed files with 289 additions and 118 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/routers/retrieval.py
+15 -13
View File
@@ -1467,11 +1467,12 @@ def save_docs_to_vector_db(
),
)
embeddings = embedding_function(
# Run async embedding in sync context
embeddings = asyncio.run(embedding_function(
list(map(lambda x: x.replace("\n", " "), texts)),
prefix=RAG_EMBEDDING_CONTENT_PREFIX,
user=user,
)
))
log.info(f"embeddings generated {len(embeddings)} for {len(texts)} items")
items = [
@@ -2262,7 +2263,7 @@ class QueryDocForm(BaseModel):
@router.post("/query/doc")
def query_doc_handler(
async def query_doc_handler(
request: Request,
form_data: QueryDocForm,
user=Depends(get_verified_user),
@@ -2275,7 +2276,7 @@ def query_doc_handler(
collection_results[form_data.collection_name] = VECTOR_DB_CLIENT.get(
collection_name=form_data.collection_name
)
return query_doc_with_hybrid_search(
return await query_doc_with_hybrid_search(
collection_name=form_data.collection_name,
collection_result=collection_results[form_data.collection_name],
query=form_data.query,
@@ -2285,8 +2286,8 @@ def query_doc_handler(
k=form_data.k if form_data.k else request.app.state.config.TOP_K,
reranking_function=(
(
lambda sentences: request.app.state.RERANKING_FUNCTION(
sentences, user=user
lambda query, documents: request.app.state.RERANKING_FUNCTION(
query, documents, user=user
)
)
if request.app.state.RERANKING_FUNCTION
@@ -2307,11 +2308,12 @@ def query_doc_handler(
user=user,
)
else:
query_embedding = await request.app.state.EMBEDDING_FUNCTION(
form_data.query, prefix=RAG_EMBEDDING_QUERY_PREFIX, user=user
)
return query_doc(
collection_name=form_data.collection_name,
query_embedding=request.app.state.EMBEDDING_FUNCTION(
form_data.query, prefix=RAG_EMBEDDING_QUERY_PREFIX, user=user
),
query_embedding=query_embedding,
k=form_data.k if form_data.k else request.app.state.config.TOP_K,
user=user,
)
@@ -2335,7 +2337,7 @@ class QueryCollectionsForm(BaseModel):
@router.post("/query/collection")
def query_collection_handler(
async def query_collection_handler(
request: Request,
form_data: QueryCollectionsForm,
user=Depends(get_verified_user),
@@ -2344,7 +2346,7 @@ def query_collection_handler(
if request.app.state.config.ENABLE_RAG_HYBRID_SEARCH and (
form_data.hybrid is None or form_data.hybrid
):
return query_collection_with_hybrid_search(
return await query_collection_with_hybrid_search(
collection_names=form_data.collection_names,
queries=[form_data.query],
embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
@@ -2379,7 +2381,7 @@ def query_collection_handler(
),
)
else:
return query_collection(
return await query_collection(
collection_names=form_data.collection_names,
queries=[form_data.query],
embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
@@ -2461,7 +2463,7 @@ if ENV == "dev":
@router.get("/ef/{text}")
async def get_embeddings(request: Request, text: Optional[str] = "Hello World!"):
return {
"result": request.app.state.EMBEDDING_FUNCTION(
"result": await request.app.state.EMBEDDING_FUNCTION(
text, prefix=RAG_EMBEDDING_QUERY_PREFIX
)
}