Update app.py

app.py

@@ -1,21 +1,27 @@
-from litellm import completion
+import os
+import litellm
+from fastapi import FastAPI
+from pydantic import BaseModel
 
 # Configuration du modèle
 model_name = "gemini/gemini-1.5-pro"
 api_key = os.environ.get("GEMINI_API_KEY")
 
-# Exemple de fonction de base
-def get_completion(prompt):
+# Configuration de LiteLLM
+litellm.api_key = api_key
+
+app = FastAPI()
+
+class PromptRequest(BaseModel):
+    prompt: str
+
+@app.post("/")
+async def predict(request: PromptRequest):
     try:
-        response = completion(
+        response = litellm.completion(
             model=model_name, 
-            messages=[{"role": "user", "content": prompt}],
-            api_key=api_key
+            messages=[{"role": "user", "content": request.prompt}]
         )
-        return response.choices[0].message.content
+        return {"response": response.choices[0].message.content}
     except Exception as e:
-        return f"Erreur: {str(e)}"
-
-# Pour les Spaces Hugging Face
-def predict(prompt):
-    return get_completion(prompt)
+        return {"error": str(e)}