Spaces:

Romanes
/

Surco

Running

Surco / app.py
fix: migrar de google-generativeai a google-genai SDK (soporte gemini-2.5-flash)
0a5be17 verified 5 months ago
29.4 kB
import asyncio
import base64
import json
import logging
import os
import uuid
from datetime import datetime
from typing import Any, Optional

import firebase_admin
import pytz
from anthropic import (
    APIConnectionError,
    APIStatusError,
    APITimeoutError,
    AsyncAnthropic,
    AuthenticationError,
    BadRequestError,
    RateLimitError,
)
from fastapi import Depends, FastAPI, Header, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from firebase_admin import auth, credentials, firestore
from google import genai as google_genai
from google.genai import errors as genai_errors
from google.genai import types as genai_types
from pydantic import BaseModel, Field, field_validator

# --- Logging con request_id ---
class _RequestIdFilter(logging.Filter):
    def filter(self, record: logging.LogRecord) -> bool:
        if not hasattr(record, "request_id"):
            record.request_id = "-"
        return True


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s [%(request_id)s]: %(message)s",
)
logger = logging.getLogger("surco_proxy")
logger.addFilter(_RequestIdFilter())

# --- Constantes ---
BOGOTA_TZ = pytz.timezone("America/Bogota")
HTTP_TIMEOUT_SECONDS = 60.0
MAX_HISTORY = 10
MAX_TOKENS = 1024

DEFAULT_CLAUDE_MODEL = "claude-sonnet-4-6"
ALLOWED_CLAUDE_MODELS = {
    "claude-sonnet-4-6",
    "claude-opus-4-7",
    "claude-haiku-4-5-20251001",
}

DEFAULT_GEMINI_MODEL = "gemini-2.5-flash"
ALLOWED_GEMINI_MODELS = {
    "gemini-2.5-flash",
    "gemini-2.5-pro",
}

AI_DAILY_LIMITS = {
    "free": 5,
    "pro": 50,
    "proTecnico": 200,
}
DEFAULT_AI_PLAN = "free"

# --- Estado global tolerante a fallos en boot ---
# IMPORTANTE: no lanzar excepciones al import-time o el Space queda en "Starting"
# indefinidamente. Cualquier fallo se reporta en GET / y en los logs.
db: Optional[firestore.Client] = None
firebase_error: Optional[str] = None

if not firebase_admin._apps:
    service_account_json = os.environ.get("FIREBASE_SERVICE_ACCOUNT")
    if not service_account_json:
        firebase_error = "Falta secret FIREBASE_SERVICE_ACCOUNT"
        logger.error(firebase_error)
    else:
        try:
            cred = credentials.Certificate(json.loads(service_account_json))
            firebase_admin.initialize_app(cred)
            db = firestore.client()
            logger.info("Firebase inicializado.")
        except (ValueError, json.JSONDecodeError) as exc:
            firebase_error = f"FIREBASE_SERVICE_ACCOUNT inválido: {exc}"
            logger.error(firebase_error)
        except Exception as exc:
            firebase_error = f"Error inicializando Firebase: {exc}"
            logger.exception("Firebase init failed")
else:
    db = firestore.client()

# --- Clientes externos como singletons ---
CLAUDE_API_KEY = os.environ.get("CLAUDE_API_KEY")
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")

if not CLAUDE_API_KEY:
    logger.warning("CLAUDE_API_KEY no configurada — /claude responderá 503.")
if not GEMINI_API_KEY:
    logger.warning("GEMINI_API_KEY no configurada — /gemini responderá 503.")

claude_client: Optional[AsyncAnthropic] = None
if CLAUDE_API_KEY:
    try:
        claude_client = AsyncAnthropic(
            api_key=CLAUDE_API_KEY, timeout=HTTP_TIMEOUT_SECONDS
        )
    except Exception as exc:
        logger.exception("No se pudo crear cliente Anthropic: %s", exc)

