llm_client.py

# llm_client.py
import os
import time
from pathlib import Path
from typing import Callable, Type, TypeVar

from dotenv import load_dotenv
from pydantic import BaseModel

load_dotenv()

T = TypeVar("T", bound=BaseModel)
MAX_RETRIES = 2
BACKOFF_BASE = 2
GEMINI_FALLBACK_MODEL = "gemini-2.5-flash"


def _with_retry(fn, *args, **kwargs):
    last_exc = None
    for attempt in range(MAX_RETRIES):
        try:
            return fn(*args, **kwargs)
        except Exception as e:
            last_exc = e
            if attempt < MAX_RETRIES - 1:
                wait = BACKOFF_BASE ** attempt
                print(f"  LLM call failed ({e}), retrying in {wait}s...")
                time.sleep(wait)
    raise last_exc


def _is_quota_exhausted(exc: Exception) -> bool:
    msg = str(exc)
    return (
        "RESOURCE_EXHAUSTED" in msg or "429" in msg or "quota" in msg.lower()
        or "503" in msg or "UNAVAILABLE" in msg or "high demand" in msg.lower()
    )


def _call_gemini(model: str, images: list[Path], prompt: str, response_model: Type[T]) -> T:
    import google.genai as genai
    from google.genai import types

    client = genai.Client(api_key=os.getenv("GEMINI_API_KEY"))

    parts = []
    for img_path in images:
        if img_path.exists():
            parts.append(types.Part.from_bytes(data=img_path.read_bytes(), mime_type="image/jpeg"))
    parts.append(types.Part.from_text(text=prompt))

    response = client.models.generate_content(
        model=model,
        contents=parts,
        config=types.GenerateContentConfig(
            response_mime_type="application/json",
            response_schema=response_model,
        ),
    )
    return response.parsed


def _call_openai_compatible(base_url: str | None, api_key: str, model: str,
                              images: list[Path], prompt: str, response_model: Type[T]) -> T:
    import base64
    from openai import OpenAI

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

    content = []
    for img_path in images:
        if img_path.exists():
            b64 = base64.b64encode(img_path.read_bytes()).decode()
            content.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{b64}"}})
    content.append({"type": "text", "text": prompt})

    # Use native OpenAI structured outputs (no instructor dependency)
    response = client.beta.chat.completions.parse(
        model=model,
        messages=[{"role": "user", "content": content}],
        response_format=response_model,
    )
    return response.choices[0].message.parsed


def get_llm_client(provider: str, model: str) -> Callable:
    """Returns call_llm(images, prompt, response_model) -> validated Pydantic instance."""

    if provider == "gemini":
        def call(images, prompt, rm):
            try:
                return _with_retry(_call_gemini, model, images, prompt, rm)
            except Exception as e:
                if _is_quota_exhausted(e):
                    print(f"  Gemini quota exhausted for {model!r}, falling back to {GEMINI_FALLBACK_MODEL!r}")
                    return _with_retry(_call_gemini, GEMINI_FALLBACK_MODEL, images, prompt, rm)
                raise
        return call

    if provider == "openai":
        def call(images, prompt, rm):
            return _with_retry(_call_openai_compatible, None,
                               os.getenv("OPENAI_API_KEY"), model, images, prompt, rm)
        return call

    if provider == "qwen":
        def call(images, prompt, rm):
            return _with_retry(_call_openai_compatible,
                               "https://openrouter.ai/api/v1",
                               os.getenv("OPENROUTER_API_KEY"),
                               model, images, prompt, rm)
        return call

    raise ValueError(f"Unknown provider: {provider!r}")