Update main.py

main.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI, Request
 from fastapi.responses import StreamingResponse, JSONResponse
-from pydantic import BaseModel, Field
+from pydantic import BaseModel
 from typing import List, Optional, Union
 import requests
 import json
@@ -49,15 +49,23 @@ def build_payload(chat: ChatRequest):
         "frequency_penalty": chat.frequency_penalty
     }
 
-def stream_generator(requested_model, payload, headers):
+def stream_generator(requested_model: str, payload: dict, headers: dict):
     with requests.post(API_URL, headers=headers, json=payload, stream=True) as r:
-        for line in r.iter_lines():
-            if line:
-                decoded = line.decode('utf-8')
-                # Rewrite the model name in streaming output
-                if BACKEND_MODEL in decoded:
-                    decoded = decoded.replace(BACKEND_MODEL, requested_model)
-                yield f"data: {decoded}\n\n"
+        for line in r.iter_lines(decode_unicode=True):
+            if line and line.startswith("data:"):
+                # Remove "data: " prefix
+                content = line[6:].strip()
+                try:
+                    # Try to parse and replace model field
+                    json_obj = json.loads(content)
+                    if "model" in json_obj and json_obj["model"] == BACKEND_MODEL:
+                        json_obj["model"] = requested_model
+                    fixed_line = f"data: {json.dumps(json_obj)}\n\n"
+                except json.JSONDecodeError:
+                    fixed_line = f"data: {content}\n\n"
+                yield fixed_line
+            elif line:
+                yield f"data: {line}\n\n"
         yield "data: [DONE]\n\n"
 
 @app.post("/v1/chat/completions")
@@ -78,7 +86,6 @@ async def proxy_chat(request: Request):
     else:
         response = requests.post(API_URL, headers=headers, json=payload)
         data = response.json()
-        # Replace model in final result
         if "model" in data and data["model"] == BACKEND_MODEL:
             data["model"] = chat_request.model
         return JSONResponse(content=data)