Spaces:

Romanes
/

Surco

Sleeping

Surco / main.py

fix: simplify gemini stream handler, remove broken retry loop

3601c4b verified 5 months ago

22.8 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",
)
_rid_filter = _RequestIdFilter()
for _h in logging.root.handlers:
    _h.addFilter(_rid_filter)
logger = logging.getLogger("surco_proxy")
logger.addFilter(_rid_filter)

# --- 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 ---
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 ---
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.2.0")

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
        return (doc.to_dict() or {}).get("plan", DEFAULT_AI_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:
        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)


def _resolve_model(requested: Optional[str], allowed: set, default: str) -> str:
    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"),
        },
    )


def _gemini_safety() -> list[genai_types.SafetySetting]:
    return [
        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"),
    ]


# --- Endpoints Claude ---
@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=payload.systemPrompt or "",
            messages=messages,
        )

        text_blocks = [b.text for b in res.content if getattr(b, "type", None) == "text"]
        if not text_blocks:
            _decrement_quota(uid)
            raise HTTPException(status_code=502, detail="Respuesta vacía del modelo")

        logger.info("Claude OK model=%s remaining=%s", model, 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"})
    except BadRequestError as exc:
        _decrement_quota(uid)
        raise HTTPException(status_code=400, detail=f"Solicitud inválida: {exc.message}") from exc
    except AuthenticationError as exc:
        _decrement_quota(uid)
        raise HTTPException(status_code=502, detail="Error de autenticación con el proveedor") from exc
    except (APITimeoutError, APIConnectionError) as exc:
        _decrement_quota(uid)
        raise HTTPException(status_code=504, detail="Tiempo de espera agotado") from exc
    except APIStatusError as exc:
        _decrement_quota(uid)
        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


# --- Endpoints Gemini ---
@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:
        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=_gemini_safety(),
        )

        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 "429" in str(exc) or "RESOURCE_EXHAUSTED" in str(exc):
                    last_rate_exc = exc
                    logger.warning("Gemini rate limit attempt %d/3", attempt + 1, extra=extra)
                else:
                    raise

        if last_rate_exc is not None:
            _decrement_quota(uid)
            return JSONResponse(status_code=429, content={"error": "upstream_rate_limit", "message": "El proveedor está saturado"})

        if not text:
            _decrement_quota(uid)
            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)
        exc_str = str(exc)
        if "401" in exc_str or "403" in exc_str or "UNAUTHENTICATED" in exc_str or "PERMISSION_DENIED" in exc_str:
            raise HTTPException(status_code=502, detail="Error de autenticación con el proveedor") from exc
        logger.error("Gemini API error: %s", exc_str[:200], 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


# --- SSE helpers ---
SSE_HEADERS = {
    "Cache-Control": "no-cache, no-transform",
    "X-Accel-Buffering": "no",
    "Connection": "keep-alive",
    "Content-Type": "text/event-stream; charset=utf-8",
}

# 2KB padding para vencer el buffering de proxies/edge.
_SSE_PADDING = ": " + ("surco" * 400) + "\n\n"


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


def _build_claude_messages(payload: AIRequest) -> list[dict[str, Any]]:
    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():
        yielded = False
        try:
            async with claude_client.messages.stream(
                model=model,
                max_tokens=MAX_TOKENS,
                system=payload.systemPrompt or "",
                messages=messages,
            ) as stream:
                async for text in stream.text_stream:
                    yielded = True
                    yield _sse({"delta": text})
            yield _sse({"done": True, "remaining": quota["remaining"], "model": model})
            logger.info("Claude stream OK model=%s", model, 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)

    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

    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=_gemini_safety(),
    )

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

        try:
            _stream = await gemini_client.aio.models.generate_content_stream(
                model=model_name, contents=contents, config=config,
            )
            async for chunk in _stream:
                if chunk.text:
                    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)
        except genai_errors.APIError as exc:
            _decrement_quota(uid)
            exc_str = str(exc)
            is_429 = "429" in exc_str or "RESOURCE_EXHAUSTED" in exc_str
            if is_429:
                logger.warning("Gemini stream rate limit/quota: %s", exc_str[:120], extra=extra)
                yield _sse({"error": "upstream_rate_limit", "retryable": True})
            else:
                logger.error("Gemini stream API error: %s", exc_str[:200], extra=extra)
                yield _sse({"error": "upstream_error", "message": exc_str})
        except asyncio.TimeoutError:
            _decrement_quota(uid)
            logger.warning("Gemini stream timeout", extra=extra)
            yield _sse({"error": "timeout", "retryable": True})
        except Exception as exc:
            _decrement_quota(uid)
            logger.exception("Gemini 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.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_client is not None,
        "streaming": True,
        "time": datetime.now(BOGOTA_TZ).isoformat(),
    }


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

	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",
	)
	_rid_filter = _RequestIdFilter()
	for _h in logging.root.handlers:
	_h.addFilter(_rid_filter)
	logger = logging.getLogger("surco_proxy")
	logger.addFilter(_rid_filter)

	# --- 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 ---
	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 ---
	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.2.0")

	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
	return (doc.to_dict() or {}).get("plan", DEFAULT_AI_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:
	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)


	def _resolve_model(requested: Optional[str], allowed: set, default: str) -> str:
	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"),
	},
	)


	def _gemini_safety() -> list[genai_types.SafetySetting]:
	return [
	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"),
	]


	# --- Endpoints Claude ---
	@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=payload.systemPrompt or "",
	messages=messages,
	)

	text_blocks = [b.text for b in res.content if getattr(b, "type", None) == "text"]
	if not text_blocks:
	_decrement_quota(uid)
	raise HTTPException(status_code=502, detail="Respuesta vacía del modelo")

	logger.info("Claude OK model=%s remaining=%s", model, 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"})
	except BadRequestError as exc:
	_decrement_quota(uid)
	raise HTTPException(status_code=400, detail=f"Solicitud inválida: {exc.message}") from exc
	except AuthenticationError as exc:
	_decrement_quota(uid)
	raise HTTPException(status_code=502, detail="Error de autenticación con el proveedor") from exc
	except (APITimeoutError, APIConnectionError) as exc:
	_decrement_quota(uid)
	raise HTTPException(status_code=504, detail="Tiempo de espera agotado") from exc
	except APIStatusError as exc:
	_decrement_quota(uid)
	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


	# --- Endpoints Gemini ---
	@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:
	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=_gemini_safety(),
	)

	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 "429" in str(exc) or "RESOURCE_EXHAUSTED" in str(exc):
	last_rate_exc = exc
	logger.warning("Gemini rate limit attempt %d/3", attempt + 1, extra=extra)
	else:
	raise

	if last_rate_exc is not None:
	_decrement_quota(uid)
	return JSONResponse(status_code=429, content={"error": "upstream_rate_limit", "message": "El proveedor está saturado"})

	if not text:
	_decrement_quota(uid)
	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)
	exc_str = str(exc)
	if "401" in exc_str or "403" in exc_str or "UNAUTHENTICATED" in exc_str or "PERMISSION_DENIED" in exc_str:
	raise HTTPException(status_code=502, detail="Error de autenticación con el proveedor") from exc
	logger.error("Gemini API error: %s", exc_str[:200], 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


	# --- SSE helpers ---
	SSE_HEADERS = {
	"Cache-Control": "no-cache, no-transform",
	"X-Accel-Buffering": "no",
	"Connection": "keep-alive",
	"Content-Type": "text/event-stream; charset=utf-8",
	}

	# 2KB padding para vencer el buffering de proxies/edge.
	_SSE_PADDING = ": " + ("surco" * 400) + "\n\n"


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


	def _build_claude_messages(payload: AIRequest) -> list[dict[str, Any]]:
	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():
	yielded = False
	try:
	async with claude_client.messages.stream(
	model=model,
	max_tokens=MAX_TOKENS,
	system=payload.systemPrompt or "",
	messages=messages,
	) as stream:
	async for text in stream.text_stream:
	yielded = True
	yield _sse({"delta": text})
	yield _sse({"done": True, "remaining": quota["remaining"], "model": model})
	logger.info("Claude stream OK model=%s", model, 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)

	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

	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=_gemini_safety(),
	)

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

	try:
	_stream = await gemini_client.aio.models.generate_content_stream(
	model=model_name, contents=contents, config=config,
	)
	async for chunk in _stream:
	if chunk.text:
	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)
	except genai_errors.APIError as exc:
	_decrement_quota(uid)
	exc_str = str(exc)
	is_429 = "429" in exc_str or "RESOURCE_EXHAUSTED" in exc_str
	if is_429:
	logger.warning("Gemini stream rate limit/quota: %s", exc_str[:120], extra=extra)
	yield _sse({"error": "upstream_rate_limit", "retryable": True})
	else:
	logger.error("Gemini stream API error: %s", exc_str[:200], extra=extra)
	yield _sse({"error": "upstream_error", "message": exc_str})
	except asyncio.TimeoutError:
	_decrement_quota(uid)
	logger.warning("Gemini stream timeout", extra=extra)
	yield _sse({"error": "timeout", "retryable": True})
	except Exception as exc:
	_decrement_quota(uid)
	logger.exception("Gemini 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.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_client is not None,
	"streaming": True,
	"time": datetime.now(BOGOTA_TZ).isoformat(),
	}


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