gemini_client: Optional[google_genai.Client] = None
if GEMINI_API_KEY:
    try:
        gemini_client = google_genai.Client(api_key=GEMINI_API_KEY)
        logger.info("Cliente Gemini (google-genai) inicializado, model=%s", DEFAULT_GEMINI_MODEL)
    except Exception as exc:
        logger.exception("No se pudo crear cliente Gemini: %s", exc)

# --- App ---
app = FastAPI(title="Surco AI Proxy", version="2.1.1")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["GET", "POST"],
    allow_headers=["Authorization", "Content-Type", "X-Request-Id"],
)


@app.middleware("http")
async def attach_request_id(request: Request, call_next):
    request_id = request.headers.get("x-request-id") or uuid.uuid4().hex[:12]
    request.state.request_id = request_id
    response = await call_next(request)
    response.headers["x-request-id"] = request_id
    return response


# --- Schemas ---
class HistoryMessage(BaseModel):
    role: str
    text: Optional[str] = None


class AIRequest(BaseModel):
    history: list[HistoryMessage] = Field(default_factory=list)
    userText: Optional[str] = None
    imageBase64: Optional[str] = None
    systemPrompt: Optional[str] = ""
    model: Optional[str] = None

    @field_validator("imageBase64")
    @classmethod
    def _validate_image(cls, v: Optional[str]) -> Optional[str]:
        if v is None:
            return v
        try:
            base64.b64decode(v[:200], validate=True)
        except Exception as exc:
            raise ValueError(f"imageBase64 inválido: {exc}") from exc
        return v


# --- Auth ---
async def verify_token(authorization: str = Header(None)) -> dict:
    if not authorization or not authorization.startswith("Bearer "):
        raise HTTPException(status_code=401, detail="Falta token")
    token = authorization.removeprefix("Bearer ").strip()
    try:
        return auth.verify_id_token(token)
    except auth.ExpiredIdTokenError as exc:
        raise HTTPException(status_code=401, detail="Token expirado") from exc
    except auth.RevokedIdTokenError as exc:
        raise HTTPException(status_code=401, detail="Token revocado") from exc
    except (auth.InvalidIdTokenError, ValueError) as exc:
        raise HTTPException(status_code=401, detail="Token inválido") from exc


# --- Cuotas ---
def _get_user_plan(uid: str) -> str:
    if db is None:
        return DEFAULT_AI_PLAN
    try:
        doc = (
            db.collection("users")
            .document(uid)
            .collection("subscription")
            .document("current")
            .get()
        )
        if not doc.exists:
            return DEFAULT_AI_PLAN
        plan = (doc.to_dict() or {}).get("plan", DEFAULT_AI_PLAN)
        if plan == "premium":
            return "pro"
        return plan
    except Exception as exc:
        logger.warning("get_user_plan fallback uid=%s: %s", uid, exc)
        return DEFAULT_AI_PLAN


def _check_and_increment_quota(uid: str) -> dict:
    plan = _get_user_plan(uid)
    limit = AI_DAILY_LIMITS.get(plan, 5)
    today = datetime.now(BOGOTA_TZ).strftime("%Y-%m-%d")
    if db is None:
        # Firebase no disponible — concedemos sin contar para no bloquear UX,
        # pero el endpoint igual va a fallar más abajo si Firebase es requerido.
        return {"allowed": True, "remaining": limit, "plan": plan}
    ref = (
        db.collection("users")
        .document(uid)
        .collection("ai_usage")
        .document(today)
    )

    @firestore.transactional
    def txn(transaction):
        snap = ref.get(transaction=transaction)
        current = (snap.to_dict() or {}).get("count", 0) if snap.exists else 0
        if current >= limit:
            return {
                "allowed": False,
                "limit": limit,
                "plan": plan,
                "used": current,
            }
        new_count = current + 1
        transaction.set(
            ref,
            {
                "count": new_count,
                "limit": limit,
                "plan": plan,
                "date": today,
            },
            merge=True,
        )
        return {"allowed": True, "remaining": limit - new_count, "plan": plan}

    return txn(db.transaction())


