qwen_test_runner/providers/claude_api.py

"""
claude_api.py — Anthropic Claude as a caption-processing provider.

Uses Anthropic's tool-calling API with forced tool choice and CAPTION_JSON_SCHEMA
as the tool's input_schema. The model is constrained to emit JSON matching the
schema — equivalent in semantic guarantee to xgrammar-constrained Qwen output,
but produced by a far more capable model.

Primary use cases:
  1. Teacher labels for SFT — generate gold structured outputs from real
     captions (COCO, LAION, Flickr30k) to fine-tune Qwen3.5-0.8B on.
  2. Comparison baseline — see what near-perfect schema/faithfulness numbers
     look like on the same eval set Qwen runs against.

Requires:
  pip install anthropic
  export ANTHROPIC_API_KEY=sk-ant-...

Cost note: at ~250 input tokens + ~200 output tokens per caption, Claude Sonnet
costs roughly $0.003/sample (~$30 per 10K captions). Cheaper models (Haiku) are
available; pass `model=...` to swap.
"""

from __future__ import annotations
import json
import os
import time
import warnings
from pathlib import Path
from typing import Optional

from ..registry import SLOT_REGISTRY
from ..schema import CAPTION_JSON_SCHEMA
from . import ProviderResult

# TaskSpec import is lazy — inside the function — to avoid a hard dependency
# at module-load time for callers that don't use the task-driven path.


# ──────────────────────────────────────────────────────────────────────────────
# .env loading — minimal stdlib parser, no python-dotenv dependency.
#
# Cowork's VM doesn't inherit the host shell's environment, so a `.env` file
# inside the mounted repo folder is the most reliable way to get the API
# key in. Same pattern works for Claude Code and plain CLI use.
# ──────────────────────────────────────────────────────────────────────────────

def _load_dotenv(path: Path) -> int:
    """Parse a .env file and set unset vars in os.environ. Returns count set.

    Existing env vars are not overwritten (host env wins over .env file).
    Supports lines like:  KEY=value  /  KEY="quoted value"  /  # comments
    """
    if not path.is_file():
        return 0
    n_set = 0
    for raw in path.read_text().splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, _, val = line.partition("=")
        key = key.strip()
        # Strip surrounding quotes if present
        val = val.strip()
        if len(val) >= 2 and val[0] == val[-1] and val[0] in ("'", '"'):
            val = val[1:-1]
        # Skip "export KEY=..." prefix if user copied a shell-style file
        if key.startswith("export "):
            key = key[len("export "):].strip()
        # Treat empty existing values as unset. Some sandboxes (e.g. Claude
        # Code / Cowork) inject blank ANTHROPIC_API_KEY into child processes
        # to mask the host's real key; .env must remain authoritative there.
        if key and not os.environ.get(key):
            os.environ[key] = val
            n_set += 1
    return n_set


def _autoload_dotenv() -> None:
    """Search common locations for a .env and load the first match.

    Priority: cwd/.env, then walk up to 3 parent dirs (for cases where the
    agent is invoked from a subdirectory of the repo).
    """
    cwd = Path.cwd().resolve()
    for candidate in [cwd, *list(cwd.parents)[:3]]:
        env_path = candidate / ".env"
        if env_path.is_file():
            _load_dotenv(env_path)
            return


# ──────────────────────────────────────────────────────────────────────────────
# System prompts — the strict/enhance distinction is encoded here.
#
# Both prompts pin the model to the registry-driven schema. The difference is
# what category fields each prompt licenses the model to populate:
#
#   strict     — descriptive only. style/mood → null. For SFT teacher labels
#                where we want grounded-only outputs to filter on.
#   enhance    — all categories. Style/mood may be inferred. For prompt-
#                enhancement training data.
# ──────────────────────────────────────────────────────────────────────────────

# Both prompts are intentionally sized so that (tool def + system) lands
# safely above Sonnet's 1024-token prompt-caching minimum. The original
# 165-token versions kept the cacheable prefix at ~1023 tokens — 1 below
# threshold — which silently disabled caching and cost ~60% more per call.
# The richer examples also tightened grounding (fewer style/mood leaks on
# strict, more consistent inference on enhance).

