Imago/app/routers/admin_ai.py

"""
Router — Admin AI : configuration, test et préférences AI par client.
"""
import asyncio
import base64
import logging
import time
from typing import Optional

import httpx
from fastapi import APIRouter, Depends, HTTPException, status
from pydantic import BaseModel, Field
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.config import settings
from app.database import get_db
from app.dependencies.auth import require_scope
from app.models.client import APIClient

logger = logging.getLogger(__name__)

router = APIRouter(prefix="/admin/api/ai", tags=["Admin AI"])


# ─────────────────────────────────────────────────────────────
# Schemas
# ─────────────────────────────────────────────────────────────

class ProviderInfo(BaseModel):
    key: str
    name: str
    configured: bool
    models: list[str] = []


class AIStatusResponse(BaseModel):
    ai_enabled: bool
    current_provider: str
    current_model: str
    providers: list[ProviderInfo]


class TestModelRequest(BaseModel):
    provider: str = Field(..., description="openrouter ou gemini")
    model: str = Field(..., description="Nom du modèle à tester")
    prompt: Optional[str] = Field(None, description="Prompt de test (optionnel)")
    image_base64: Optional[str] = Field(None, description="Image base64 pour test vision (optionnel)")


class TestModelResponse(BaseModel):
    success: bool
    provider: str
    model: str
    response: Optional[str] = None
    tokens: Optional[dict] = None
    latency_ms: float = 0
    error: Optional[str] = None


# ─────────────────────────────────────────────────────────────
# Constantes — Modèles disponibles par fournisseur
# ─────────────────────────────────────────────────────────────

OPENROUTER_VISION_MODELS = [
    "google/gemini-2.5-flash-lite",
    "google/gemini-2.5-flash",
    "google/gemini-2.5-pro-preview",
    "qwen/qwen2.5-vl-72b-instruct",
    "anthropic/claude-3.5-sonnet",
    "anthropic/claude-3-opus",
    "openai/gpt-4o",
    "openai/gpt-4o-mini",
    "meta-llama/llama-4-maverick",
]

GEMINI_MODELS = [
    "gemini-2.5-flash-lite",
    "gemini-2.5-flash",
    "gemini-2.5-pro-preview",
    "gemini-2.0-flash",
]


def _build_providers() -> list[ProviderInfo]:
    """Construit la liste des fournisseurs disponibles selon la config .env."""
    providers = []

    # OpenRouter
    or_configured = bool(settings.OPENROUTER_API_KEY)
    providers.append(ProviderInfo(
        key="openrouter",
        name="OpenRouter",
        configured=or_configured,
        models=OPENROUTER_VISION_MODELS if or_configured else [],
    ))

    # Gemini
    gemini_configured = bool(settings.GEMINI_API_KEY)
    providers.append(ProviderInfo(
        key="gemini",
        name="Google Gemini",
        configured=gemini_configured,
        models=GEMINI_MODELS if gemini_configured else [],
    ))

    return providers


# ─────────────────────────────────────────────────────────────
# GET /admin/api/ai/status
# ─────────────────────────────────────────────────────────────

@router.get("/status", response_model=AIStatusResponse)
async def get_ai_status():
    """Retourne la configuration AI actuelle et les fournisseurs/modèles disponibles."""
    active_model = (
        settings.OPENROUTER_MODEL if settings.AI_PROVIDER == "openrouter"
        else settings.GEMINI_MODEL
    )
    return AIStatusResponse(
        ai_enabled=settings.AI_ENABLED,
        current_provider=settings.AI_PROVIDER,
        current_model=active_model,
        providers=_build_providers(),
    )


# ─────────────────────────────────────────────────────────────
# POST /admin/api/ai/test
# ─────────────────────────────────────────────────────────────