def _decrement_quota(uid: str) -> None:
    if db is None:
        return
    today = datetime.now(BOGOTA_TZ).strftime("%Y-%m-%d")
    ref = (
        db.collection("users")
        .document(uid)
        .collection("ai_usage")
        .document(today)
    )
    try:
        ref.update({"count": firestore.Increment(-1)})
    except Exception as exc:
        logger.warning("decrement_quota error uid=%s: %s", uid, exc)


_MODEL_ALIASES: dict[str, str] = {
    "gemini-1.5-flash": "gemini-2.5-flash",
    "gemini-1.5-pro": "gemini-2.5-pro",
    "gemini-1.5-flash-8b": "gemini-2.5-flash",
}


def _resolve_model(requested: Optional[str], allowed: set, default: str) -> str:
    if requested:
        aliased = _MODEL_ALIASES.get(requested)
        if aliased:
            logger.info("Modelo %r obsoleto → alias %r", requested, aliased)
            requested = aliased
    if requested and requested in allowed:
        return requested
    if requested:
        logger.info("Modelo %r no permitido; usando default %r", requested, default)
    return default


def _quota_response(quota: dict) -> JSONResponse:
    return JSONResponse(
        status_code=429,
        content={
            "error": "rate_limit_exceeded",
            "message": "Límite diario alcanzado",
            "limit": quota.get("limit"),
            "plan": quota.get("plan"),
        },
    )


# --- Endpoints ---
@app.post("/claude")
async def claude_proxy(
    request: Request,
    payload: AIRequest,
    user: dict = Depends(verify_token),
):
    rid = request.state.request_id
    extra = {"request_id": rid}
    uid = user["uid"]

    if claude_client is None:
        raise HTTPException(status_code=503, detail="Claude no disponible")

    quota = _check_and_increment_quota(uid)
    if not quota["allowed"]:
        return _quota_response(quota)

    model = _resolve_model(payload.model, ALLOWED_CLAUDE_MODELS, DEFAULT_CLAUDE_MODEL)

    try:
        messages: list[dict[str, Any]] = []
        for msg in payload.history[-MAX_HISTORY:]:
            role = "user" if msg.role == "user" else "assistant"
            messages.append({"role": role, "content": msg.text or "[imagen]"})

        user_content: list[dict[str, Any]] = []
        if payload.imageBase64:
            user_content.append(
                {
                    "type": "image",
                    "source": {
                        "type": "base64",
                        "media_type": "image/jpeg",
                        "data": payload.imageBase64,
                    },
                }
            )
        user_content.append(
            {"type": "text", "text": payload.userText or "Analiza esta imagen"}
        )
        messages.append({"role": "user", "content": user_content})

        res = await claude_client.messages.create(
            model=model,
            max_tokens=MAX_TOKENS,
            system=_build_claude_system(payload),
            messages=messages,
        )

        text_blocks = [
            b.text for b in res.content if getattr(b, "type", None) == "text"
        ]
        if not text_blocks:
            _decrement_quota(uid)
            logger.error("Claude respondió sin bloques de texto", extra=extra)
            raise HTTPException(status_code=502, detail="Respuesta vacía del modelo")

        _log_claude_usage(
            logger,
            getattr(res, "usage", None),
            route="messages",
            model=model,
            remaining=quota["remaining"],
            extra=extra,
        )
        return {
            "text": "\n".join(text_blocks),
            "remaining": quota["remaining"],
            "model": model,
        }

    except RateLimitError:
        _decrement_quota(uid)
        logger.warning("Claude upstream rate limit", extra=extra)
        return JSONResponse(
            status_code=429,
            content={
                "error": "upstream_rate_limit",
                "message": "El proveedor está saturado, intenta en un minuto",
            },
        )
    except BadRequestError as exc:
        _decrement_quota(uid)
        logger.warning("Claude bad request: %s", exc, extra=extra)
        raise HTTPException(status_code=400, detail=f"Solicitud inválida: {exc.message}") from exc
    except AuthenticationError as exc:
        _decrement_quota(uid)
        logger.error("Claude auth error — revisa CLAUDE_API_KEY", extra=extra)
        raise HTTPException(status_code=502, detail="Error de autenticación con el proveedor") from exc
    except (APITimeoutError, APIConnectionError) as exc:
        _decrement_quota(uid)
        logger.warning("Claude conexión: %s", exc, extra=extra)
        raise HTTPException(status_code=504, detail="Tiempo de espera agotado") from exc
    except APIStatusError as exc:
        _decrement_quota(uid)
        logger.error("Claude API error %s: %s", exc.status_code, exc, extra=extra)
        raise HTTPException(status_code=502, detail="Error del proveedor") from exc
    except HTTPException:
        raise
    except Exception as exc:
        _decrement_quota(uid)
        logger.exception("Claude error inesperado", extra=extra)
        raise HTTPException(
            status_code=500, detail=f"Error interno: {type(exc).__name__}"
        ) from exc