PROMPT_STRICT = """You are a caption-structuring assistant. Given an image caption,
emit JSON matching the provided schema. Your sole job is to extract structured
information explicitly stated in the input — never embellish, infer, or imagine
details that aren't there.

RULES:
- `subjects`: every entity named in the caption. Each subject has a `name`
  (a noun phrase from the caption) and optional `attributes` — adjectives or
  descriptors the caption explicitly attaches to that subject (color, age,
  expression, material, count, etc.).
- `actions`: verb phrases describing what's happening. Prefer the caption's
  own wording when possible.
- `setting`: use "indoor" or "outdoor" if the caption indicates it (kitchen,
  restaurant → indoor; park, beach → outdoor). Use "unknown" if no cue.
- `style`: ALWAYS null in strict mode. Do not infer style from content.
- `mood`: ALWAYS null in strict mode. Do not infer mood from content.
- Empty lists `[]` and `null` are correct outputs — DO NOT invent content
  just to populate a field.

EXAMPLES:
- "a young girl in a red dress" → subjects: [{name: "girl", attributes: ["young"]},
  {name: "dress", attributes: ["red"]}]; setting: "unknown"
- "a cat sleeping on a sofa" → subjects: [{name: "cat", attributes: []},
  {name: "sofa", attributes: []}], actions: ["sleeping on a sofa"];
  setting: "indoor"
- "the beach at sunset" → subjects: [{name: "beach", attributes: []}];
  setting: "outdoor"

WHAT TO AVOID:
- Adding subjects not named in the caption (no inferred "people" or
  "background figures" unless the caption mentions them).
- Inferring attributes the caption does not state (do not add "fluffy" for
  a cat just because cats are typically fluffy).
- Setting `style` or `mood` to anything other than null.

Call the `emit_caption_schema` tool with your structured output.""".strip()


PROMPT_ENHANCE = """You are a caption-structuring assistant. Given an image caption,
emit JSON matching the provided schema. Extract grounded content faithfully, and
where the schema licenses inference, draw it from the caption's content rather
than imagining unrelated detail.

RULES:
- `subjects`, `actions`, subject `attributes`: ONLY content explicitly named
  in the caption. Use noun phrases from the caption itself; do not invent
  entities or descriptors.
- `setting`: "indoor" or "outdoor" if the caption indicates or strongly
  implies it (kitchen, restaurant → indoor; park, beach → outdoor);
  otherwise "unknown".
- `style`: you MAY infer a visual style (e.g. "photorealistic",
  "watercolor", "cyberpunk illustration", "vintage film", "anime") when
  the caption suggests one. Leave null if there's no signal.
- `mood`: you MAY infer a mood from the caption's content (e.g. "tense",
  "celebratory", "melancholy", "playful", "serene"). Leave null if neutral.

EXAMPLES:
- "an oil painting of a stormy sea" → subjects: [{name: "sea",
  attributes: ["stormy"]}]; setting: "outdoor"; style: "oil painting";
  mood: "tense"
- "a child laughing at a birthday party" → subjects: [{name: "child",
  attributes: []}], actions: ["laughing"]; setting: "indoor";
  style: null; mood: "celebratory"
- "a sketch of a hand holding a pencil" → subjects: [{name: "hand",
  attributes: []}, {name: "pencil", attributes: []}],
  actions: ["holding a pencil"]; setting: "unknown"; style: "sketch";
  mood: null

Call the `emit_caption_schema` tool with your structured output.""".strip()


# ──────────────────────────────────────────────────────────────────────────────
# Pricing table (per million tokens). Update when Anthropic publishes new rates.
# Used only to estimate cost_usd in ProviderResult — not authoritative.
# ──────────────────────────────────────────────────────────────────────────────

_PRICING = {
    # model_id_substring → (input_$/Mtok, output_$/Mtok)
    "claude-opus-4":    (15.0, 75.0),
    "claude-sonnet-4":  ( 3.0, 15.0),
    "claude-haiku-4":   ( 0.80,  4.0),
    # legacy 3.x models, in case someone pins to them
    "claude-3-5-sonnet":(3.0, 15.0),
    "claude-3-5-haiku": (0.80, 4.0),
}


