Upload 9 files

app.py

@@ -1,6 +1,6 @@
 import gradio as gr
 from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 from typing import List, Optional
 import asyncio
@@ -8,6 +8,7 @@ import random
 from datetime import datetime
 import threading
 import uvicorn
+import json
 
 from src.config import get_api_keys
 from src.model_tester import ModelTester
@@ -70,6 +71,12 @@ async def chat_completions(request: ChatCompletionRequest):
     prompt = request.messages[-1].content if request.messages else ""
     model_hint = request.model
     
+    if request.stream:
+        return StreamingResponse(
+            stream_chat(request.model, prompt, request.messages),
+            media_type="text/event-stream"
+        )
+    
     result = await model_tester.chat_completion(prompt, model_hint)
     
     if not result.get("success"):
@@ -104,6 +111,62 @@ async def chat_completions(request: ChatCompletionRequest):
     }
 
 
+async def stream_chat(model_hint: Optional[str], prompt: str, messages: list):
+    model_hint = model_hint or ""
+    
+    result = await model_tester.chat_completion(prompt, model_hint)
+    
+    if not result.get("success"):
+        yield f'data: {{"error": "{result.get("error", "Request failed")}"}}\n\n'
+        yield "data: [DONE]\n\n"
+        return
+    
+    response_data = result.get("response", {})
+    content = ""
+    if "choices" in response_data and response_data["choices"]:
+        content = response_data["choices"][0].get("message", {}).get("content", "")
+    
+    model_id = result.get("model", model_hint or "unknown")
+    completion_id = f"chatcmpl-{random.randint(100000, 999999)}"
+    created = int(datetime.now().timestamp())
+    
+    # 流式输出每个字
+    for i, char in enumerate(content):
+        chunk = {
+            "id": completion_id,
+            "object": "chat.completion.chunk",
+            "created": created,
+            "model": model_id,
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {
+                        "content": char
+                    },
+                    "finish_reason": None
+                }
+            ]
+        }
+        yield f"data: {json.dumps(chunk)}\n\n"
+    
+    # 发送完成信号
+    final_chunk = {
+        "id": completion_id,
+        "object": "chat.completion.chunk",
+        "created": created,
+        "model": model_id,
+        "choices": [
+            {
+                "index": 0,
+                "delta": {},
+                "finish_reason": "stop"
+            }
+        ]
+    }
+    yield f"data: {json.dumps(final_chunk)}\n\n"
+    yield "data: [DONE]\n\n"
+
+
 @fastapi_app.get("/health")
 async def health():
     return {"status": "ok"}