@app.post("/gemini")
async def gemini_proxy(
    request: Request,
    payload: AIRequest,
    user: dict = Depends(verify_token),
):
    rid = request.state.request_id
    extra = {"request_id": rid}
    uid = user["uid"]

    if not gemini_client:
        raise HTTPException(status_code=503, detail="Gemini no disponible")

    quota = _check_and_increment_quota(uid)
    if not quota["allowed"]:
        return _quota_response(quota)

    model_name = DEFAULT_GEMINI_MODEL
    logger.info("Gemini model=%s uid=%s", model_name, uid, extra=extra)

    try:
        # Build contents (history + current turn)
        contents: list[dict[str, Any]] = []
        for msg in payload.history[-MAX_HISTORY:]:
            role = "user" if msg.role == "user" else "model"
            contents.append({"role": role, "parts": [{"text": msg.text or "[imagen]"}]})

        current_parts: list[Any] = []
        if payload.imageBase64:
            try:
                img_bytes = base64.b64decode(payload.imageBase64, validate=True)
            except Exception as exc:
                _decrement_quota(uid)
                raise HTTPException(
                    status_code=400, detail=f"imageBase64 inválido: {exc}"
                ) from exc
            current_parts.append({"inline_data": {"mime_type": "image/jpeg", "data": img_bytes}})
        current_parts.append({"text": payload.userText or "Analiza esta imagen"})
        contents.append({"role": "user", "parts": current_parts})

        config = genai_types.GenerateContentConfig(
            system_instruction=payload.systemPrompt or None,
            safety_settings=[
                genai_types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="BLOCK_NONE"),
                genai_types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="BLOCK_NONE"),
                genai_types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="BLOCK_NONE"),
                genai_types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="BLOCK_NONE"),
            ],
        )

        text: Optional[str] = None
        last_rate_exc: Optional[Exception] = None
        for attempt in range(3):
            if attempt > 0:
                await asyncio.sleep(2 ** (attempt - 1))
            try:
                response = await gemini_client.aio.models.generate_content(
                    model=model_name,
                    contents=contents,
                    config=config,
                )
                text = response.text
                last_rate_exc = None
                break
            except genai_errors.APIError as exc:
                if exc.status_code == 429:
                    last_rate_exc = exc
                    logger.warning("Gemini rate limit attempt %d/3: %s", attempt + 1, exc, extra=extra)
                else:
                    raise

        if last_rate_exc is not None:
            _decrement_quota(uid)
            logger.warning("Gemini upstream rate limit exhausted", extra=extra)
            return JSONResponse(
                status_code=429,
                content={
                    "error": "upstream_rate_limit",
                    "message": "El proveedor está saturado, intenta en un minuto",
                },
            )

        if not text:
            _decrement_quota(uid)
            logger.error("Gemini respondió sin texto", extra=extra)
            raise HTTPException(status_code=502, detail="Respuesta vacía del modelo")

        logger.info("Gemini OK model=%s remaining=%s", model_name, quota["remaining"], extra=extra)
        return {"text": text, "remaining": quota["remaining"], "model": model_name}

    except genai_errors.APIError as exc:
        _decrement_quota(uid)
        if exc.status_code == 400:
            logger.warning("Gemini bad request: %s", exc, extra=extra)
            raise HTTPException(status_code=400, detail=f"Solicitud inválida: {exc}") from exc
        if exc.status_code in (401, 403):
            logger.error("Gemini auth error — revisa GEMINI_API_KEY", extra=extra)
            raise HTTPException(status_code=502, detail="Error de autenticación con el proveedor") from exc
        if exc.status_code == 504:
            logger.warning("Gemini timeout", extra=extra)
            raise HTTPException(status_code=504, detail="Tiempo de espera agotado") from exc
        logger.error("Gemini API error %s: %s", exc.status_code, exc, extra=extra)
        raise HTTPException(status_code=502, detail="Error del proveedor") from exc
    except HTTPException:
        raise
    except Exception as exc:
        _decrement_quota(uid)
        logger.exception("Gemini error inesperado", extra=extra)
        raise HTTPException(
            status_code=500, detail=f"Error interno: {type(exc).__name__}"
        ) from exc


