core/api_client.py

"""
core/api_client.py — MiniCPM-V API wrapper.

Handles:
  • OpenAI-compatible streaming via the `openai` SDK
  • API key resolution (env var → UI field → public fallback)
  • Friendly error messages for 401 / 429 / network failures
"""

import os
from collections.abc import Generator

from openai import OpenAI, APIStatusError, APIConnectionError
from PIL import Image

from config import API_BASE_URL, PUBLIC_API_KEY, MODELS
from core.image_utils import pil_to_data_url


def _resolve_key(ui_key: str) -> str:
    """
    Priority: env var MINICPM_API_KEY → UI field → hardcoded public key.
    """
    return (
        os.environ.get("MINICPM_API_KEY", "").strip()
        or (ui_key or "").strip()
        or PUBLIC_API_KEY
    )


def stream_description(
    image: Image.Image,
    prompt: str,
    model_label: str,
    max_tokens: int,
    temperature: float,
    api_key: str,
) -> Generator[str, None, None]:
    """
    Stream the model's response token-by-token.

    Yields:
        Cumulative response string — each yield is the full text so far,
        so Gradio can update the textbox incrementally.
    """
    if image is None:
        yield "⚠️  Please upload an image first."
        return

    key      = _resolve_key(api_key)
    model_id = MODELS[model_label]
    data_url = pil_to_data_url(image)

    client = OpenAI(api_key=key, base_url=API_BASE_URL)

    try:
        stream = client.chat.completions.create(
            model=model_id,
            messages=[{
                "role": "user",
                "content": [
                    {"type": "image_url", "image_url": {"url": data_url}},
                    {"type": "text",      "text": prompt},
                ],
            }],
            max_tokens=max_tokens,
            temperature=temperature,
            stream=True,
        )

        result = ""
        for chunk in stream:
            delta = chunk.choices[0].delta.content or ""
            if delta:
                result += delta
                yield result

    except APIStatusError as e:
        if e.status_code == 401:
            yield (
                "🔐 **API key rejected (401).**\n\n"
                "The public key may have hit its rate limit.  \n"
                "Get a free personal key at https://modelbest.cn and paste it "
                "in the **API Key** field."
            )
        elif e.status_code == 429:
            yield (
                "⏳ **Rate limit reached (429).**\n\n"
                "The free public key is shared — wait a moment and retry, "
                "or use your own key from https://modelbest.cn"
            )
        else:
            yield f"❌ API error {e.status_code}:\n\n```\n{e.message}\n```"

    except APIConnectionError:
        yield (
            "❌ **Cannot reach the API** (`api.modelbest.cn`).  \n"
            "Check your internet connection and try again."
        )

    except Exception as e:
        yield f"❌ Unexpected error:\n\n```\n{e}\n```"