import os
import litellm
from fastapi import FastAPI, HTTPException, Header, Depends
from pydantic import BaseModel

# Configuration du modèle
model_name = "gemini/gemini-1.5-pro"
api_key = os.environ.get("GEMINI_API_KEY")

# Configuration de LiteLLM
litellm.api_key = api_key

app = FastAPI()

class PromptRequest(BaseModel):
    prompt: str
    model: str = "gemini/gemini-1.5-pro"  # Valeur par défaut

# Prompt de base
base_prompt = """
Tu es un développeur logiciel expérimenté. Tu dois toujours répondre en français d'une manière réfléchie et concise. 
"""

# Clé API pour l'authentification
USER_API_KEY = "YOUR_API_KEY"  # Remplacez par une clé secrète

async def verify_api_key(x_api_key: str = Header(...)):
    if x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Clé API invalide")

@app.post("/", dependencies=[Depends(verify_api_key)])
async def predict(request: PromptRequest):
    try:
        # Ajout du prompt de base à la requête de l'utilisateur
        full_prompt = base_prompt + request.prompt 

        response = litellm.completion(
            model=request.model, 
            messages=[{"role": "user", "content": full_prompt}]
        )
        return {"response": response.choices[0].message.content, "model": request.model}
    except Exception as e:
        return {"error": str(e)}