Fully support Open-WebUI

Co-authored-by: Tim Baek <tim@openwebui.com>

Author: vmesel

docker-compose-open-webui.yml

@@ -0,0 +1,58 @@
+version: '3.3'
+services:
+  db:
+    image: pgvector/pgvector:pg15
+    restart: always
+    ports:
+      - '5432:5432'
+    environment:
+      POSTGRES_USER: talkdai
+      POSTGRES_PASSWORD: talkdai
+      POSTGRES_DB: talkdai
+    volumes:
+       - ./data/db:/var/lib/postgresql/data
+       - ./psql/db-ext-vector.sql:/docker-entrypoint-initdb.d/db-ext-vector.sql
+    healthcheck:
+      test: ["CMD", "pg_isready", "-d", "talkdai", "-U", "talkdai"]
+      interval: 10s
+      timeout: 5s
+      retries: 5
+  dialog:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    stdin_open: true
+    tty: true
+    volumes:
+      - ./:/app
+      - ./.empty:/app/data/db
+      - ./static:/app/static
+      - ./sample_data:/app/src/sample_data
+    ports:
+      - '8000:8000'
+    depends_on:
+      db:
+        condition: service_healthy
+    environment:
+      - DATABASE_URL=postgresql://talkdai:talkdai@db:5432/talkdai
+      - STATIC_FILE_LOCATION=/app/static
+    env_file:
+      - .env
+  openwebui:
+    image: ghcr.io/open-webui/open-webui:main
+    ports:
+      - '3000:8080'
+    environment:
+      - OPENAI_API_KEYS=FAKE-KEY;
+      - OPENAI_API_BASE_URLS=http://dialog:8000/openai;
+      - ENABLE_OPENAI_API=true
+    volumes:
+      - open-webui:/app/backend/data
+    depends_on:
+      db:
+        condition: service_healthy
+      dialog:
+        condition: service_started
+
+volumes:
+  open-webui:

docker-compose.yml

@@ -38,21 +38,6 @@ services:
       - STATIC_FILE_LOCATION=/app/static
     env_file:
       - .env
-  openwebui:
-    image: ghcr.io/open-webui/open-webui:main
-    ports:
-      - '3000:8080'
-    environment:
-      - OPENAI_API_KEYS=FAKE-KEY;
-      - OPENAI_API_BASE_URLS=http://dialog:8000/openai;
-      - ENABLE_OPENAI_API=true
-    volumes:
-      - open-webui:/app/backend/data
-    depends_on:
-      db:
-        condition: service_healthy
-      dialog:
-        condition: service_started
 
 volumes:
   open-webui:

src/dialog/routers/openai.py

@@ -5,12 +5,16 @@
 
 from dialog.db import engine, get_session
 from dialog_lib.db.models import Chat as ChatEntity, ChatMessages
-from dialog.schemas import OpenAIChat, OpenAIChatCompletion, OpenAIModel
+from dialog.schemas import (
+    OpenAIChat, OpenAIChatCompletion, OpenAIModel, OpenAIMessage,
+    OpenAIStreamChoice, OpenAIStreamSchema
+)
 from dialog.llm import process_user_message
 
 from sqlalchemy.orm import Session
 
 from fastapi import APIRouter, Depends
+from fastapi.responses import StreamingResponse
 
 open_ai_api_router = APIRouter()
 