def _estimate_cost(
    model_id: str,
    n_in: int,
    n_out: int,
    n_cache_create: int = 0,
    n_cache_read: int = 0,
) -> float:
    """Cost in USD for a single call. Cache write at 1.25x input, read at 0.10x."""
    rates = next((v for k, v in _PRICING.items() if k in model_id), None)
    if rates is None:
        return 0.0
    in_rate, out_rate = rates
    return (
        (n_in / 1_000_000) * in_rate
        + (n_cache_create / 1_000_000) * (in_rate * 1.25)
        + (n_cache_read / 1_000_000) * (in_rate * 0.10)
        + (n_out / 1_000_000) * out_rate
    )


# ──────────────────────────────────────────────────────────────────────────────
# Schema slimming — Pydantic's model_json_schema() is verbose by default
# (per-field "title", "default", "$defs/$ref" for nested types). Anthropic's
# tool-use enforcer ignores those cosmetic fields, but they cost input tokens
# on every uncached call. Stripping them cuts the tool definition roughly in
# half while keeping the constraints (types, enums, maxLength, maxItems, etc).
# ──────────────────────────────────────────────────────────────────────────────

_STRIP_KEYS = {"title", "default", "description"}


def _slim_schema(schema: dict) -> dict:
    """Return a copy of `schema` with cosmetic keys removed and $defs inlined.

    Drops: title, default, description (the model's tool-use enforcement
    doesn't read them). Inlines local $defs/$ref pairs so nested types like
    SubjectValue appear directly under their parent property. Constraints
    (types, enums, anyOf, maxLength, minLength, maxItems, required) are kept.
    """
    import copy
    schema = copy.deepcopy(schema)
    defs = schema.pop("$defs", {})

    def resolve(node):
        if isinstance(node, dict):
            if "$ref" in node and node["$ref"].startswith("#/$defs/"):
                name = node["$ref"][len("#/$defs/"):]
                return resolve(defs.get(name, {}))
            return {k: resolve(v) for k, v in node.items() if k not in _STRIP_KEYS}
        if isinstance(node, list):
            return [resolve(x) for x in node]
        return node

    return resolve(schema)


# ──────────────────────────────────────────────────────────────────────────────
# Provider
# ──────────────────────────────────────────────────────────────────────────────

