OpenWebUI/pipelines
1
0
Fork 0
mirror of https://github.com/open-webui/pipelines synced 2025-05-23 05:24:09 +00:00
Code Issues Actions 1 Packages Projects Releases Wiki Activity

fix

Browse Source
This commit is contained in:
Timothy J. Baek 2024-05-21 18:47:46 -07:00
parent 56239abd6b
commit 39da1c982a
2 changed files with 6 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
config.py
View File

@ -10,7 +10,3 @@ try:
load_dotenv(find_dotenv("./.env")) load_dotenv(find_dotenv("./.env"))
except ImportError: except ImportError:
print("dotenv not installed, skipping...") print("dotenv not installed, skipping...")
MODEL_ID = os.environ.get("MODEL_ID", "plugin-id")
MODEL_NAME = os.environ.get("MODEL_NAME", "Plugin Model")

9
main.py
View File

@ -14,7 +14,6 @@ import uuid
from utils import get_last_user_message, stream_message_template from utils import get_last_user_message, stream_message_template
from schemas import OpenAIChatCompletionForm from schemas import OpenAIChatCompletionForm
from config import MODEL_ID, MODEL_NAME
import os import os
import importlib.util import importlib.util
@ -61,6 +60,7 @@ async def lifespan(app: FastAPI):
yield yield
for pipeline in PIPELINES.values(): for pipeline in PIPELINES.values():
if hasattr(pipeline["module"], "on_shutdown"): if hasattr(pipeline["module"], "on_shutdown"):
await pipeline["module"].on_shutdown() await pipeline["module"].on_shutdown()
@ -124,6 +124,7 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
) )
def job(): def job():
print(form_data.model)
get_response = app.state.PIPELINES[form_data.model]["module"].get_response get_response = app.state.PIPELINES[form_data.model]["module"].get_response
if form_data.stream: if form_data.stream:
@ -131,6 +132,8 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
def stream_content(): def stream_content():
res = get_response(user_message, messages=form_data.messages) res = get_response(user_message, messages=form_data.messages)
print(res)
if isinstance(res, str): if isinstance(res, str):
message = stream_message_template(res) message = stream_message_template(res)
yield f"data: {json.dumps(message)}\n\n" yield f"data: {json.dumps(message)}\n\n"
@ -145,7 +148,7 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
"id": f"{form_data.model}-{str(uuid.uuid4())}", "id": f"{form_data.model}-{str(uuid.uuid4())}",
"object": "chat.completion.chunk", "object": "chat.completion.chunk",
"created": int(time.time()), "created": int(time.time()),
"model": MODEL_ID, "model": form_data.model,
"choices": [ "choices": [
{ {
"index": 0, "index": 0,
@ -175,7 +178,7 @@ async def generate_openai_chat_completion(form_data: OpenAIChatCompletionForm):
"id": f"{form_data.model}-{str(uuid.uuid4())}", "id": f"{form_data.model}-{str(uuid.uuid4())}",
"object": "chat.completion", "object": "chat.completion",
"created": int(time.time()), "created": int(time.time()),
"model": MODEL_ID, "model": form_data.model,
"choices": [ "choices": [
{ {
"index": 0, "index": 0,