models.py

"""
models.py
=========
The model-agnostic layer. This is the part that lets you use ANY provider for
either the drawing model or the chorus, and swap them by editing config.yaml
rather than touching code.

HOW IT WORKS (plain language):
Your program only ever speaks ONE internal language: "here is a system prompt,
here is the conversation so far, here (optionally) is an image of the canvas —
give me back some text." Each provider (Anthropic, OpenAI, ...) has its own
slightly different way of being asked that. An "adapter" is a small translator
for one provider. Add a provider = add one adapter. Nothing else changes.

The drawing model MUST support images (it sees its canvas). The chorus models
do NOT need images (they're blind by design), so they can be cheaper/smaller.

API KEYS: read from environment variables so you never paste secrets into
files. e.g. export ANTHROPIC_API_KEY=your-anthropic-api-key-here,
export OPENAI_API_KEY=your-openai-api-key-here, or
export OPENROUTER_API_KEY=your-openrouter-api-key-here.
"""

from __future__ import annotations
import os
import base64
import io
from dataclasses import dataclass
from typing import Optional
from PIL import Image

from env_utils import load_dotenv


load_dotenv()


# ---------------------------------------------------------------------------
# A standard message format we use everywhere internally.
# role is "user" or "assistant"; text is the content; image is optional.
# ---------------------------------------------------------------------------
@dataclass
class Msg:
    role: str
    text: str


def _img_to_b64(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    return base64.b64encode(buf.getvalue()).decode("utf-8")


# ===========================================================================
# Base adapter — defines the one method every provider must implement.
# ===========================================================================
class BaseAdapter:
    def __init__(self, model: str, temperature: float = 1.0):
        self.model = model
        self.temperature = temperature

    def complete(self, system: str, messages: list[Msg],
                 image: Optional[Image.Image] = None,
                 max_tokens: int = 1500) -> str:
        raise NotImplementedError


# ===========================================================================
# Anthropic adapter
# ===========================================================================
class AnthropicAdapter(BaseAdapter):
    def __init__(self, model: str, temperature: float = 1.0):
        super().__init__(model, temperature)
        import anthropic  # imported lazily so you only need installed SDKs
        self.client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

    def complete(self, system, messages, image=None, max_tokens=1500):
        api_messages = []
        for i, m in enumerate(messages):
            # attach the image to the most recent user message, if provided
            is_last_user = (i == len(messages) - 1 and m.role == "user")
            if image is not None and is_last_user:
                api_messages.append({
                    "role": "user",
                    "content": [
                        {"type": "image", "source": {
                            "type": "base64", "media_type": "image/png",
                            "data": _img_to_b64(image)}},
                        {"type": "text", "text": m.text},
                    ],
                })
            else:
                api_messages.append({"role": m.role, "content": m.text})

        resp = self.client.messages.create(
            model=self.model, system=system, messages=api_messages,
            max_tokens=max_tokens, temperature=self.temperature,
        )
        return "".join(b.text for b in resp.content if b.type == "text")


# ===========================================================================
# OpenAI adapter
# ===========================================================================
class OpenAIAdapter(BaseAdapter):
    def __init__(self, model: str, temperature: float = 1.0):
        super().__init__(model, temperature)
        from openai import OpenAI
        self.client = OpenAI(api_key=os.environ["OPENAI_API_KEY"])

    def complete(self, system, messages, image=None, max_tokens=1500):
        api_messages = [{"role": "system", "content": system}]
        for i, m in enumerate(messages):
            is_last_user = (i == len(messages) - 1 and m.role == "user")
            if image is not None and is_last_user:
                api_messages.append({
                    "role": "user",
                    "content": [
                        {"type": "text", "text": m.text},
                        {"type": "image_url", "image_url": {
                            "url": f"data:image/png;base64,{_img_to_b64(image)}"}},
                    ],
                })
            else:
                api_messages.append({"role": m.role, "content": m.text})

        resp = self.client.chat.completions.create(
            model=self.model, messages=api_messages,
            max_tokens=max_tokens, temperature=self.temperature,
        )
        return resp.choices[0].message.content or ""


# ===========================================================================
# OpenRouter adapter — OpenAI-compatible API with many routed models
# ===========================================================================
class OpenRouterAdapter(BaseAdapter):
    def __init__(self, model: str, temperature: float = 1.0):
        super().__init__(model, temperature)
        from openai import OpenAI

        headers = {}
        referer = os.environ.get("OPENROUTER_HTTP_REFERER")
        title = os.environ.get("OPENROUTER_APP_TITLE")
        if referer:
            headers["HTTP-Referer"] = referer
        if title:
            headers["X-Title"] = title

        self.client = OpenAI(
            api_key=os.environ["OPENROUTER_API_KEY"],
            base_url="https://openrouter.ai/api/v1",
            default_headers=headers or None,
        )

    def complete(self, system, messages, image=None, max_tokens=1500):
        api_messages = [{"role": "system", "content": system}]
        for i, m in enumerate(messages):
            is_last_user = (i == len(messages) - 1 and m.role == "user")
            if image is not None and is_last_user:
                api_messages.append({
                    "role": "user",
                    "content": [
                        {"type": "text", "text": m.text},
                        {"type": "image_url", "image_url": {
                            "url": f"data:image/png;base64,{_img_to_b64(image)}"}},
                    ],
                })
            else:
                api_messages.append({"role": m.role, "content": m.text})

        resp = self.client.chat.completions.create(
            model=self.model, messages=api_messages,
            max_tokens=max_tokens, temperature=self.temperature,
        )
        return resp.choices[0].message.content or ""


# ===========================================================================
# Factory: turn a config dict into the right adapter.
# config looks like: {"provider": "anthropic", "model": "...", "temperature": 1.0}
# ===========================================================================
ADAPTERS = {
    "anthropic": AnthropicAdapter,
    "openai": OpenAIAdapter,
    "openrouter": OpenRouterAdapter,
}


def make_model(cfg: dict) -> BaseAdapter:
    provider = cfg.get("provider", "anthropic").lower()
    if provider not in ADAPTERS:
        raise ValueError(
            f"Unknown provider '{provider}'. "
            f"Available: {list(ADAPTERS)}. "
            f"To add one, write a new adapter class in models.py.")
    return ADAPTERS[provider](
        model=cfg["model"],
        temperature=float(cfg.get("temperature", 1.0)),
    )