@router.post("/test", response_model=TestModelResponse)
async def test_model(
    body: TestModelRequest,
    _=Depends(require_scope("admin")),
):
    """Teste un modèle AI avec un prompt simple (et optionnellement une image)."""
    t0 = time.monotonic()

    if body.provider == "openrouter":
        result = await _test_openrouter(body.model, body.prompt, body.image_base64)
    elif body.provider == "gemini":
        result = await _test_gemini(body.model, body.prompt, body.image_base64)
    else:
        raise HTTPException(
            status_code=status.HTTP_400_BAD_REQUEST,
            detail=f"Fournisseur inconnu : {body.provider}",
        )

    result["latency_ms"] = round((time.monotonic() - t0) * 1000, 1)
    result["provider"] = body.provider
    result["model"] = body.model
    return result


async def _test_openrouter(model: str, prompt: Optional[str], image_b64: Optional[str]) -> dict:
    """Appel test vers OpenRouter."""
    if not settings.OPENROUTER_API_KEY:
        return {"success": False, "error": "Clé API OpenRouter non configurée"}

    test_prompt = prompt or "Réponds uniquement avec le mot 'OK' pour confirmer que tu fonctionnes."
    headers = {
        "Authorization": f"Bearer {settings.OPENROUTER_API_KEY}",
        "Content-Type": "application/json",
        "HTTP-Referer": "imago-test",
        "X-Title": "Imago AI Test",
    }

    messages = []
    content_payload = [{"type": "text", "text": test_prompt}]
    if image_b64:
        content_payload.append({
            "type": "image_url",
            "image_url": {"url": f"data:image/jpeg;base64,{image_b64}"},
        })
    messages.append({"role": "user", "content": content_payload})

    payload = {
        "model": model,
        "messages": messages,
        "max_tokens": 64,
    }

    try:
        async with httpx.AsyncClient() as client:
            resp = await client.post(
                "https://openrouter.ai/api/v1/chat/completions",
                json=payload,
                headers=headers,
                timeout=60,
            )
            resp.raise_for_status()
            data = resp.json()

        text = data.get("choices", [{}])[0].get("message", {}).get("content", "")
        usage = data.get("usage", {})
        return {
            "success": True,
            "response": text[:500],
            "tokens": {
                "prompt": usage.get("prompt_tokens"),
                "completion": usage.get("completion_tokens"),
            },
        }
    except httpx.HTTPStatusError as e:
        return {"success": False, "error": f"HTTP {e.response.status_code}: {e.response.text[:300]}"}
    except Exception as e:
        return {"success": False, "error": str(e)[:500]}


async def _test_gemini(model: str, prompt: Optional[str], image_b64: Optional[str]) -> dict:
    """Appel test vers Google Gemini."""
    if not settings.GEMINI_API_KEY:
        return {"success": False, "error": "Clé API Gemini non configurée"}

    try:
        from google import genai
        from google.genai import types
    except ImportError:
        return {"success": False, "error": "SDK Google GenAI non installé"}

    test_prompt = prompt or "Réponds uniquement avec le mot 'OK' pour confirmer que tu fonctionnes."

    client = genai.Client(api_key=settings.GEMINI_API_KEY)
    contents = []
    if image_b64:
        contents.append(types.Part.from_bytes(
            data=base64.b64decode(image_b64),
            mime_type="image/jpeg",
        ))
    contents.append(test_prompt)

    try:
        response = await asyncio.wait_for(
            asyncio.to_thread(
                client.models.generate_content,
                model=model,
                contents=contents,
                config=types.GenerateContentConfig(max_output_tokens=64),
            ),
            timeout=60,
        )
        text = getattr(response, "text", "")
        usage = getattr(response, "usage_metadata", None)
        tokens = None
        if usage:
            tokens = {
                "prompt": getattr(usage, "prompt_token_count", None),
                "completion": getattr(usage, "candidates_token_count", None),
            }
        return {"success": True, "response": text[:500], "tokens": tokens}
    except Exception as e:
        return {"success": False, "error": str(e)[:500]}