ti.sil/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2025-12-15 13:55:19 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

refac/enh: async process chat handling

Browse source
This commit is contained in:
Timothy Jaeryang Baek 2025-08-19 01:24:53 +04:00
parent 4bc77b544e
commit d6f709574e
2 changed files with 66 additions and 58 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

115
backend/open_webui/main.py
View file

@ -57,6 +57,7 @@ from open_webui.utils.logger import start_logger
from open_webui.socket.main import ( from open_webui.socket.main import (
app as socket_app, app as socket_app,
periodic_usage_pool_cleanup, periodic_usage_pool_cleanup,
get_event_emitter,
get_models_in_use, get_models_in_use,
get_active_user_ids, get_active_user_ids,
) )
@ -466,6 +467,7 @@ from open_webui.utils.redis import get_redis_connection
from open_webui.tasks import ( from open_webui.tasks import (
redis_task_command_listener, redis_task_command_listener,
list_task_ids_by_item_id, list_task_ids_by_item_id,
create_task,
stop_task, stop_task,
list_tasks, list_tasks,
) # Import from tasks.py ) # Import from tasks.py
@ -1473,65 +1475,78 @@ async def chat_completion(
request.state.metadata = metadata request.state.metadata = metadata
form_data["metadata"] = metadata form_data["metadata"] = metadata
form_data, metadata, events = await process_chat_payload(
request, form_data, user, metadata, model
)
except Exception as e: except Exception as e:
log.debug(f"Error processing chat payload: {e}") log.debug(f"Error processing chat metadata: {e}")
if metadata.get("chat_id") and metadata.get("message_id"):
# Update the chat message with the error
try:
Chats.upsert_message_to_chat_by_id_and_message_id(
metadata["chat_id"],
metadata["message_id"],
{
"error": {"content": str(e)},
},
)
except:
pass
raise HTTPException( raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST, status_code=status.HTTP_400_BAD_REQUEST,
detail=str(e), detail=str(e),
) )
try: async def process_chat(request, form_data, user, metadata, model):
response = await chat_completion_handler(request, form_data, user) try:
if metadata.get("chat_id") and metadata.get("message_id"): form_data, metadata, events = await process_chat_payload(
try: request, form_data, user, metadata, model
Chats.upsert_message_to_chat_by_id_and_message_id( )
metadata["chat_id"],
metadata["message_id"],
{
"model": model_id,
},
)
except:
pass
return await process_chat_response( response = await chat_completion_handler(request, form_data, user)
request, response, form_data, user, metadata, model, events, tasks if metadata.get("chat_id") and metadata.get("message_id"):
) try:
except Exception as e: Chats.upsert_message_to_chat_by_id_and_message_id(
log.debug(f"Error in chat completion: {e}") metadata["chat_id"],
if metadata.get("chat_id") and metadata.get("message_id"): metadata["message_id"],
# Update the chat message with the error {
try: "model": model_id,
Chats.upsert_message_to_chat_by_id_and_message_id( },
metadata["chat_id"], )
metadata["message_id"], except:
{ pass
"error": {"content": str(e)},
},
)
except:
pass
raise HTTPException( return await process_chat_response(
status_code=status.HTTP_400_BAD_REQUEST, request, response, form_data, user, metadata, model, events, tasks
detail=str(e), )
except asyncio.CancelledError:
log.info("Chat processing was cancelled")
try:
event_emitter = get_event_emitter(metadata)
await event_emitter(
{"type": "task-cancelled"},
)
except Exception as e:
pass
except Exception as e:
log.debug(f"Error processing chat payload: {e}")
if metadata.get("chat_id") and metadata.get("message_id"):
# Update the chat message with the error
try:
Chats.upsert_message_to_chat_by_id_and_message_id(
metadata["chat_id"],
metadata["message_id"],
{
"error": {"content": str(e)},
},
)
except:
pass
raise HTTPException(
status_code=status.HTTP_400_BAD_REQUEST,
detail=str(e),
)
if (
metadata.get("session_id")
and metadata.get("chat_id")
and metadata.get("message_id")
):
# Asynchronous Chat Processing
task_id, _ = await create_task(
request.app.state.redis,
process_chat(request, form_data, user, metadata, model),
id=metadata["chat_id"],
) )
return {"status": True, "task_id": task_id}
else:
return await process_chat(request, form_data, user, metadata, model)
# Alias for chat_completion (Legacy) # Alias for chat_completion (Legacy)

9
backend/open_webui/utils/middleware.py
View file

@ -86,7 +86,6 @@ from open_webui.utils.filter import (
from open_webui.utils.code_interpreter import execute_code_jupyter from open_webui.utils.code_interpreter import execute_code_jupyter
from open_webui.utils.payload import apply_model_system_prompt_to_body from open_webui.utils.payload import apply_model_system_prompt_to_body
from open_webui.tasks import create_task
from open_webui.config import ( from open_webui.config import (
CACHE_DIR, CACHE_DIR,
@ -2600,13 +2599,7 @@ async def process_chat_response(
if response.background is not None: if response.background is not None:
await response.background() await response.background()
# background_tasks.add_task(response_handler, response, events) return await response_handler(response, events)
task_id, _ = await create_task(
request.app.state.redis,
response_handler(response, events),
id=metadata["chat_id"],
)
return {"status": True, "task_id": task_id}
else: else:
# Fallback to the original response # Fallback to the original response