UniqAI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-01-18 08:40:58 +00:00
Code Issues Actions 5 Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy J. Baek 2024-07-11 10:40:10 -07:00
parent a62b0924df
commit 0ef27bfc5e
2 changed files with 52 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

31
backend/apps/socket/main.py
View File

@ -137,3 +137,34 @@ async def disconnect(sid):
await sio.emit("user-count", {"count": len(USER_POOL)})
else:
print(f"Unknown session ID {sid} disconnected")
async def get_event_emitter(request_info):
async def __event_emitter__(event_data):
await sio.emit(
"chat-events",
{
"chat_id": request_info["chat_id"],
"message_id": request_info["id"],
"data": event_data,
},
to=request_info["session_id"],
)
return __event_emitter__
async def get_event_call(request_info):
async def __event_call__(event_data):
response = await sio.call(
"chat-events",
{
"chat_id": request_info["chat_id"],
"message_id": request_info["id"],
"data": event_data,
},
to=request_info["session_id"],
)
return response
return __event_call__

57
backend/main.py
View File

@ -29,7 +29,7 @@ from starlette.middleware.sessions import SessionMiddleware
from starlette.responses import StreamingResponse, Response, RedirectResponse
from apps.socket.main import sio, app as socket_app
from apps.socket.main import sio, app as socket_app, get_event_emitter, get_event_call
from apps.ollama.main import (
app as ollama_app,
get_all_models as get_ollama_models,
@ -632,24 +632,12 @@ class ChatCompletionMiddleware(BaseHTTPMiddleware):
message_id = body["id"]
del body["id"]
async def __event_emitter__(data):
await sio.emit(
"chat-events",
{
"chat_id": chat_id,
"message_id": message_id,
"data": data,
},
to=session_id,
)
async def __event_call__(data):
response = await sio.call(
"chat-events",
{"chat_id": chat_id, "message_id": message_id, "data": data},
to=session_id,
)
return response
__event_emitter__ = await get_event_emitter(
{"chat_id": chat_id, "message_id": message_id, "session_id": session_id}
)
__event_call__ = await get_event_call(
{"chat_id": chat_id, "message_id": message_id, "session_id": session_id}
)
# Initialize data_items to store additional data to be sent to the client
data_items = []
@ -1107,24 +1095,21 @@ async def chat_completed(form_data: dict, user=Depends(get_verified_user)):
else:
pass
async def __event_emitter__(event_data):
await sio.emit(
"chat-events",
{
"chat_id": data["chat_id"],
"message_id": data["id"],
"data": event_data,
},
to=data["session_id"],
)
__event_emitter__ = await get_event_emitter(
{
"chat_id": data["chat_id"],
"message_id": data["id"],
"session_id": data["session_id"],
}
)
async def __event_call__(event_data):
response = await sio.call(
"chat-events",
{"chat_id": data["chat_id"], "message_id": data["id"], "data": event_data},
to=data["session_id"],
)
return response
__event_call__ = await get_event_call(
{
"chat_id": data["chat_id"],
"message_id": data["id"],
"session_id": data["session_id"],
}
)
def get_priority(function_id):
function = Functions.get_function_by_id(function_id)