mirror of
https://github.com/open-webui/open-webui
synced 2025-06-09 07:56:42 +00:00
refac
This commit is contained in:
parent
c83e68282d
commit
2b7f9d14d0
@ -10,8 +10,8 @@ import inspect
|
|||||||
import uuid
|
import uuid
|
||||||
import asyncio
|
import asyncio
|
||||||
|
|
||||||
from fastapi import Request
|
from fastapi import Request, status
|
||||||
from starlette.responses import Response, StreamingResponse
|
from starlette.responses import Response, StreamingResponse, JSONResponse
|
||||||
|
|
||||||
|
|
||||||
from open_webui.models.users import UserModel
|
from open_webui.models.users import UserModel
|
||||||
@ -82,10 +82,6 @@ async def generate_direct_chat_completion(
|
|||||||
if form_data.get("stream"):
|
if form_data.get("stream"):
|
||||||
q = asyncio.Queue()
|
q = asyncio.Queue()
|
||||||
|
|
||||||
# Define a generator to stream responses
|
|
||||||
async def event_generator():
|
|
||||||
nonlocal q
|
|
||||||
|
|
||||||
async def message_listener(sid, data):
|
async def message_listener(sid, data):
|
||||||
"""
|
"""
|
||||||
Handle received socket messages and push them into the queue.
|
Handle received socket messages and push them into the queue.
|
||||||
@ -96,37 +92,6 @@ async def generate_direct_chat_completion(
|
|||||||
sio.on(channel, message_listener)
|
sio.on(channel, message_listener)
|
||||||
|
|
||||||
# Start processing chat completion in background
|
# Start processing chat completion in background
|
||||||
await event_emitter(
|
|
||||||
{
|
|
||||||
"type": "request:chat:completion",
|
|
||||||
"data": {
|
|
||||||
"form_data": form_data,
|
|
||||||
"model": models[form_data["model"]],
|
|
||||||
"channel": channel,
|
|
||||||
"session_id": session_id,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
)
|
|
||||||
|
|
||||||
try:
|
|
||||||
while True:
|
|
||||||
data = await q.get() # Wait for new messages
|
|
||||||
if isinstance(data, dict):
|
|
||||||
if "error" in data:
|
|
||||||
raise Exception(data["error"])
|
|
||||||
|
|
||||||
if "done" in data and data["done"]:
|
|
||||||
break # Stop streaming when 'done' is received
|
|
||||||
|
|
||||||
yield f"data: {json.dumps(data)}\n\n"
|
|
||||||
elif isinstance(data, str):
|
|
||||||
yield data
|
|
||||||
finally:
|
|
||||||
del sio.handlers["/"][channel] # Remove the listener
|
|
||||||
|
|
||||||
# Return the streaming response
|
|
||||||
return StreamingResponse(event_generator(), media_type="text/event-stream")
|
|
||||||
else:
|
|
||||||
res = await event_caller(
|
res = await event_caller(
|
||||||
{
|
{
|
||||||
"type": "request:chat:completion",
|
"type": "request:chat:completion",
|
||||||
@ -139,7 +104,51 @@ async def generate_direct_chat_completion(
|
|||||||
}
|
}
|
||||||
)
|
)
|
||||||
|
|
||||||
print(res)
|
print("res", res)
|
||||||
|
|
||||||
|
if res.get("status", False):
|
||||||
|
# Define a generator to stream responses
|
||||||
|
async def event_generator():
|
||||||
|
nonlocal q
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
data = await q.get() # Wait for new messages
|
||||||
|
if isinstance(data, dict):
|
||||||
|
if "done" in data and data["done"]:
|
||||||
|
break # Stop streaming when 'done' is received
|
||||||
|
|
||||||
|
yield f"data: {json.dumps(data)}\n\n"
|
||||||
|
elif isinstance(data, str):
|
||||||
|
yield data
|
||||||
|
except Exception as e:
|
||||||
|
log.debug(f"Error in event generator: {e}")
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Define a background task to run the event generator
|
||||||
|
async def background():
|
||||||
|
try:
|
||||||
|
del sio.handlers["/"][channel]
|
||||||
|
except Exception as e:
|
||||||
|
pass
|
||||||
|
|
||||||
|
# Return the streaming response
|
||||||
|
return StreamingResponse(
|
||||||
|
event_generator(), media_type="text/event-stream", background=background
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
raise Exception(str(res))
|
||||||
|
else:
|
||||||
|
res = await event_caller(
|
||||||
|
{
|
||||||
|
"type": "request:chat:completion",
|
||||||
|
"data": {
|
||||||
|
"form_data": form_data,
|
||||||
|
"model": models[form_data["model"]],
|
||||||
|
"channel": channel,
|
||||||
|
"session_id": session_id,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
if "error" in res:
|
if "error" in res:
|
||||||
raise Exception(res["error"])
|
raise Exception(res["error"])
|
||||||
|
@ -279,8 +279,17 @@
|
|||||||
OPENAI_API_URL
|
OPENAI_API_URL
|
||||||
);
|
);
|
||||||
|
|
||||||
if (res && res.ok) {
|
if (res) {
|
||||||
|
// raise if the response is not ok
|
||||||
|
if (!res.ok) {
|
||||||
|
throw await res.json();
|
||||||
|
}
|
||||||
|
|
||||||
if (form_data?.stream ?? false) {
|
if (form_data?.stream ?? false) {
|
||||||
|
cb({
|
||||||
|
status: true
|
||||||
|
});
|
||||||
|
|
||||||
// res will either be SSE or JSON
|
// res will either be SSE or JSON
|
||||||
const reader = res.body.getReader();
|
const reader = res.body.getReader();
|
||||||
const decoder = new TextDecoder();
|
const decoder = new TextDecoder();
|
||||||
@ -316,29 +325,12 @@
|
|||||||
}
|
}
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
console.error('chatCompletion', error);
|
console.error('chatCompletion', error);
|
||||||
|
cb(error);
|
||||||
if (form_data?.stream ?? false) {
|
|
||||||
$socket.emit(channel, {
|
|
||||||
error: error
|
|
||||||
});
|
|
||||||
} else {
|
|
||||||
cb({
|
|
||||||
error: error
|
|
||||||
});
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
} catch (error) {
|
} catch (error) {
|
||||||
console.error('chatCompletion', error);
|
console.error('chatCompletion', error);
|
||||||
if (form_data?.stream ?? false) {
|
cb(error);
|
||||||
$socket.emit(channel, {
|
|
||||||
error: error
|
|
||||||
});
|
|
||||||
} else {
|
|
||||||
cb({
|
|
||||||
error: error
|
|
||||||
});
|
|
||||||
}
|
|
||||||
} finally {
|
} finally {
|
||||||
$socket.emit(channel, {
|
$socket.emit(channel, {
|
||||||
done: true
|
done: true
|
||||||
|
Loading…
Reference in New Issue
Block a user