Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in / Register
Toggle navigation
Menu
Open sidebar
chenpangpang
open-webui
Commits
23e69bcd
Commit
23e69bcd
authored
Jul 24, 2024
by
Timothy J. Baek
Browse files
enh: AsyncGenerator support
parent
edff071c
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
24 additions
and
2 deletions
+24
-2
backend/apps/webui/main.py
backend/apps/webui/main.py
+24
-2
No files found.
backend/apps/webui/main.py
View file @
23e69bcd
...
...
@@ -54,7 +54,7 @@ import uuid
import
time
import
json
from
typing
import
Iterator
,
Generator
,
Optional
from
typing
import
Iterator
,
Generator
,
AsyncGenerator
,
Optional
from
pydantic
import
BaseModel
app
=
FastAPI
()
...
...
@@ -411,6 +411,25 @@ async def generate_function_chat_completion(form_data, user):
yield
f
"data:
{
json
.
dumps
(
finish_message
)
}
\n\n
"
yield
f
"data: [DONE]"
if
isinstance
(
res
,
AsyncGenerator
):
async
for
line
in
res
:
if
isinstance
(
line
,
BaseModel
):
line
=
line
.
model_dump_json
()
line
=
f
"data:
{
line
}
"
if
isinstance
(
line
,
dict
):
line
=
f
"data:
{
json
.
dumps
(
line
)
}
"
try
:
line
=
line
.
decode
(
"utf-8"
)
except
:
pass
if
line
.
startswith
(
"data:"
):
yield
f
"
{
line
}
\n\n
"
else
:
line
=
stream_message_template
(
form_data
[
"model"
],
line
)
yield
f
"data:
{
json
.
dumps
(
line
)
}
\n\n
"
return
StreamingResponse
(
stream_content
(),
media_type
=
"text/event-stream"
)
else
:
...
...
@@ -434,9 +453,12 @@ async def generate_function_chat_completion(form_data, user):
message
=
""
if
isinstance
(
res
,
str
):
message
=
res
if
isinstance
(
res
,
Generator
):
el
if
isinstance
(
res
,
Generator
):
for
stream
in
res
:
message
=
f
"
{
message
}{
stream
}
"
elif
isinstance
(
res
,
AsyncGenerator
):
async
for
stream
in
res
:
message
=
f
"
{
message
}{
stream
}
"
return
{
"id"
:
f
"
{
form_data
[
'model'
]
}
-
{
str
(
uuid
.
uuid4
())
}
"
,
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment