Update main.py

main.py

@@ -4,17 +4,24 @@ from pydantic import BaseModel
 from typing import List, Optional, Union
 import requests
 import json
+import logging
 
 app = FastAPI()
 
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger("proxy")
+
+# TypeGPT API settings
 API_URL = "https://api.typegpt.net/v1/chat/completions"
 API_KEY = "sk-XzS5hhsa3vpIcRLz3prQirBQXOx2hPydPzSpzdRcE1YddnNm"
 BACKEND_MODEL = "pixtral-large-latest"
 
-# Load virtual model -> system prompt mappings
+# Load model-system-prompt mappings
 with open("model_map.json", "r") as f:
     MODEL_PROMPTS = json.load(f)
 
+# Request schema
 class Message(BaseModel):
     role: str
     content: str
@@ -30,62 +37,68 @@ class ChatRequest(BaseModel):
     presence_penalty: Optional[float] = 0.0
     frequency_penalty: Optional[float] = 0.0
 
+# Build payload to send to actual backend API
 def build_payload(chat: ChatRequest):
     system_prompt = MODEL_PROMPTS.get(chat.model, "You are a helpful assistant.")
-
-    messages = [{"role": "system", "content": system_prompt}] + [
-        {"role": msg.role, "content": msg.content} for msg in chat.messages
-    ]
-
     return {
         "model": BACKEND_MODEL,
-        "messages": messages,
         "stream": chat.stream,
         "temperature": chat.temperature,
         "top_p": chat.top_p,
         "n": chat.n,
         "stop": chat.stop,
         "presence_penalty": chat.presence_penalty,
-        "frequency_penalty": chat.frequency_penalty
+        "frequency_penalty": chat.frequency_penalty,
+        "messages": [{"role": "system", "content": system_prompt}] + [
+            {"role": msg.role, "content": msg.content} for msg in chat.messages
+        ]
     }
 
+# Properly forward streaming data and replace model
 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):
-            if line and line.startswith("data:"):
-                # Remove "data: " prefix
+            if not line:
+                continue
+            if line.startswith("data:"):
                 content = line[6:].strip()
+                if content == "[DONE]":
+                    yield "data: [DONE]\n\n"
+                    continue
                 try:
-                    # Try to parse and replace model field
                     json_obj = json.loads(content)
-                    if "model" in json_obj and json_obj["model"] == BACKEND_MODEL:
+                    if json_obj.get("model") == BACKEND_MODEL:
                         json_obj["model"] = requested_model
-                    fixed_line = f"data: {json.dumps(json_obj)}\n\n"
+                    yield 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"
+                    logger.warning("Invalid JSON in stream chunk: %s", content)
+            else:
+                logger.debug("Non-data stream line skipped: %s", line)
 
+# Main endpoint
 @app.post("/v1/chat/completions")
 async def proxy_chat(request: Request):
-    body = await request.json()
-    chat_request = ChatRequest(**body)
-    payload = build_payload(chat_request)
-    headers = {
-        "Authorization": f"Bearer {API_KEY}",
-        "Content-Type": "application/json"
-    }
+    try:
+        body = await request.json()
+        chat_request = ChatRequest(**body)
+        payload = build_payload(chat_request)
+        headers = {
+            "Authorization": f"Bearer {API_KEY}",
+            "Content-Type": "application/json"
+        }
+
+        if chat_request.stream:
+            return StreamingResponse(
+                stream_generator(chat_request.model, payload, headers),
+                media_type="text/event-stream"
+            )
+        else:
+            response = requests.post(API_URL, headers=headers, json=payload)
+            data = response.json()
+            if "model" in data and data["model"] == BACKEND_MODEL:
+                data["model"] = chat_request.model
+            return JSONResponse(content=data)
 
-    if chat_request.stream:
-        return StreamingResponse(
-            stream_generator(chat_request.model, payload, headers),
-            media_type="text/event-stream"
-        )
-    else:
-        response = requests.post(API_URL, headers=headers, json=payload)
-        data = response.json()
-        if "model" in data and data["model"] == BACKEND_MODEL:
-            data["model"] = chat_request.model
-        return JSONResponse(content=data)
+    except Exception as e:
+        logger.error("Error in proxy_chat: %s", str(e))
+        return JSONResponse(content={"error": "Internal server error."}, status_code=500)