class ClaudeProvider:
    """Caption-processing provider backed by Anthropic's Claude API.

    Loads the anthropic SDK lazily so importing the package doesn't fail
    when anthropic isn't installed and the user just wants the Qwen path.
    """

    def __init__(
        self,
        model: str = "claude-sonnet-4-6",
        api_key: Optional[str] = None,
        max_retries: int = 3,
        retry_backoff: float = 2.0,
    ):
        try:
            import anthropic
        except ImportError as e:
            raise ImportError(
                "ClaudeProvider requires the `anthropic` package. "
                "Install with: pip install anthropic"
            ) from e

        # Cowork/Claude Code don't inherit the host shell environment — load
        # a project-local .env if one exists. Has no effect when the key is
        # already in os.environ.
        _autoload_dotenv()

        api_key = api_key or os.environ.get("ANTHROPIC_API_KEY")
        if not api_key:
            raise RuntimeError(
                "No Anthropic API key. Set ANTHROPIC_API_KEY in your shell, "
                "drop a `.env` file with ANTHROPIC_API_KEY=... in the repo "
                "root, or pass api_key= to ClaudeProvider(...)."
            )

        self.client = anthropic.Anthropic(api_key=api_key)
        self.model = model
        self.max_retries = max_retries
        self.retry_backoff = retry_backoff

        # The tool definition is the JSON Schema generated by the registry.
        # Forcing tool use guarantees the output is schema-valid.
        #
        # NOTE: we keep the *verbose* pydantic schema rather than passing
        # _slim_schema(CAPTION_JSON_SCHEMA) here. The slim form saves ~110
        # tokens per call, but it also pushes the (tools + system) cacheable
        # prefix below Sonnet's 1024-token minimum, which silently disables
        # prompt caching — losing ~60% in subsequent-call savings. The
        # _slim_schema helper stays available for cases where caching is
        # off (e.g. one-shot calls) or for models with a smaller minimum.
        self._tool = {
            "name": "emit_caption_schema",
            "description": (
                "Emit the structured caption representation. The input_schema "
                "follows the qwen-test-runner slot registry."
            ),
            "input_schema": CAPTION_JSON_SCHEMA,
        }

    def process(
        self,
        caption: str,
        prompt: str = "strict",
        max_tokens: int = 1024,
        task=None,  # Optional[TaskSpec] — takes precedence over prompt= when set
    ) -> ProviderResult:
        """Convert one caption to schema-conformant JSON.

        Two modes (use one):

          task=<TaskSpec>     — task-driven: system prompt + tool schema come
                                from the TaskSpec. mode_tag = "claude_<task.name>".
                                Preferred for the per-task SFT pipeline.

          prompt="strict"|"enhance"|"<custom>" — legacy path. Uses the module-level
                                PROMPT_STRICT/PROMPT_ENHANCE constants and the
                                universal CAPTION_JSON_SCHEMA as the tool's
                                input_schema. Kept for back-compat with data_gen.py.
        """
        if task is not None:
            # Task-driven path. We build a tool definition locally rather than
            # reuse self._tool so the per-task schema overlay takes effect.
            tool = {
                "name": "emit_caption_schema",
                "description": self._tool["description"],
                "input_schema": task.tool_schema,
            }
            sys_prompt = task.system_prompt
            mode_tag = f"claude_{task.name}"
        else:
            tool = self._tool
            if prompt == "strict":
                sys_prompt = PROMPT_STRICT
                mode_tag = "claude_strict"
            elif prompt == "enhance":
                sys_prompt = PROMPT_ENHANCE
                mode_tag = "claude_enhance"
            else:
                sys_prompt = prompt
                mode_tag = "claude_custom"

        response = self._call_with_retry(
            system=sys_prompt,
            user=caption,
            max_tokens=max_tokens,
            tool=tool,
        )

        # Find the tool_use block. Forced tool_choice means there's always
        # exactly one — but we extract by type, not position, for safety.
        tool_input = None
        for block in response.content:
            if block.type == "tool_use" and block.name == "emit_caption_schema":
                tool_input = block.input
                break
        if tool_input is None:
            raise RuntimeError(
                f"Claude returned no tool_use block. Stop reason: {response.stop_reason!r}"
            )

        raw_json = json.dumps(tool_input, separators=(",", ":"))

        usage = response.usage
        n_in = usage.input_tokens
        n_out = usage.output_tokens
        n_cache_create = getattr(usage, "cache_creation_input_tokens", 0) or 0
        n_cache_read = getattr(usage, "cache_read_input_tokens", 0) or 0
        cost = _estimate_cost(self.model, n_in, n_out, n_cache_create, n_cache_read)

        return ProviderResult(
            mode=mode_tag,
            raw_text=raw_json,
            backend="claude",
            n_input_tokens=n_in,
            n_output_tokens=n_out,
            cost_usd=cost,
            n_cache_creation_tokens=n_cache_create,
            n_cache_read_tokens=n_cache_read,
        )

    def _call_with_retry(self, system: str, user: str, max_tokens: int, tool=None):
        """Anthropic call with exponential backoff on rate-limit / transient errors.

        tool: optional override for the tool definition. Defaults to self._tool
              (the universal schema). Task-driven callers pass their own.
        """
        import anthropic  # already imported in __init__, just re-bind name

        tool = tool if tool is not None else self._tool
        last_err: Optional[Exception] = None
        for attempt in range(self.max_retries):
            try:
                # `cache_control` on the last system block marks the cache
                # breakpoint. Everything before/including it (tools + system)
                # is cached; the user message remains variable per-call.
                # Sonnet's minimum cacheable prefix is 1024 tokens — our
                # tool def + system prompt sits just above that.
                return self.client.messages.create(
                    model=self.model,
                    max_tokens=max_tokens,
                    system=[{
                        "type": "text",
                        "text": system,
                        "cache_control": {"type": "ephemeral"},
                    }],
                    tools=[tool],
                    tool_choice={"type": "tool", "name": "emit_caption_schema"},
                    messages=[{"role": "user", "content": f"Caption: {user}"}],
                )
            except (anthropic.RateLimitError, anthropic.APIStatusError) as e:
                last_err = e
                sleep_s = self.retry_backoff ** attempt
                warnings.warn(
                    f"Claude API error (attempt {attempt + 1}/{self.max_retries}): "
                    f"{type(e).__name__}: {e}. Sleeping {sleep_s:.1f}s."
                )
                time.sleep(sleep_s)
        # All retries exhausted
        raise RuntimeError(f"Claude API failed after {self.max_retries} retries") from last_err