perf: 50x performance improvement for external embeddings (#19296)
* 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:
@@ -993,37 +993,32 @@ async def chat_completion_files_handler(
|
||||
queries = [get_last_user_message(body["messages"])]
|
||||
|
||||
try:
|
||||
# Offload get_sources_from_items to a separate thread
|
||||
loop = asyncio.get_running_loop()
|
||||
with ThreadPoolExecutor() as executor:
|
||||
sources = await loop.run_in_executor(
|
||||
executor,
|
||||
lambda: get_sources_from_items(
|
||||
request=request,
|
||||
items=files,
|
||||
queries=queries,
|
||||
embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
|
||||
query, prefix=prefix, user=user
|
||||
),
|
||||
k=request.app.state.config.TOP_K,
|
||||
reranking_function=(
|
||||
(
|
||||
lambda query, documents: request.app.state.RERANKING_FUNCTION(
|
||||
query, documents, user=user
|
||||
)
|
||||
)
|
||||
if request.app.state.RERANKING_FUNCTION
|
||||
else None
|
||||
),
|
||||
k_reranker=request.app.state.config.TOP_K_RERANKER,
|
||||
r=request.app.state.config.RELEVANCE_THRESHOLD,
|
||||
hybrid_bm25_weight=request.app.state.config.HYBRID_BM25_WEIGHT,
|
||||
hybrid_search=request.app.state.config.ENABLE_RAG_HYBRID_SEARCH,
|
||||
full_context=all_full_context
|
||||
or request.app.state.config.RAG_FULL_CONTEXT,
|
||||
user=user,
|
||||
),
|
||||
)
|
||||
# Directly await async get_sources_from_items (no thread needed - fully async now)
|
||||
sources = await get_sources_from_items(
|
||||
request=request,
|
||||
items=files,
|
||||
queries=queries,
|
||||
embedding_function=lambda query, prefix: request.app.state.EMBEDDING_FUNCTION(
|
||||
query, prefix=prefix, user=user
|
||||
),
|
||||
k=request.app.state.config.TOP_K,
|
||||
reranking_function=(
|
||||
(
|
||||
lambda query, documents: request.app.state.RERANKING_FUNCTION(
|
||||
query, documents, user=user
|
||||
)
|
||||
)
|
||||
if request.app.state.RERANKING_FUNCTION
|
||||
else None
|
||||
),
|
||||
k_reranker=request.app.state.config.TOP_K_RERANKER,
|
||||
r=request.app.state.config.RELEVANCE_THRESHOLD,
|
||||
hybrid_bm25_weight=request.app.state.config.HYBRID_BM25_WEIGHT,
|
||||
hybrid_search=request.app.state.config.ENABLE_RAG_HYBRID_SEARCH,
|
||||
full_context=all_full_context
|
||||
or request.app.state.config.RAG_FULL_CONTEXT,
|
||||
user=user,
|
||||
)
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user