qwen_test_runner/tasks.py

"""
tasks.py — Task registry.

Each TaskSpec declares a single "what kind of JSON should the model emit"
behaviour and owns everything needed to drive Claude (the teacher) and Qwen
(the student) toward it: a system prompt, a tool-schema overlay on top of
the universal CAPTION_JSON_SCHEMA, and a validator hook for post-schema
checks (grounding for task_1, regex pattern for task_2 and task_3).

Three tasks (as of v0.2):

    task_1 — hallucination_reduction
        Grounded literal extraction. Subject/action/attribute values come
        from the caption verbatim. Style and mood are forbidden (null).
        The schema does not enable inference; the validator runs grounding
        check (substring + token match against input caption).

    task_2 — useful_generalization
        Encouraged categorical abstraction. Every string value is a
        bracketed canonical generic like [pet], [vehicle], [color], [playing].
        Schema constrains values to regex /^\\[[a-z_]+\\]$/.
        Validator just enforces the format; semantic correctness is
        a soft target — the open vocabulary is curated post-hoc from
        what the model actually emits.

    task_3 — generic_symbolism
        Pure positional placeholders. subjects[].name → [ENTITY_N],
        actions[] → [ACTION_N], setting → [INDOOR|OUTDOOR|UNKNOWN],
        attributes → [ATTRIBUTE_N]. Numbering is within-slot, starts at 1,
        monotonically increasing. Style and mood are nullable typed
        placeholders.

Adding a task is one TASK_REGISTRY entry. The pipeline (prompt_maker.py)
iterates TASK_REGISTRY; downstream consumers (ClaudeProvider, the qwen
tester) look tasks up by name.
"""
from __future__ import annotations

import copy
import re
from dataclasses import dataclass, field
from typing import Callable, Optional

from .schema import CAPTION_JSON_SCHEMA


# ──────────────────────────────────────────────────────────────────────────────
# TaskSpec
# ──────────────────────────────────────────────────────────────────────────────

@dataclass(frozen=True)
class TaskSpec:
    """Declarative definition of one differentiation mode.

    Fields:
      name             — stable task id used in row tags + filenames
      description      — one-liner for logs and row meta
      system_prompt    — the task's system prompt (Claude + Qwen)
      tool_schema      — a JSON Schema dict, fully built (with overlays applied).
                          Passed as input_schema to Claude's tool def.
      value_pattern    — optional regex every emitted string value must match.
                          Used by both Claude (via schema 'pattern') AND the
                          evaluator (post-hoc check on Qwen outputs).
      validate         — optional post-hoc validator. Signature:
                            (caption, parsed_args_dict) -> list[str]
                          Returns a list of reject reasons; empty list = pass.
    """
    name: str
    description: str
    system_prompt: str
    tool_schema: dict
    value_pattern: Optional[str] = None
    validate: Optional[Callable] = None


# ──────────────────────────────────────────────────────────────────────────────
# Schema-overlay helpers (used to build per-task tool_schema from base)
# ──────────────────────────────────────────────────────────────────────────────

def _deep_merge(base: dict, overlay: dict) -> dict:
    """Recursively merge overlay into a copy of base. Overlay wins on conflicts."""
    out = copy.deepcopy(base)
    for k, v in overlay.items():
        if isinstance(v, dict) and isinstance(out.get(k), dict):
            out[k] = _deep_merge(out[k], v)
        else:
            out[k] = copy.deepcopy(v)
    return out


def _apply_string_pattern(schema: dict, pattern: str) -> dict:
    """Return a copy of schema with `pattern` applied to every string-typed leaf.

    Walks the schema and adds {'pattern': pattern} to every node where
    type=='string' (including inside anyOf branches). Skips closed enums
    — those are already constrained.
    """
    out = copy.deepcopy(schema)

    def walk(node):
        if isinstance(node, dict):
            # If this node is a string type without an enum, attach pattern
            if node.get("type") == "string" and "enum" not in node:
                node["pattern"] = pattern
            # Recurse into children
            for v in node.values():
                walk(v)
        elif isinstance(node, list):
            for item in node:
                walk(item)

    walk(out)
    return out


# ──────────────────────────────────────────────────────────────────────────────
# Task 1: hallucination_reduction
#
# Schema overlay forces style and mood to const null so Claude cannot emit
# anything else. The system prompt also forbids them — belt and suspenders.
# Grounding check is the validator (uses evaluator.ground_check).
# ──────────────────────────────────────────────────────────────────────────────

_TASK1_OVERLAY = {
    "properties": {
        "style": {"const": None},
        "mood":  {"const": None},
    },
}