# --- Endpoints streaming (SSE) ---
SSE_HEADERS = {
    "Cache-Control": "no-cache, no-transform",
    "X-Accel-Buffering": "no",
    "Connection": "keep-alive",
    "Content-Type": "text/event-stream; charset=utf-8",
    # Defeat any compression middleware in HF's edge that would buffer the body.
    "Content-Encoding": "identity",
}

# 2KB of padding sent as an SSE comment to defeat proxy/edge buffering.
# Many reverse proxies (incl. HF Spaces' edge) hold the response until they
# accumulate ~2KB before flushing the first chunk. The padding ensures the
# client sees deltas as soon as they arrive.
_SSE_PADDING = ":" + (" " * 2048) + "\n\n"


def _sse(event: dict) -> str:
    return f"data: {json.dumps(event, ensure_ascii=False)}\n\n"


def _build_claude_system(payload: AIRequest) -> Any:
    """Envuelve el system prompt en un bloque cacheable de Anthropic.
    Devuelve "" cuando no hay prompt (Anthropic acepta string vacío). Cuando hay
    prompt, lo envuelve en un bloque con cache_control=ephemeral para que turnos
    consecutivos paguen sólo $0.30/M input en vez de $3/M. Si el prompt es <1024
    tokens Anthropic no cachea — el flag se ignora silenciosamente, sin error.
    """
    if not payload.systemPrompt:
        return ""
    return [
        {
            "type": "text",
            "text": payload.systemPrompt,
            "cache_control": {"type": "ephemeral"},
        }
    ]


def _log_claude_usage(logger_, usage, *, route: str, model: str, remaining: Any, extra: dict) -> None:
    """Loguea tokens input/output/cache de la respuesta de Claude.
    Sin esto no hay forma de saber si el prompt cache está pegando. El ratio
    cache_read / (cache_read + input_tokens) debería tender a >70% en sesiones
    con varios turnos seguidos; si es <30% algo invalida el cache.
    """
    cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0
    cache_create = getattr(usage, "cache_creation_input_tokens", 0) or 0
    input_tok = getattr(usage, "input_tokens", 0) or 0
    output_tok = getattr(usage, "output_tokens", 0) or 0
    logger_.info(
        "Claude %s OK model=%s in=%d out=%d cache_read=%d cache_create=%d remaining=%s",
        route, model, input_tok, output_tok, cache_read, cache_create, remaining,
        extra=extra,
    )


