Update main.py

main.py

@@ -1,8 +1,9 @@
 from fastapi import FastAPI, Request
-from pydantic import BaseModel
+from fastapi.responses import StreamingResponse, JSONResponse
+from pydantic import BaseModel, Field
+from typing import List, Optional, Union
 import requests
 import json
-import os
 
 app = FastAPI()
 
@@ -10,6 +11,7 @@ API_URL = "https://api.typegpt.net/v1/chat/completions"
 API_KEY = "sk-XzS5hhsa3vpIcRLz3prQirBQXOx2hPydPzSpzdRcE1YddnNm"
 BACKEND_MODEL = "pixtral-large-latest"
 
+# Load virtual model -> system prompt mappings
 with open("model_map.json", "r") as f:
     MODEL_PROMPTS = json.load(f)
 
@@ -19,26 +21,64 @@ class Message(BaseModel):
 
 class ChatRequest(BaseModel):
     model: str
-    messages: list[Message]
+    messages: List[Message]
+    stream: Optional[bool] = False
+    temperature: Optional[float] = 1.0
+    top_p: Optional[float] = 1.0
+    n: Optional[int] = 1
+    stop: Optional[Union[str, List[str]]] = None
+    presence_penalty: Optional[float] = 0.0
+    frequency_penalty: Optional[float] = 0.0
 
-@app.post("/v1/chat/completions")
-async def openai_compatible(request: ChatRequest):
-    user_messages = [m for m in request.messages if m.role == "user"]
-    if not user_messages:
-        return {"error": "No user message provided."}
+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
+    ]
 
-    model_prompt = MODEL_PROMPTS.get(request.model, "You are a helpful AI assistant.")
-    payload = {
+    return {
         "model": BACKEND_MODEL,
-        "messages": [
-            {"role": "system", "content": model_prompt}
-        ] + [{"role": "user", "content": m.content} for m in user_messages]
+        "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
     }
 
+def stream_generator(requested_model, payload, headers):
+    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"
+        yield "data: [DONE]\n\n"
+
+@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"
     }
 
-    response = requests.post(API_URL, headers=headers, json=payload)
-    return response.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()
+        # Replace model in final result
+        if "model" in data and data["model"] == BACKEND_MODEL:
+            data["model"] = chat_request.model
+        return JSONResponse(content=data)