@@ -49,27 +53,48 @@ async def ask_question_to_llm(message: OpenAIChat, session: Session = Depends(ge
 
     duration = datetime.datetime.now() - start_time
     logging.info(f"Request processing time: {duration}")
-    chat_completion = OpenAIChatCompletion(
-        choices=[
-            {
-                "finish_reason": "stop",
-                "index": 0,
-                "message": {
-                    "content": ai_message["text"],
-                    "role": "assistant"
-                },
-                "logprobs": None
+    generated_message = ai_message["text"]
+    if not message.stream:
+        chat_completion = OpenAIChatCompletion(
+            choices=[
+                {
+                    "finish_reason": "stop",
+                    "index": 0,
+                    "message": OpenAIMessage(**{
+                        "content": generated_message,
+                        "role": "assistant"
+                    }),
+                    "logprobs": None
+                }
+            ],
+            created=int(datetime.datetime.now().timestamp()),
+            id=f"talkdai-{str(uuid4())}",
+            model="talkd-ai",
+            object="chat.completion",
+            usage={
+                "completion_tokens": None,
+                "prompt_tokens": None,
+                "total_tokens": None
             }
-        ],
-        created=int(datetime.datetime.now().timestamp()),
-        id=f"talkdai-{str(uuid4())}",
-        model="talkd-ai",
-        object="chat.completion",
-        usage={
-            "completion_tokens": None,
-            "prompt_tokens": None,
-            "total_tokens": None
-        }
-    )
-    logging.info(f"Chat completion: {chat_completion}")
-    return chat_completion
+        )
+        logging.info(f"Chat completion: {chat_completion}")
+        return chat_completion
+
+    def gen():
+        for word in f"{generated_message} +END".split():
+            # Yield Streaming Response on each word
+            message_part = OpenAIStreamChoice(
+                index=0,
+                delta={
+                    "content": f"{word} "
+                } if word != "+END" else {}
+            )
+
+            message_stream = OpenAIStreamSchema(
+                id=f"talkdai-{str(uuid4())}",
+                choices=[message_part]
+            )
+            logging.info(f"data: {message_stream.model_dump_json()}")
+            yield f"data: {message_stream.model_dump_json()}\n\n"
+
+    return StreamingResponse(gen(), media_type='text/event-stream')

src/dialog/schemas/openai_schemas.py

@@ -1,3 +1,4 @@
+import datetime
 from typing import List, Optional
 from pydantic import BaseModel, ConfigDict
 
@@ -9,6 +10,7 @@ class OpenAIMessage(BaseModel):
 class OpenAIChat(BaseModel):
     model: str
     messages: List[OpenAIMessage]
+    stream: bool = False
 
 
 class OpenAIChoices(BaseModel):
@@ -28,7 +30,7 @@ class OpenAIChatCompletion(BaseModel):
     choices: List[OpenAIChoices]
     created: float
     id: str
-    model: str = "talkdai"
+    model: str = "talkd-ai"
     object: str = "chat.completion"
     usage: OpenAIUsageDict
 
@@ -37,4 +39,19 @@ class OpenAIModel(BaseModel):
     id: str
     object: str
     created: int
-    owned_by: str
+    owned_by: str
+
+
+class OpenAIStreamChoice(BaseModel):
+    index: int
+    delta: dict
+    logprobs: Optional[str] = None
+    finish_reason: Optional[str] = None
+
+class OpenAIStreamSchema(BaseModel):
+    id: str
+    object: str = "chat.completion.chunk"
+    created: int = int(datetime.datetime.now().timestamp())
+    model: str = "talkd-ai"
+    system_fingerprint: str = None
+    choices: List[OpenAIStreamChoice]

src/tests/test_views.py

@@ -58,7 +58,7 @@ def test_customized_openai_models_response(client):
     for i in ["id", "object", "created", "owned_by"]:
         assert i in response.json()[0]
 
-def test_customized_openai_chat_completion_response(client, llm_mock_openai_router):
+def test_customized_openai_chat_completion_response_stream_false(client, llm_mock_openai_router):
     os.environ["LLM_CLASS"] = "dialog.llm.agents.default.DialogLLM"
     response = client.post("/openai/chat/completions", json={
         "model": "talkdai",
@@ -67,7 +67,8 @@ def test_customized_openai_chat_completion_response(client, llm_mock_openai_rout
                 "role": "user",
                 "content": "Hello"
             }
-        ]
+        ],
+        "stream": False
     })
     assert response.status_code == 200
     for i in ["choices", "created", "id", "model", "object", "usage"]: