Update main.py

main.py

@@ -8,16 +8,16 @@ import logging
 
 app = FastAPI()
 
-# Setup logging
+# Logging setup
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("proxy")
 
-# TypeGPT API settings
+# Configuration
 API_URL = "https://api.typegpt.net/v1/chat/completions"
 API_KEY = "sk-XzS5hhsa3vpIcRLz3prQirBQXOx2hPydPzSpzdRcE1YddnNm"
 BACKEND_MODEL = "pixtral-large-latest"
 
-# Load model -> system prompt mappings
+# Load system prompt mappings
 with open("model_map.json", "r", encoding="utf-8") as f:
     MODEL_PROMPTS = json.load(f)
 
@@ -37,10 +37,9 @@ class ChatRequest(BaseModel):
     presence_penalty: Optional[float] = 0.0
     frequency_penalty: Optional[float] = 0.0
 
-# Construct payload with enforced system prompt
+# Build request to backend with injected system prompt
 def build_payload(chat: ChatRequest):
     system_prompt = MODEL_PROMPTS.get(chat.model, "You are a helpful assistant.")
-    # Strip user system messages
     filtered_messages = [msg for msg in chat.messages if msg.role != "system"]
     payload_messages = [{"role": "system", "content": system_prompt}] + [
         {"role": msg.role, "content": msg.content} for msg in filtered_messages
@@ -57,7 +56,7 @@ def build_payload(chat: ChatRequest):
         "frequency_penalty": chat.frequency_penalty
     }
 
-# Properly streamed UTF-8 chunks with model rewrite
+# Streaming chunk handler with model replacement and UTF-8 fix
 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(decode_unicode=True):
@@ -72,13 +71,13 @@ def stream_generator(requested_model: str, payload: dict, headers: dict):
                     json_obj = json.loads(content)
                     if json_obj.get("model") == BACKEND_MODEL:
                         json_obj["model"] = requested_model
-                    yield f"data: {json.dumps(json_obj, ensure_ascii=False)}\n\n"
+                    yield "data: " + json.dumps(json_obj, ensure_ascii=False) + "\n\n"
                 except json.JSONDecodeError:
                     logger.warning("Invalid JSON in stream chunk: %s", content)
             else:
                 logger.debug("Non-data stream line skipped: %s", line)
 
-# Main endpoint
+# Main API endpoint
 @app.post("/v1/chat/completions")
 async def proxy_chat(request: Request):
     try:
@@ -102,8 +101,15 @@ async def proxy_chat(request: Request):
             data = response.json()
             if "model" in data and data["model"] == BACKEND_MODEL:
                 data["model"] = chat_request.model
-            return JSONResponse(content=data, media_type="application/json; charset=utf-8")
+            return JSONResponse(
+                content=data,
+                media_type="application/json; charset=utf-8",
+                headers={"Content-Type": "application/json; charset=utf-8"}
+            )
 
     except Exception as e:
         logger.error("Error in /v1/chat/completions: %s", str(e))
-        return JSONResponse(content={"error": "Internal server error."}, status_code=500)
+        return JSONResponse(
+            content={"error": "Internal server error."},
+            status_code=500
+        )