OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-26 18:26:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

refac: task ollama stream support

Browse Source
This commit is contained in:
Timothy J. Baek
2024-09-21 01:07:57 +02:00
parent 41926172d3
commit 3a0a1aca11
3 changed files with 69 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
backend/open_webui/utils/misc.py
View File

@@ -105,17 +105,25 @@ def openai_chat_message_template(model: str):
}
def openai_chat_chunk_message_template(model: str, message: str) -> dict:
def openai_chat_chunk_message_template(
model: str, message: Optional[str] = None
) -> dict:
template = openai_chat_message_template(model)
template["object"] = "chat.completion.chunk"
template["choices"][0]["delta"] = {"content": message}
if message:
template["choices"][0]["delta"] = {"content": message}
else:
template["choices"][0]["finish_reason"] = "stop"
return template
def openai_chat_completion_message_template(model: str, message: str) -> dict:
def openai_chat_completion_message_template(
model: str, message: Optional[str] = None
) -> dict:
template = openai_chat_message_template(model)
template["object"] = "chat.completion"
template["choices"][0]["message"] = {"content": message, "role": "assistant"}
if message:
template["choices"][0]["message"] = {"content": message, "role": "assistant"}
template["choices"][0]["finish_reason"] = "stop"
return template

24
backend/open_webui/utils/response.py
View File

@@ -1,10 +1,9 @@
from open_webui.utils.task import prompt_template
import json
from open_webui.utils.misc import (
openai_chat_chunk_message_template,
openai_chat_completion_message_template,
)
from typing import Callable, Optional
def convert_response_ollama_to_openai(ollama_response: dict) -> dict:
model = ollama_response.get("model", "ollama")
@@ -12,3 +11,22 @@ def convert_response_ollama_to_openai(ollama_response: dict) -> dict:
response = openai_chat_completion_message_template(model, message_content)
return response
async def convert_streaming_response_ollama_to_openai(ollama_streaming_response):
async for data in ollama_streaming_response.body_iterator:
data = json.loads(data)
model = data.get("model", "ollama")
message_content = data.get("message", {}).get("content", "")
done = data.get("done", False)
data = openai_chat_chunk_message_template(
model, message_content if not done else None
)
line = f"data: {json.dumps(data)}\n\n"
if done:
line += "data: [DONE]\n\n"
yield line