implement streaming parameter

Author: user786875

pyproject.toml

@@ -21,13 +21,13 @@ python-multipart = "^0.0.6"
 python-dotenv = "^1.0.0"
 faker = "^18.11.1"
 requests = "^2.31.0"
+openai = "^0.27.8"
 
 [tool.poetry.group.dev.dependencies]
 mypy = "^1.4.0"
 black = "^23.3.0"
 isort = "^5.12.0"
 pytest = "^7.3.2"
-openai = "^0.27.8"
 flake8 = "^6.0.0"
 types-python-jose = "^3.3.4.7"
 types-passlib = "^1.7.7.12"

vector_embedding_server/openai_like_api_models.py

@@ -15,6 +15,10 @@ class Usage(BaseModel):
     total_tokens: int
 
 
+class CompletionUsage(Usage):
+    completion_tokens: int
+
+
 class EmbeddingResponse(BaseModel):
     object: str
     data: list[EmbeddingData]
@@ -85,4 +89,4 @@ class ChatCompletionResponse(BaseModel):
     created: int
 
     choices: list[Choice]
-    usage: Usage
+    usage: CompletionUsage

vector_embedding_server/server.py

@@ -1,11 +1,12 @@
 import json
 import os
 from pathlib import Path
+from typing import Iterator, cast
 
-import requests
+import openai
 from dotenv import load_dotenv
 from fastapi import Depends, FastAPI, Request
-from fastapi.responses import HTMLResponse
+from fastapi.responses import HTMLResponse, StreamingResponse
 from fastapi.templating import Jinja2Templates
 
 from vector_embedding_server.auth import (
@@ -23,6 +24,7 @@
     EmbeddingResponse,
     Usage,
 )
+from vector_embedding_server.streaming_models import ChatCompletionStreamingResponse
 
 load_dotenv()
 
@@ -31,6 +33,9 @@
 HASHED_PASSWORD = os.environ["HASHED_PASSWORD"]
 LANGUAGE_MODEL_SERVER = os.environ["LANGUAGE_MODEL_SERVER"]
 
+openai.api_base = f"{LANGUAGE_MODEL_SERVER}/v1"
+openai.api_key = "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"
+
 
 FAKE_USERS_DB = {
     USERNAME: User(
@@ -84,16 +89,30 @@ async def create_embedding(
 
 
 @app.post("/v1/chat/completions", response_model=ChatCompletionResponse)
-def chat_completion_proxy(
+async def chat_completion_proxy(
     chat_completion_input: ChatCompletionInput,
     current_user: str = Depends(get_current_user_wrapper(FAKE_USERS_DB)),
 ) -> ChatCompletionResponse:
-    response = requests.post(
-        url=f"{LANGUAGE_MODEL_SERVER}/v1/chat/completions",
-        json=json.loads(chat_completion_input.json()),
+    response = openai.ChatCompletion.create(  # type: ignore
+        **json.loads(chat_completion_input.json())
+    )
+    if not chat_completion_input.stream:
+        return ChatCompletionResponse(**response)
+
+    def event_stream() -> Iterator[bytes]:
+        for chunk in response:
+            resp = ChatCompletionStreamingResponse(**chunk)
+            if resp.choices[0].finish_reason is None:
+                yield ("data: " + resp.json() + "\r\n\r\n").encode("utf-8")
+            else:
+                yield ("data: " + resp.json() + "\r\n\r\ndata: [DONE]\r\n\r\n").encode(
+                    "utf-8"
+                )
+
+    return cast(
+        ChatCompletionResponse,
+        StreamingResponse(event_stream(), media_type="text/event-stream"),
     )
-    response.raise_for_status()
-    return ChatCompletionResponse.parse_obj(response.json())
 
 
 @app.get("/docs", response_class=HTMLResponse)

vector_embedding_server/streaming_models.py

@@ -0,0 +1,26 @@
+from typing import Optional
+
+from pydantic import BaseModel
+
+from .openai_like_api_models import CompletionUsage, MessageRole
+
+
+class StreamingMessage(BaseModel):
+    role: Optional[MessageRole]
+    content: str
+
+
+class StreamingChoice(BaseModel):
+    index: int
+    message: StreamingMessage
+    finish_reason: Optional[str]
+    delta: StreamingMessage
+
+
+class ChatCompletionStreamingResponse(BaseModel):
+    id: str
+    object: str
+    created: int
+
+    choices: list[StreamingChoice]
+    usage: Optional[CompletionUsage]