""" 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= — task-driven: system prompt + tool schema come from the TaskSpec. mode_tag = "claude_". Preferred for the per-task SFT pipeline. prompt="strict"|"enhance"|"" — 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