Update request format: OpenAI-like message handling

app.py

@@ -5,22 +5,12 @@ import os
 import re
 import requests
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-from typing import List, Dict, Any, Optional, Tuple
+from pydantic import BaseModel, Field
+from typing import List, Optional, Literal, Union
 
 load_dotenv()
 
 app = FastAPI()
-ranked_models = [
-    "llama-3.3-70b-versatile",
-    "llama3-70b-8192",
-    "meta-llama/llama-4-maverick-17b-128e-instruct",
-    "meta-llama/llama-4-scout-17b-16e-instruct",
-    "mistral-saba-24b",
-    "gemma2-9b-it",
-    "llama-3.1-8b-instant",
-    "llama3-8b-8192"
-]
 
 api_keys = []
 
@@ -36,12 +26,32 @@ app.add_middleware(
     allow_origins=["*"]
 )
 
+class ChatMessage(BaseModel):
+    role: Literal["system", "user", "assistant", "tool"]
+    content: Optional[str]  # Null pour certains messages (ex: tool calls)
+    name: Optional[str] = None
+    function_call: Optional[dict] = None  # Déprécié
+    tool_call_id: Optional[str] = None
+    tool_calls: Optional[List[dict]] = None
+
 class ChatRequest(BaseModel):
-    models: Optional[List[Any]] = []
-    query: str
+    models: Optional[List[str]] = []
+    messages: List[ChatMessage]
+    temperature: Optional[float] = Field(default=1.0, ge=0.0, le=2.0)
+    top_p: Optional[float] = Field(default=1.0, ge=0.0, le=1.0)
+    n: Optional[int] = Field(default=1, ge=1)
+    stream: Optional[bool] = False
+    stop: Optional[Union[str, List[str]]] = None
+    max_tokens: Optional[int] = None
+    presence_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
+    frequency_penalty: Optional[float] = Field(default=0.0, ge=-2.0, le=2.0)
+    logit_bias: Optional[dict] = None
+    user: Optional[str] = None
+    tools: Optional[List[dict]] = None
+    tool_choice: Optional[Union[str, dict]] = None
 
-class ChatResponse(BaseModel):
-    output: str
+def clean_message(msg: ChatMessage) -> dict:
+    return {k: v for k, v in msg.model_dump().items() if v is not None}
 
 @app.get("/")
 def main_page():
@@ -50,41 +60,29 @@ def main_page():
 @app.post("/chat")
 def ask_groq_llm(req: ChatRequest):
     models = req.models
-    query = req.query
+    if len(models) == 1 and models[0] == "":
+        raise HTTPException(400, detail="Empty model field")
+    messages = [clean_message(m) for m in req.messages]
     looping = True
-    if models == []:
-        while looping:
-            for model in ranked_models:
-                for key in api_keys:
-                    resp = requests.post("https://api.groq.com/openai/v1/chat/completions", verify=False, headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"}, data=json.dumps({"model": model, "messages": [{"role": "user", "content": query}]}))
-                    if resp.status_code == 200:
-                        respJson = resp.json()
-                        print("Asked to", model, "with the key ID", str(api_keys.index(key)+1), ":", query)
-                        return {"error": False, "content": respJson["choices"]}
-                    print(resp.status_code, resp.text)
-            looping = False
-        return {"error": True, "content": "Aucun des modèles, ni des clés ne fonctionne, patientez ...."}
-    elif len(models) == 1:
+    if len(models) == 1:
         while looping:
             for key in api_keys:
-                resp = requests.post("https://api.groq.com/openai/v1/chat/completions", verify=False, headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"}, data=json.dumps({"model": models[0], "messages": [{"role": "user", "content": query}]}))
+                resp = requests.post("https://api.groq.com/openai/v1/chat/completions", verify=False, headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"}, data=json.dumps({"model": models[0], "messages": messages}))
                 if resp.status_code == 200:
                     respJson = resp.json()
-                    print("Asked to", model, "with the key ID", str(api_keys.index(key)+1), ":", query)
+                    print("Asked to", model, "with the key ID", str(api_keys.index(key)+1), ":", messages)
                     return {"error": False, "content": respJson["choices"]}
                 print(resp.status_code, resp.text)
             looping = False
         return {"error": True, "content": "Aucun des modèles, ni des clés ne fonctionne, patientez ...."}
     else:
         while looping:
-            order = {val: ind for ind, val in enumerate(ranked_models)}
-            sorted_models = sorted(models, key=lambda x: order.get(x, float('inf')))
-            for model in sorted_models:
+            for model in models:
                 for key in api_keys:
-                    resp = requests.post("https://api.groq.com/openai/v1/chat/completions", verify=False, headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"}, data=json.dumps({"model": model, "messages": [{"role": "user", "content": query}]}))
+                    resp = requests.post("https://api.groq.com/openai/v1/chat/completions", verify=False, headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"}, data=json.dumps({"model": model, "messages": messages}))
                     if resp.status_code == 200:
                         respJson = resp.json()
-                        print("Asked to", model, "with the key ID", str(api_keys.index(key)+1), ":", query)
+                        print("Asked to", model, "with the key ID", str(api_keys.index(key)+1), ":", messages)
                         return {"error": False, "content": respJson["choices"]}
                     print(resp.status_code, resp.text)
             looping = False