Changed to LiteLLM

app.py

@@ -1,12 +1,16 @@
 import json
+import traceback
 from fastapi import FastAPI, HTTPException
 from dotenv import load_dotenv
 import os
 import re
-import requests
+from huggingface_hub import ChatCompletionInputMessage, ChatCompletionInputTool
+import litellm
+litellm.ssl_verify = False
+from litellm.router import Router
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
-from typing import List, Optional, Literal, Union
+from typing import List, Optional, Literal, Type, Union
 
 load_dotenv()
 
@@ -18,6 +22,65 @@ for k,v in os.environ.items():
     if re.match(r'^GROQ_\d+$', k):
         api_keys.append(v)
 
+models_data = {
+    "allam-2-7b": {"rpm": 30, "rpd": 7000, "tpm": 6000},
+    "compound-beta": {"rpm": 15, "rpd": 200, "tpm": 70000},
+    "compound-beta-mini": {"rpm": 15, "rpd": 200, "tpm": 70000},
+    "deepseek-r1-distill-llama-70b": {"rpm": 30, "rpd": 1000, "tpm": 6000},
+    "gemma2-9b-it": {"rpm": 30, "rpd": 14400, "tpm": 15000, "tpd": 500000},
+    "llama-3.1-8b-instant": {"rpm": 30, "rpd": 14400, "tpm": 6000, "tpd": 500000},
+    "llama-3.3-70b-versatile": {"rpm": 30, "rpd": 1000, "tpm": 12000, "tpd": 100000},
+    "llama3-70b-8192": {"rpm": 30, "rpd": 14400, "tpm": 6000, "tpd": 500000},
+    "llama3-8b-8192": {"rpm": 30, "rpd": 14400, "tpm": 6000, "tpd": 500000},
+}
+
+model_list = [
+    {
+        "model_name": f"{model_name}_{key_idx}" if key_idx != 0 else f"{model_name}",  # Nom unique par clé
+        "litellm_params": {
+            "model": f"groq/{model_name}",
+            "api_key": api_key
+        },
+        "timeout": 120,
+        "max_retries": 5
+    }
+    for model_name, config in models_data.items()
+    for key_idx, api_key in enumerate(api_keys)
+]
+
+def generate_fallbacks_per_key():
+    fallbacks = []  # Liste de dictionnaires au lieu d'un dictionnaire
+    excluded_models = {"compound-beta", "compound-beta-mini"}
+    
+    for model_name in models_data.keys():
+        if model_name in excluded_models:
+            continue
+            
+        # Pour chaque version d'un modèle, les fallbacks sont les autres versions du même modèle
+        for key_idx in range(len(api_keys)):
+            current_model = f"{model_name}_{key_idx}" if key_idx != 0 else f"{model_name}"
+            fallback_versions = [
+                f"{model_name}_{other_key_idx}" if other_key_idx != 0 else f"{model_name}"
+                for other_key_idx in range(len(api_keys)) 
+                if other_key_idx != key_idx
+            ]
+            
+            # Format attendu par LiteLLM
+            fallbacks.append({
+                current_model: fallback_versions
+            })
+    
+    return fallbacks
+
+fallbacks = generate_fallbacks_per_key()
+
+router = Router(
+    model_list=model_list,
+    fallbacks=fallbacks,
+    num_retries=5,
+    retry_after=10
+)
+
 app.add_middleware(
     CORSMiddleware,
     allow_credentials=True,
@@ -26,64 +89,56 @@ 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[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
+    models: List[str]
+    messages: List[ChatCompletionInputMessage]
+    tools: Optional[List[ChatCompletionInputTool]] = None
+    temperature: Optional[float] = 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
+    n: Optional[int] = None
+    stream: Optional[bool] = None
+    stop: Optional[List[str]] = None
 
-def clean_message(msg: ChatMessage) -> dict:
-    return {k: v for k, v in msg.model_dump().items() if v is not None}
+def clean_message(msg) -> dict:
+    """Convertit un message en dictionnaire, gérant différents types d'objets"""
+    if hasattr(msg, 'model_dump'):
+        # Pour les objets Pydantic
+        return {k: v for k, v in msg.model_dump().items() if v is not None}
+    elif hasattr(msg, '__dict__'):
+        # Pour les objets avec attributs
+        return {k: v for k, v in msg.__dict__.items() if v is not None}
+    elif isinstance(msg, dict):
+        # Si c'est déjà un dictionnaire
+        return {k: v for k, v in msg.items() if v is not None}
+    else:
+        # Conversion générique
+        return dict(msg)
 
 @app.get("/")
 def main_page():
     return {"status": "ok"}
 
 @app.post("/chat")
-def ask_groq_llm(req: ChatRequest):
+def chat_with_groq(req: ChatRequest):
     models = req.models
     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 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": messages}))
-                if resp.status_code == 200:
-                    respJson = resp.json()
-                    print("Asked to", models[0], "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 ...."}
+        resp = router.completion(model=models[0], messages=messages, **req.model_dump(exclude={"models", "messages"}, exclude_defaults=True, exclude_none=True))
+        try:
+            print("Asked to", models[0], ":", messages)
+            return {"error": False, "content": resp.choices[0].message.content}
+        except Exception as e:
+            traceback.print_exception(e)
+            return {"error": True, "content": "Aucune clé ne fonctionne avec le modèle sélectionné, patientez ...."}
     else:
-        while looping:
-            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": messages}))
-                    if resp.status_code == 200:
-                        respJson = resp.json()
-                        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 ...."}
+        for model in models:
+            resp = router.completion(model=model, messages=messages, **req.model_dump(exclude={"models", "messages"}, exclude_defaults=True, exclude_none=True))
+            try:
+                print("Asked to", models[0], ":", messages)
+                return {"error": False, "content": resp.choices[0].message.content}
+            except Exception as e:
+                traceback.print_exception(e)
+                continue
+        return {"error": True, "content": "Aucune clé ne fonctionne avec le modèle sélectionné, patientez ...."}