_TASK1_PROMPT = """You are a caption-structuring assistant. Given an image-synthesis
prompt, emit structured JSON via the emit_caption_schema tool. Your job is
GROUNDED LITERAL EXTRACTION — extract structured information that is
explicitly stated in the input, never embellish, infer, or imagine details.

RULES:
- subjects: every entity named in the caption. Each subject has a name (a noun
  phrase taken from the caption) and optional attributes (adjectives/descriptors
  the caption explicitly attaches to that subject: color, age, expression,
  material, count, etc.).
- actions: verb phrases describing what is happening. Use caption wording.
- setting: "indoor" or "outdoor" if the caption indicates it (kitchen,
  restaurant → indoor; park, beach → outdoor). Otherwise "unknown".
- style: ALWAYS null. The schema does not permit any other value here.
- mood: ALWAYS null. The schema does not permit any other value here.
- Empty lists [] and null are correct outputs — DO NOT invent content to fill
  any field. Schema-valid empty is better than schema-valid invented.

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:
- Inventing subjects, attributes, or actions not in the caption.
- Inferring style or mood — the schema rejects anything but null for these.

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


# ──────────────────────────────────────────────────────────────────────────────
# Task 2: useful_generalization
#
# All open-vocab string values must match /^\[[a-z_]+\]$/ — bracketed lowercase
# generics like [pet], [vehicle], [playing], [outdoor_scene].
# setting's enum is replaced with bracketed versions for consistency.
# Style and mood remain null (style/mood are out of scope for this task too).
# ──────────────────────────────────────────────────────────────────────────────

_TASK2_PATTERN = r"^\[[a-z_]+\]$"

_TASK2_SETTING_ENUM = ["[indoor]", "[outdoor]", "[unknown]"]

# Build task_2's schema: apply pattern to all open strings, then overlay
# setting's enum + force style/mood null.
def _build_task2_schema() -> dict:
    s = _apply_string_pattern(CAPTION_JSON_SCHEMA, _TASK2_PATTERN)
    overlay = {
        "properties": {
            "setting": {
                "enum": _TASK2_SETTING_ENUM,
                "default": "[unknown]",
            },
            "style": {"const": None},
            "mood":  {"const": None},
        },
    }
    s = _deep_merge(s, overlay)
    # The 'setting' enum was overwritten; remove its old pattern (closed vocab
    # doesn't need it, and pattern + enum can confuse some validators).
    if "pattern" in s["properties"]["setting"]:
        del s["properties"]["setting"]["pattern"]
    return s


_TASK2_PROMPT = """You are a caption-structuring assistant. Given an image-synthesis
prompt, emit structured JSON via the emit_caption_schema tool. Your job is
USEFUL GENERALIZATION — abstract every concrete noun, adjective, and verb to
a small canonical CATEGORICAL GENERIC in [bracket_word] form.

RULES:
- Every open-vocabulary string value MUST be in [lowercase_with_underscores]
  format, between square brackets. The schema enforces this.
- subjects: list of bracketed generics that abstract caption entities.
  Prefer the smallest sensible category ([pet] over [golden_retriever];
  [clothing] over [red_dress]; [tool] over [pencil]).
- attributes: bracketed generic descriptors ([color], [young], [shiny]).
- actions: bracketed generic verbs ([playing], [eating], [waiting]).
- setting: choose [indoor], [outdoor], or [unknown].
- style: ALWAYS null in this task.
- mood: ALWAYS null in this task.

EXAMPLES:
- "a golden retriever catching a red frisbee in a sunny park" →
  subjects: [{name: "[pet]", attributes: []},
             {name: "[toy]", attributes: ["[color]"]}]
  actions: ["[playing]"]
  setting: "[outdoor]"

- "a young girl in a red dress" →
  subjects: [{name: "[person]", attributes: ["[young]"]},
             {name: "[clothing]", attributes: ["[color]"]}]
  actions: []
  setting: "[unknown]"

- "an architect at his desk reviewing blueprints" →
  subjects: [{name: "[person]", attributes: []},
             {name: "[furniture]", attributes: []},
             {name: "[document]", attributes: []}]
  actions: ["[working]"]
  setting: "[indoor]"

The aim is to teach a categorical view of caption content. Pick generics that
group similar specifics together. Different captions producing similar generic
structures is GOOD — that is the point.

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


# ──────────────────────────────────────────────────────────────────────────────
# Task 3: generic_symbolism
#
# Numbered typed placeholders. Each slot has its own type prefix and integer
# index (1-based, monotonic within slot). Captures positional structure with
# zero semantic content.
# ──────────────────────────────────────────────────────────────────────────────

_TASK3_ENTITY_PATTERN     = r"^\[ENTITY_\d+\]$"
_TASK3_ATTRIBUTE_PATTERN  = r"^\[ATTRIBUTE_\d+\]$"
_TASK3_ACTION_PATTERN     = r"^\[ACTION_\d+\]$"