def _build_claude_messages(payload: AIRequest) -> list[dict[str, Any]]:
    """Construye messages con alternancia user/assistant estricta + turno actual."""
    messages: list[dict[str, Any]] = []
    for msg in payload.history[-MAX_HISTORY:]:
        role = "user" if msg.role == "user" else "assistant"
        if not messages or messages[-1]["role"] != role:
            messages.append({"role": role, "content": msg.text or "[imagen]"})
    if messages and messages[-1]["role"] == "user":
        messages.pop()

    user_content: list[dict[str, Any]] = []
    if payload.imageBase64:
        user_content.append(
            {
                "type": "image",
                "source": {
                    "type": "base64",
                    "media_type": "image/jpeg",
                    "data": payload.imageBase64,
                },
            }
        )
    user_content.append(
        {"type": "text", "text": payload.userText or "Analiza esta imagen"}
    )
    messages.append({"role": "user", "content": user_content})
    return messages


@app.post("/claude/stream")
async def claude_stream(
    request: Request,
    payload: AIRequest,
    user: dict = Depends(verify_token),
):
    rid = request.state.request_id
    extra = {"request_id": rid}
    uid = user["uid"]

    if claude_client is None:
        raise HTTPException(status_code=503, detail="Claude no disponible")

    quota = _check_and_increment_quota(uid)
    if not quota["allowed"]:
        return _quota_response(quota)

    model = _resolve_model(payload.model, ALLOWED_CLAUDE_MODELS, DEFAULT_CLAUDE_MODEL)
    messages = _build_claude_messages(payload)

    async def event_stream():
        # Flush proxy buffer ASAP so the client sees deltas in real time.
        yield _SSE_PADDING
        yield _sse({"start": True, "model": model})
        await asyncio.sleep(0)

        yielded = False
        try:
            final_message = None
            async with claude_client.messages.stream(
                model=model,
                max_tokens=MAX_TOKENS,
                system=_build_claude_system(payload),
                messages=messages,
            ) as stream:
                async for text in stream.text_stream:
                    yielded = True
                    yield _sse({"delta": text})
                    await asyncio.sleep(0)
                final_message = await stream.get_final_message()
            yield _sse(
                {
                    "done": True,
                    "remaining": quota["remaining"],
                    "model": model,
                }
            )
            _log_claude_usage(
                logger,
                getattr(final_message, "usage", None),
                route="stream",
                model=model,
                remaining=quota["remaining"],
                extra=extra,
            )
        except RateLimitError:
            if not yielded:
                _decrement_quota(uid)
            logger.warning("Claude stream rate limit", extra=extra)
            yield _sse({"error": "upstream_rate_limit", "retryable": True})
        except (APITimeoutError, APIConnectionError) as exc:
            if not yielded:
                _decrement_quota(uid)
            logger.warning("Claude stream conexión: %s", exc, extra=extra)
            yield _sse({"error": "timeout", "retryable": True})
        except (AuthenticationError, BadRequestError, APIStatusError) as exc:
            if not yielded:
                _decrement_quota(uid)
            logger.error("Claude stream upstream error: %s", exc, extra=extra)
            yield _sse({"error": "upstream_error", "message": str(exc)})
        except Exception as exc:
            if not yielded:
                _decrement_quota(uid)
            logger.exception("Claude stream error inesperado", extra=extra)
            yield _sse({"error": "internal", "message": type(exc).__name__})

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers=SSE_HEADERS,
    )


