Update main.py

main.py

@@ -1,55 +1,44 @@
+from fastapi import FastAPI, Request
+from pydantic import BaseModel
 import requests
+import json
+import os
+
+app = FastAPI()
 
 API_URL = "https://api.typegpt.net/v1/chat/completions"
 API_KEY = "sk-XzS5hhsa3vpIcRLz3prQirBQXOx2hPydPzSpzdRcE1YddnNm"
 BACKEND_MODEL = "pixtral-large-latest"
 
-# Define virtual models with system prompts
-SYSTEM_PROMPTS = {
-    "gpt-3.5": "You are ChatGPT, a helpful assistant.",
-    "claude": "You are Claude, known for being careful, ethical, and thoughtful.",
-    "mistral": "You are Mistral, known for giving concise and smart answers.",
-    "bard": "You are Bard, a friendly and knowledgeable Google AI assistant."
-}
+with open("model_map.json", "r") as f:
+    MODEL_PROMPTS = json.load(f)
 
-def generate_response(model_label: str, user_prompt: str):
-    if model_label not in SYSTEM_PROMPTS:
-        raise ValueError(f"Unknown model label '{model_label}'. Valid options: {', '.join(SYSTEM_PROMPTS.keys())}")
+class Message(BaseModel):
+    role: str
+    content: str
 
-    system_prompt = SYSTEM_PROMPTS[model_label]
+class ChatRequest(BaseModel):
+    model: str
+    messages: list[Message]
 
-    headers = {
-        "Authorization": f"Bearer {API_KEY}",
-        "Content-Type": "application/json"
-    }
+@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."}
 
+    model_prompt = MODEL_PROMPTS.get(request.model, "You are a helpful AI assistant.")
     payload = {
         "model": BACKEND_MODEL,
         "messages": [
-            {"role": "system", "content": system_prompt},
-            {"role": "user", "content": user_prompt}
-        ]
+            {"role": "system", "content": model_prompt}
+        ] + [{"role": "user", "content": m.content} for m in user_messages]
     }
 
-    response = requests.post(API_URL, headers=headers, json=payload)
+    headers = {
+        "Authorization": f"Bearer {API_KEY}",
+        "Content-Type": "application/json"
+    }
 
-    if response.status_code == 200:
-        result = response.json()
-        return result["choices"][0]["message"]["content"]
-    else:
-        raise Exception(f"API error {response.status_code}: {response.text}")
-
-if __name__ == "__main__":
-    print("🧠 Choose a virtual model:")
-    for name in SYSTEM_PROMPTS:
-        print(f" - {name}")
-    model = input("Model: ").strip()
-
-    print("\n👤 Enter your message:")
-    user_prompt = input("> ").strip()
-
-    try:
-        print("\n🤖 AI Response:")
-        print(generate_response(model, user_prompt))
-    except Exception as e:
-        print(f"❌ Error: {e}")
+    response = requests.post(API_URL, headers=headers, json=payload)
+    return response.json()