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/retrieval/utils.py
+224 -62
View File
@@ -3,6 +3,8 @@ import os
from typing import Optional, Union
import requests
import aiohttp
import asyncio
import hashlib
from concurrent.futures import ThreadPoolExecutor
import time
@@ -27,6 +29,7 @@ from open_webui.models.notes import Notes
from open_webui.retrieval.vector.main import GetResult
from open_webui.utils.access_control import has_access
from open_webui.utils.headers import include_user_info_headers
from open_webui.utils.misc import get_message_list
from open_webui.retrieval.web.utils import get_web_loader
@@ -87,15 +90,16 @@ class VectorSearchRetriever(BaseRetriever):
embedding_function: Any
top_k: int
def _get_relevant_documents(
async def _aget_relevant_documents(
self,
query: str,
*,
run_manager: CallbackManagerForRetrieverRun,
) -> list[Document]:
embedding = await self.embedding_function(query, RAG_EMBEDDING_QUERY_PREFIX)
result = VECTOR_DB_CLIENT.search(
collection_name=self.collection_name,
vectors=[self.embedding_function(query, RAG_EMBEDDING_QUERY_PREFIX)],
vectors=[embedding],
limit=self.top_k,
)
@@ -186,7 +190,7 @@ def get_enriched_texts(collection_result: GetResult) -> list[str]:
return enriched_texts
def query_doc_with_hybrid_search(
async def query_doc_with_hybrid_search(
collection_name: str,
collection_result: GetResult,
query: str,
@@ -262,7 +266,7 @@ def query_doc_with_hybrid_search(
base_compressor=compressor, base_retriever=ensemble_retriever
)
result = compression_retriever.invoke(query)
result = await compression_retriever.ainvoke(query)
distances = [d.metadata.get("score") for d in result]
documents = [d.page_content for d in result]
@@ -381,7 +385,7 @@ def get_all_items_from_collections(collection_names: list[str]) -> dict:
return merge_get_results(results)
def query_collection(
async def query_collection(
collection_names: list[str],
queries: list[str],
embedding_function,
@@ -406,7 +410,7 @@ def query_collection(
return None, e
# Generate all query embeddings (in one call)
query_embeddings = embedding_function(queries, prefix=RAG_EMBEDDING_QUERY_PREFIX)
query_embeddings = await embedding_function(queries, prefix=RAG_EMBEDDING_QUERY_PREFIX)
log.debug(
f"query_collection: processing {len(queries)} queries across {len(collection_names)} collections"
)
@@ -433,7 +437,7 @@ def query_collection(
return merge_and_sort_query_results(results, k=k)
def query_collection_with_hybrid_search(
async def query_collection_with_hybrid_search(
collection_names: list[str],
queries: list[str],
embedding_function,
@@ -465,9 +469,9 @@ def query_collection_with_hybrid_search(
f"Starting hybrid search for {len(queries)} queries in {len(collection_names)} collections..."
)
def process_query(collection_name, query):
async def process_query(collection_name, query):
try:
result = query_doc_with_hybrid_search(
result = await query_doc_with_hybrid_search(
collection_name=collection_name,
collection_result=collection_results[collection_name],
query=query,
@@ -493,9 +497,8 @@ def query_collection_with_hybrid_search(
for q in queries
]
with ThreadPoolExecutor() as executor:
future_results = [executor.submit(process_query, cn, q) for cn, q in tasks]
task_results = [future.result() for future in future_results]
# Run all queries in parallel using asyncio.gather
task_results = await asyncio.gather(*[process_query(cn, q) for cn, q in tasks])
for result, err in task_results:
if err is not None:
@@ -511,6 +514,146 @@ def query_collection_with_hybrid_search(
return merge_and_sort_query_results(results, k=k)
async def generate_openai_batch_embeddings_async(
model: str,
texts: list[str],
url: str = "https://api.openai.com/v1",
key: str = "",
prefix: str = None,
user: UserModel = None,
) -> Optional[list[list[float]]]:
try:
log.debug(
f"generate_openai_batch_embeddings_async:model {model} batch size: {len(texts)}"
)
form_data = {"input": texts, "model": model}
if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str):
form_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {key}",
}
if ENABLE_FORWARD_USER_INFO_HEADERS and user:
headers = include_user_info_headers(headers, user)
async with aiohttp.ClientSession() as session:
async with session.post(f"{url}/embeddings", headers=headers, json=form_data) as r:
r.raise_for_status()
data = await r.json()
if "data" in data:
return [item["embedding"] for item in data["data"]]
else:
raise Exception("Something went wrong :/")
except Exception as e:
log.exception(f"Error generating openai batch embeddings: {e}")
return None
async def generate_azure_openai_batch_embeddings_async(
model: str,
texts: list[str],
url: str,
key: str = "",
version: str = "",
prefix: str = None,
user: UserModel = None,
) -> Optional[list[list[float]]]:
try:
log.debug(
f"generate_azure_openai_batch_embeddings_async:deployment {model} batch size: {len(texts)}"
)
form_data = {"input": texts}
if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str):
form_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
full_url = f"{url}/openai/deployments/{model}/embeddings?api-version={version}"
headers = {
"Content-Type": "application/json",
"api-key": key,
}
if ENABLE_FORWARD_USER_INFO_HEADERS and user:
headers = include_user_info_headers(headers, user)
async with aiohttp.ClientSession() as session:
async with session.post(full_url, headers=headers, json=form_data) as r:
r.raise_for_status()
data = await r.json()
if "data" in data:
return [item["embedding"] for item in data["data"]]
else:
raise Exception("Something went wrong :/")
except Exception as e:
log.exception(f"Error generating azure openai batch embeddings: {e}")
return None
async def generate_ollama_batch_embeddings_async(
model: str,
texts: list[str],
url: str,
key: str = "",
prefix: str = None,
user: UserModel = None,
) -> Optional[list[list[float]]]:
try:
log.debug(
f"generate_ollama_batch_embeddings_async:model {model} batch size: {len(texts)}"
)
form_data = {"input": texts, "model": model}
if isinstance(RAG_EMBEDDING_PREFIX_FIELD_NAME, str) and isinstance(prefix, str):
form_data[RAG_EMBEDDING_PREFIX_FIELD_NAME] = prefix
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {key}",
}
if ENABLE_FORWARD_USER_INFO_HEADERS and user:
headers = include_user_info_headers(headers, user)
async with aiohttp.ClientSession() as session:
async with session.post(f"{url}/api/embed", headers=headers, json=form_data) as r:
r.raise_for_status()
data = await r.json()
if "embeddings" in data:
return data["embeddings"]
else:
raise Exception("Something went wrong :/")
except Exception as e:
log.exception(f"Error generating ollama batch embeddings: {e}")
return None
async def generate_multiple_async(query, prefix, user, func_async, embedding_batch_size):
if isinstance(query, list):
# Create batches
batches = [
query[i : i + embedding_batch_size]
for i in range(0, len(query), embedding_batch_size)
]
log.debug(f"generate_multiple_async: Processing {len(batches)} batches in parallel")
# Execute all batches in parallel
tasks = [
func_async(batch, prefix=prefix, user=user)
for batch in batches
]
batch_results = await asyncio.gather(*tasks)
# Flatten results
embeddings = []
for batch_embeddings in batch_results:
if isinstance(batch_embeddings, list):
embeddings.extend(batch_embeddings)
log.debug(f"generate_multiple_async: Generated {len(embeddings)} embeddings from {len(batches)} parallel batches")
return embeddings
else:
return await func_async([query], prefix=prefix, user=user)
def get_embedding_function(
embedding_engine,
embedding_model,
@@ -521,40 +664,58 @@ def get_embedding_function(
azure_api_version=None,
):
if embedding_engine == "":
return lambda query, prefix=None, user=None: embedding_function.encode(
query, **({"prompt": prefix} if prefix else {})
).tolist()
# Sentence transformers: CPU-bound sync operation
# Run in thread pool to prevent blocking async event loop
def sync_encode(query, prefix=None):
return embedding_function.encode(
query, **({"prompt": prefix} if prefix else {})
).tolist()
async def async_wrapper(query, prefix=None, user=None):
# Run CPU-bound operation in thread pool
return await asyncio.to_thread(sync_encode, query, prefix)
return async_wrapper
elif embedding_engine in ["ollama", "openai", "azure_openai"]:
func = lambda query, prefix=None, user=None: generate_embeddings(
engine=embedding_engine,
model=embedding_model,
text=query,
prefix=prefix,
url=url,
key=key,
user=user,
azure_api_version=azure_api_version,
)
# Create async function based on engine
if embedding_engine == "openai":
func_async = lambda texts, prefix=None, user=None: generate_openai_batch_embeddings_async(
model=embedding_model,
texts=texts,
url=url,
key=key,
prefix=prefix,
user=user,
)
elif embedding_engine == "azure_openai":
func_async = lambda texts, prefix=None, user=None: generate_azure_openai_batch_embeddings_async(
model=embedding_model,
texts=texts,
url=url,
key=key,
version=azure_api_version,
prefix=prefix,
user=user,
)
elif embedding_engine == "ollama":
func_async = lambda texts, prefix=None, user=None: generate_ollama_batch_embeddings_async(
model=embedding_model,
texts=texts,
url=url,
key=key,
prefix=prefix,
user=user,
)
def generate_multiple(query, prefix, user, func):
# Return async function for parallel batch processing (FastAPI-compatible)
async def embedding_wrapper(query, prefix=None, user=None):
if isinstance(query, list):
embeddings = []
for i in range(0, len(query), embedding_batch_size):
batch_embeddings = func(
query[i : i + embedding_batch_size],
prefix=prefix,
user=user,
)
if isinstance(batch_embeddings, list):
embeddings.extend(batch_embeddings)
return embeddings
return await generate_multiple_async(query, prefix, user, func_async, embedding_batch_size)
else:
return func(query, prefix, user)
result = await func_async([query], prefix=prefix, user=user)
return result[0] if result else None
return lambda query, prefix=None, user=None: generate_multiple(
query, prefix, user, func
)
return embedding_wrapper
else:
raise ValueError(f"Unknown embedding engine: {embedding_engine}")
@@ -572,7 +733,7 @@ def get_reranking_function(reranking_engine, reranking_model, reranking_function
)
def get_sources_from_items(
async def get_sources_from_items(
request,
items,
queries,
@@ -800,7 +961,7 @@ def get_sources_from_items(
query_result = None # Initialize to None
if hybrid_search:
try:
query_result = query_collection_with_hybrid_search(
query_result = await query_collection_with_hybrid_search(
collection_names=collection_names,
queries=queries,
embedding_function=embedding_function,
@@ -818,7 +979,7 @@ def get_sources_from_items(
# fallback to non-hybrid search
if not hybrid_search and query_result is None:
query_result = query_collection(
query_result = await query_collection(
collection_names=collection_names,
queries=queries,
embedding_function=embedding_function,
@@ -1056,26 +1217,25 @@ def generate_embeddings(
else:
text = f"{prefix}{text}"
# Run async embedding generation synchronously
if engine == "ollama":
embeddings = generate_ollama_batch_embeddings(
**{
"model": model,
"texts": text if isinstance(text, list) else [text],
"url": url,
"key": key,
"prefix": prefix,
"user": user,
}
)
embeddings = asyncio.run(generate_ollama_batch_embeddings_async(
model=model,
texts=text if isinstance(text, list) else [text],
url=url,
key=key,
prefix=prefix,
user=user,
))
return embeddings[0] if isinstance(text, str) else embeddings
elif engine == "openai":
embeddings = generate_openai_batch_embeddings(
embeddings = asyncio.run(generate_openai_batch_embeddings_async(
model, text if isinstance(text, list) else [text], url, key, prefix, user
)
))
return embeddings[0] if isinstance(text, str) else embeddings
elif engine == "azure_openai":
azure_api_version = kwargs.get("azure_api_version", "")
embeddings = generate_azure_openai_batch_embeddings(
embeddings = asyncio.run(generate_azure_openai_batch_embeddings_async(
model,
text if isinstance(text, list) else [text],
url,
@@ -1083,7 +1243,7 @@ def generate_embeddings(
azure_api_version,
prefix,
user,
)
))
return embeddings[0] if isinstance(text, str) else embeddings
@@ -1104,7 +1264,7 @@ class RerankCompressor(BaseDocumentCompressor):
extra = "forbid"
arbitrary_types_allowed = True
def compress_documents(
async def acompress_documents(
self,
documents: Sequence[Document],
query: str,
@@ -1114,12 +1274,14 @@ class RerankCompressor(BaseDocumentCompressor):
scores = None
if reranking:
scores = self.reranking_function(query, documents)
scores = self.reranking_function(
[(query, doc.page_content) for doc in documents]
)
else:
from sentence_transformers import util
query_embedding = self.embedding_function(query, RAG_EMBEDDING_QUERY_PREFIX)
document_embedding = self.embedding_function(
query_embedding = await self.embedding_function(query, RAG_EMBEDDING_QUERY_PREFIX)
document_embedding = await self.embedding_function(
[doc.page_content for doc in documents], RAG_EMBEDDING_CONTENT_PREFIX
)
scores = util.cos_sim(query_embedding, document_embedding)[0]