@app.post("/gemini/stream")
async def gemini_stream(
    request: Request,
    payload: AIRequest,
    user: dict = Depends(verify_token),
):
    rid = request.state.request_id
    extra = {"request_id": rid}
    uid = user["uid"]

    if not gemini_client:
        raise HTTPException(status_code=503, detail="Gemini no disponible")

    quota = _check_and_increment_quota(uid)
    if not quota["allowed"]:
        return _quota_response(quota)

    model_name = DEFAULT_GEMINI_MODEL
    logger.info("Gemini stream model=%s uid=%s", model_name, uid, extra=extra)

    # Decode image upfront so we fail before opening the stream
    img_bytes: Optional[bytes] = None
    if payload.imageBase64:
        try:
            img_bytes = base64.b64decode(payload.imageBase64, validate=True)
        except Exception as exc:
            _decrement_quota(uid)
            raise HTTPException(
                status_code=400, detail=f"imageBase64 inválido: {exc}"
            ) from exc

    # Build contents once (shared across retry attempts)
    contents: list[dict[str, Any]] = []
    for msg in payload.history[-MAX_HISTORY:]:
        role = "user" if msg.role == "user" else "model"
        if not contents or contents[-1]["role"] != role:
            contents.append({"role": role, "parts": [{"text": msg.text or "[imagen]"}]})

    current_parts: list[Any] = []
    if img_bytes is not None:
        current_parts.append({"inline_data": {"mime_type": "image/jpeg", "data": img_bytes}})
    current_parts.append({"text": payload.userText or "Analiza esta imagen"})
    contents.append({"role": "user", "parts": current_parts})

    config = genai_types.GenerateContentConfig(
        system_instruction=payload.systemPrompt or None,
        safety_settings=[
            genai_types.SafetySetting(category="HARM_CATEGORY_HARASSMENT", threshold="BLOCK_NONE"),
            genai_types.SafetySetting(category="HARM_CATEGORY_HATE_SPEECH", threshold="BLOCK_NONE"),
            genai_types.SafetySetting(category="HARM_CATEGORY_SEXUALLY_EXPLICIT", threshold="BLOCK_NONE"),
            genai_types.SafetySetting(category="HARM_CATEGORY_DANGEROUS_CONTENT", threshold="BLOCK_NONE"),
        ],
    )

    async def event_stream():
        yield _SSE_PADDING
        yield _sse({"start": True, "model": model_name})
        await asyncio.sleep(0)

        yielded = False
        last_rate_exc: Optional[Exception] = None

        for attempt in range(3):
            if attempt > 0:
                await asyncio.sleep(2 ** (attempt - 1))
            last_rate_exc = None
            try:
                async for chunk in gemini_client.aio.models.generate_content_stream(
                    model=model_name,
                    contents=contents,
                    config=config,
                ):
                    if chunk.text:
                        yielded = True
                        yield _sse({"delta": chunk.text})
                        await asyncio.sleep(0)

                yield _sse({"done": True, "remaining": quota["remaining"], "model": model_name})
                logger.info("Gemini stream OK model=%s", model_name, extra=extra)
                return

            except genai_errors.APIError as exc:
                if exc.status_code == 429:
                    last_rate_exc = exc
                    if yielded:
                        break
                    logger.warning("Gemini stream rate limit attempt %d/3: %s", attempt + 1, exc, extra=extra)
                else:
                    if not yielded:
                        _decrement_quota(uid)
                    logger.error("Gemini stream API error %s: %s", exc.status_code, exc, extra=extra)
                    yield _sse({"error": "upstream_error", "message": str(exc)})
                    return
            except asyncio.TimeoutError:
                if not yielded:
                    _decrement_quota(uid)
                logger.warning("Gemini stream timeout", extra=extra)
                yield _sse({"error": "timeout", "retryable": True})
                return
            except Exception as exc:
                if not yielded:
                    _decrement_quota(uid)
                logger.exception("Gemini stream error inesperado", extra=extra)
                yield _sse({"error": "internal", "message": type(exc).__name__})
                return

        if not yielded:
            _decrement_quota(uid)
        logger.warning("Gemini stream rate limit exhausted", extra=extra)
        yield _sse({"error": "upstream_rate_limit", "retryable": True})

    return StreamingResponse(
        event_stream(),
        media_type="text/event-stream",
        headers=SSE_HEADERS,
    )


@app.get("/")
def health():
    return {
        "status": "Surco Online",
        "version": app.version,
        "firebase": db is not None,
        "firebase_error": firebase_error,
        "claude": claude_client is not None,
        "gemini": GEMINI_API_KEY is not None,
        "streaming": True,
        "time": datetime.now(BOGOTA_TZ).isoformat(),
    }


@app.get("/healthz")
def healthz():
    return {"ok": True}