_TASK3_SETTING_ENUM = ["[INDOOR]", "[OUTDOOR]", "[UNKNOWN]"]


def _build_task3_schema() -> dict:
    s = copy.deepcopy(CAPTION_JSON_SCHEMA)
    # subjects[].name → ENTITY pattern; subjects[].attributes[] → ATTRIBUTE pattern
    subj = s["$defs"]["SubjectValue"]["properties"]
    subj["name"]["pattern"] = _TASK3_ENTITY_PATTERN
    subj["attributes"]["items"]["pattern"] = _TASK3_ATTRIBUTE_PATTERN
    # actions[] → ACTION pattern
    s["properties"]["actions"]["items"]["pattern"] = _TASK3_ACTION_PATTERN
    # setting → bracketed UPPERCASE enum
    s["properties"]["setting"]["enum"] = _TASK3_SETTING_ENUM
    s["properties"]["setting"]["default"] = "[UNKNOWN]"
    # style and mood: must be null in this task too (placeholder structure
    # doesn't have a meaningful "style" position — keep nullable for symmetry).
    s["properties"]["style"] = {"const": None}
    s["properties"]["mood"]  = {"const": None}
    return s


_TASK3_PROMPT = """You are a caption-structuring assistant. Given an image-synthesis
prompt, emit structured JSON via the emit_caption_schema tool. Your job is
PURE STRUCTURAL SYMBOLISM — convert every entity to a numbered typed
placeholder. The output captures positional roles only, with zero semantic
content.

FORMAT:
- subjects[i].name      → [ENTITY_N]      (N = 1, 2, 3, ... in caption order)
- subjects[i].attributes[j] → [ATTRIBUTE_N] (N restarts at 1 within each subject)
- actions[i]            → [ACTION_N]      (N = 1, 2, 3, ... in caption order)
- setting               → [INDOOR], [OUTDOOR], or [UNKNOWN] (uppercase)
- style                 → null
- mood                  → null

NUMBERING RULES:
- N is a positive integer starting at 1.
- Within a slot, numbering is monotonically increasing with no gaps.
- Each occurrence of a real entity → one ENTITY_N; do not collapse duplicates.

EXAMPLES:

- "a golden retriever catching a red frisbee in a sunny park" →
  subjects: [{name: "[ENTITY_1]", attributes: [],
              "..."},
             {name: "[ENTITY_2]", attributes: ["[ATTRIBUTE_1]"]}]
  actions: ["[ACTION_1]"]
  setting: "[OUTDOOR]"
  (ENTITY_1=retriever, ATTRIBUTE_1 on ENTITY_1=golden was DROPPED because
  the caption attached "golden" to the retriever; we keep that as attributes.
  Wait — corrected: golden retriever has attribute "golden" → ATTRIBUTE_1.
  frisbee has attribute "red" → ATTRIBUTE_1 (restart per subject).)

- "two children playing chess" →
  subjects: [{name: "[ENTITY_1]", attributes: ["[ATTRIBUTE_1]"]},
             {name: "[ENTITY_2]", attributes: []}]
  actions: ["[ACTION_1]"]
  setting: "[UNKNOWN]"
  (ENTITY_1=children, ATTRIBUTE_1=two on children; ENTITY_2=chess)

- "the beach at sunset" →
  subjects: [{name: "[ENTITY_1]", attributes: []}]
  actions: []
  setting: "[OUTDOOR]"

The aim is to teach the model to think about caption STRUCTURE divorced from
content. Two completely different captions with the same shape should produce
the same JSON.

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


# ──────────────────────────────────────────────────────────────────────────────
# Validators
# ──────────────────────────────────────────────────────────────────────────────

def _validate_task1(caption: str, args: dict) -> list[str]:
    """Grounding check. Imported lazily to avoid circular import with evaluator."""
    from .evaluator import parse_safely, ground_check
    import json
    parse = parse_safely(json.dumps(args))
    if not parse.schema_valid or parse.parsed is None:
        return [f"schema: {parse.error}"]
    report = ground_check(parse.parsed, caption)
    if report.grounding_rate < 1.0:
        return [f"hallucinated: {h[1]!r} at {h[0]}" for h in report.hallucinated]
    return []


_TASK2_VALUE_RE = re.compile(_TASK2_PATTERN)
_TASK3_ENTITY_RE = re.compile(_TASK3_ENTITY_PATTERN)
_TASK3_ATTRIBUTE_RE = re.compile(_TASK3_ATTRIBUTE_PATTERN)
_TASK3_ACTION_RE = re.compile(_TASK3_ACTION_PATTERN)


def _safe_match(regex: re.Pattern, value) -> bool:
    """Match-or-False without crashing on non-string inputs.

    Claude occasionally emits dicts where strings are expected (e.g.
    actions=[{'type':'action','text':'...'}]). The schema's tool_use
    enforcement *usually* catches this, but failures slip through often
    enough that the validator must not crash on them.
    """
    return isinstance(value, str) and regex.fullmatch(value) is not None


def _validate_task2(caption: str, args: dict) -> list[str]:
    """Every open-vocab string must match the bracketed-generic pattern."""
    errs: list[str] = []
    if not isinstance(args, dict):
        return [f"args is not a dict: {type(args).__name__}"]
    for i, subj in enumerate(args.get("subjects") or []):
        if not isinstance(subj, dict):
            errs.append(f"subjects[{i}] is not a dict: {type(subj).__name__}")
            continue
        if not _safe_match(_TASK2_VALUE_RE, subj.get("name")):
            errs.append(f"subjects[{i}].name not bracketed: {subj.get('name')!r}")
        for j, attr in enumerate(subj.get("attributes") or []):
            if not _safe_match(_TASK2_VALUE_RE, attr):
                errs.append(f"subjects[{i}].attributes[{j}] not bracketed: {attr!r}")
    for i, a in enumerate(args.get("actions") or []):
        if not _safe_match(_TASK2_VALUE_RE, a):
            errs.append(f"actions[{i}] not bracketed: {a!r}")
    setting = args.get("setting")
    if setting is not None and setting not in _TASK2_SETTING_ENUM:
        errs.append(f"setting not in enum: {setting!r}")
    return errs


def _validate_task3(caption: str, args: dict) -> list[str]:
    """Typed numbered placeholders + monotonic numbering within slot."""
    errs: list[str] = []
    if not isinstance(args, dict):
        return [f"args is not a dict: {type(args).__name__}"]
    # subjects.name → ENTITY_N, monotonic
    for i, subj in enumerate(args.get("subjects") or []):
        if not isinstance(subj, dict):
            errs.append(f"subjects[{i}] is not a dict: {type(subj).__name__}")
            continue
        name = subj.get("name")
        if not _safe_match(_TASK3_ENTITY_RE, name):
            errs.append(f"subjects[{i}].name not [ENTITY_N]: {name!r}")
            continue
        if name != f"[ENTITY_{i + 1}]":
            errs.append(f"subjects[{i}].name should be [ENTITY_{i + 1}], got {name!r}")
        for j, attr in enumerate(subj.get("attributes") or []):
            if not _safe_match(_TASK3_ATTRIBUTE_RE, attr):
                errs.append(f"subjects[{i}].attributes[{j}] not [ATTRIBUTE_N]: {attr!r}")
                continue
            if attr != f"[ATTRIBUTE_{j + 1}]":
                errs.append(
                    f"subjects[{i}].attributes[{j}] should be [ATTRIBUTE_{j + 1}], got {attr!r}"
                )
    # actions: ACTION_N, monotonic
    for i, a in enumerate(args.get("actions") or []):
        if not _safe_match(_TASK3_ACTION_RE, a):
            errs.append(f"actions[{i}] not [ACTION_N]: {a!r}")
            continue
        if a != f"[ACTION_{i + 1}]":
            errs.append(f"actions[{i}] should be [ACTION_{i + 1}], got {a!r}")
    setting = args.get("setting")
    if setting is not None and setting not in _TASK3_SETTING_ENUM:
        errs.append(f"setting not in enum: {setting!r}")
    return errs


# ──────────────────────────────────────────────────────────────────────────────
# THE REGISTRY
# ──────────────────────────────────────────────────────────────────────────────

TASK_REGISTRY: dict[str, TaskSpec] = {
    "task_1": TaskSpec(
        name="task_1",
        description="hallucination_reduction: grounded literal extraction; null style/mood",
        system_prompt=_TASK1_PROMPT,
        tool_schema=_deep_merge(CAPTION_JSON_SCHEMA, _TASK1_OVERLAY),
        value_pattern=None,
        validate=_validate_task1,
    ),
    "task_2": TaskSpec(
        name="task_2",
        description="useful_generalization: bracketed categorical generics",
        system_prompt=_TASK2_PROMPT,
        tool_schema=_build_task2_schema(),
        value_pattern=_TASK2_PATTERN,
        validate=_validate_task2,
    ),
    "task_3": TaskSpec(
        name="task_3",
        description="generic_symbolism: numbered typed placeholders",
        system_prompt=_TASK3_PROMPT,
        tool_schema=_build_task3_schema(),
        value_pattern=None,  # multiple patterns per slot, handled in validator
        validate=_validate_task3,
    ),
}


def get_task(name: str) -> TaskSpec:
    if name not in TASK_REGISTRY:
        raise KeyError(f"unknown task: {name!r}. known: {list(TASK_REGISTRY)}")
    return TASK_REGISTRY[name]


def task_names() -> list[str]:
    return list(TASK_REGISTRY.keys())