poe_analyze.py

import os
import uuid
import base64
import requests
import fastapi_poe as fp
import json

from .s3_storage import s3_enabled, upload_bytes
TEXT_BOT_NAME = "gpt-5.2-instant"

OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions"
# Use flux.2-pro (documented for image gen). Override with OPENROUTER_IMAGE_MODEL env.

# OPENROUTER_IMAGE_MODEL = os.getenv("OPENROUTER_IMAGE_MODEL", "black-forest-labs/flux.2-pro")
_primary_image_model = os.getenv("OPENROUTER_IMAGE_MODEL", "google/gemini-2.5-flash-image")
_fallback_model_env = os.getenv("OPENROUTER_IMAGE_MODELS", "").strip()
if _fallback_model_env:
    OPENROUTER_IMAGE_MODELS = [m.strip() for m in _fallback_model_env.split(",") if m.strip()]
else:
    # Put region-restricted model first, then a commonly available fallback.
    OPENROUTER_IMAGE_MODELS = [_primary_image_model, "black-forest-labs/flux-1-schnell"]
# -----------------------------
# Poe: 文字 + image_prompt
# -----------------------------
def analyze_with_poe(prompt: str) -> dict:
    api_key = os.getenv("POE_API_KEY")
    if not api_key:
        raise RuntimeError("POE_API_KEY missing")

    message = fp.ProtocolMessage(role="user", content=prompt)
    text = ""

    for partial in fp.get_bot_response_sync(
        messages=[message],
        bot_name=TEXT_BOT_NAME,
        api_key=api_key
    ):
        if hasattr(partial, "text") and partial.text:
            text += partial.text

    try:
        return json.loads(text)
    except json.JSONDecodeError:
        raise RuntimeError(f"Poe did not return valid JSON:\n{text[:500]}")

# -----------------------------
# OpenRouter: 圖片生成
# -----------------------------
def generate_kitten_image(image_prompt: str, logger=None) -> dict:
    api_key = os.getenv("OPENROUTER_API_KEY")
    upload_dir = os.getenv("UPLOAD_DIR", "./uploads")

    if not api_key:
        return {"success": False, "error": "OPENROUTER_API_KEY missing"}

    if not s3_enabled():
        os.makedirs(upload_dir, exist_ok=True)

    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
        "HTTP-Referer": "https://pawmatch.ai",
        "X-Title": "PawMatch",
    }

    last_error = "OpenRouter image generation failed"
    tried_models = []

    for model in OPENROUTER_IMAGE_MODELS:
        payload = {
            "model": model,
            "modalities": ["image", "text"],
            "messages": [
                {"role": "user", "content": image_prompt}
            ],
        }
        tried_models.append(model)

        try:
            r = requests.post(
                OPENROUTER_URL,
                headers=headers,
                json=payload,
                timeout=90
            )

            if not r.ok:
                err_body = r.text[:500] if r.text else "(no body)"
                last_error = f"OpenRouter {r.status_code}: {err_body}"
                if logger:
                    logger.warning(f"OpenRouter response for model {model}: {last_error}")

                # Region-locked model: continue to next fallback.
                if r.status_code == 403 and "not available in your region" in err_body.lower():
                    continue

                # Try fallback for provider/transient failures as well.
                if r.status_code in (429, 500, 502, 503, 504):
                    continue

                # Non-retryable response for this request.
                break

            data = r.json()
            message = data["choices"][0]["message"]
            images = message.get("images")

            if not images:
                last_error = "OpenRouter returned no images"
                if logger:
                    logger.warning(f"No images returned for model {model}. Message: {message}")
                continue

            image_url = images[0]["image_url"]["url"]

            if image_url.startswith("data:"):
                b64 = image_url.split(",", 1)[1]
                img_bytes = base64.b64decode(b64)
            else:
                # Some models return a direct URL instead of data URL.
                img_resp = requests.get(image_url, timeout=30)
                img_resp.raise_for_status()
                img_bytes = img_resp.content

            if s3_enabled():
                info = upload_bytes(
                    prefix="AI_PHOTO",
                    ext="png",
                    body=img_bytes,
                    content_type="image/png",
                )
                if logger:
                    logger.info(f"Image uploaded to S3 via model {model}: {info['s3_uri']}")
                return {
                    "success": True,
                    "s3_uri": info["s3_uri"],
                    "image_url": None,
                    "model": model,
                    "generation_id": info["object_id"],
                }

            generation_id = uuid.uuid4().hex
            filename = f"kitten_{generation_id}.png"
            path = os.path.join(upload_dir, filename)
            with open(path, "wb") as f:
                f.write(img_bytes)

            if logger:
                logger.info(f"Image saved via model {model}: {path}")

            return {
                "success": True,
                "s3_uri": None,
                "image_url": f"/uploads/{filename}",
                "model": model,
                "generation_id": generation_id,
            }

        except Exception as e:
            last_error = str(e)
            if logger:
                logger.exception(f"OpenRouter image generation failed on model {model}")
            continue

    return {
        "success": False,
        "error": last_error,
        "tried_models": tried_models
    }