Initial commit

skillopt/envs/searchqa/__init__.py

@@ -0,0 +1 @@
+"""SearchQA environment package for ReflACT."""

skillopt/envs/searchqa/adapter.py

@@ -0,0 +1,250 @@
+"""SearchQA environment adapter for ReflACT."""
+from __future__ import annotations
+
+import json
+import os
+
+from skillopt.gradient.deep_probe import generate_deep_probe_instruction
+from skillopt.datasets.base import BatchSpec
+from skillopt.envs.base import EnvAdapter
+from skillopt.envs.searchqa.dataloader import SearchQADataLoader
+from skillopt.envs.searchqa.rollout import run_batch
+from skillopt.gradient.reflect import run_minibatch_reflect
+from skillopt.model import get_student_backend
+
+
+class SearchQAAdapter(EnvAdapter):
+    """SearchQA environment adapter."""
+
+    def __init__(
+        self,
+        split_dir: str = "",
+        data_path: str = "",
+        split_mode: str = "ratio",
+        split_ratio: str = "2:1:7",
+        split_seed: int = 42,
+        split_output_dir: str = "",
+        max_turns: int = 1,
+        exec_timeout: int = 120,
+        workers: int = 64,
+        analyst_workers: int = 16,
+        failure_only: bool = False,
+        minibatch_size: int = 8,
+        edit_budget: int = 4,
+        seed: int = 42,
+        limit: int = 0,
+        exec_timeout: int = 600,
+        use_deep_reflect: bool = False,
+        deep_reflect_failures: int = 4,
+        deep_reflect_successes: int = 2,
+    ) -> None:
+        self.max_turns = max_turns
+        self.exec_timeout = exec_timeout
+        self.workers = workers
+        self.analyst_workers = analyst_workers
+        self.failure_only = failure_only
+        self.minibatch_size = minibatch_size
+        self.edit_budget = edit_budget
+        self.exec_timeout = exec_timeout
+        self.use_deep_reflect = use_deep_reflect
+        self.deep_reflect_failures = deep_reflect_failures
+        self.deep_reflect_successes = deep_reflect_successes
+        self.dataloader = SearchQADataLoader(
+            split_dir=split_dir,
+            data_path=data_path,
+            split_mode=split_mode,
+            split_ratio=split_ratio,
+            split_seed=split_seed,
+            split_output_dir=split_output_dir,
+            seed=seed,
+            limit=limit,
+        )
+
+    def setup(self, cfg: dict) -> None:
+        super().setup(cfg)
+        self.dataloader.setup(cfg)
+
+    def get_dataloader(self):
+        return self.dataloader
+
+    def build_env_from_batch(self, batch: BatchSpec, **kwargs):
+        return list(batch.payload or [])
+
+    def build_train_env(self, batch_size: int, seed: int, **kwargs):
+        batch = self.dataloader.build_train_batch(batch_size=batch_size, seed=seed, **kwargs)
+        return self.build_env_from_batch(batch, **kwargs)
+
+    def build_eval_env(self, env_num: int, split: str, seed: int, **kwargs):
+        batch = self.dataloader.build_eval_batch(env_num=env_num, split=split, seed=seed, **kwargs)
+        return self.build_env_from_batch(batch, **kwargs)
+
+    def rollout(
+        self,
+        env_manager,  # actually list[dict] for SearchQA
+        skill_content: str,
+        out_dir: str,
+        **kwargs,
+    ) -> list[dict]:
+        """Run QA agent on items. Resume-aware."""
+        items: list[dict] = env_manager  # type alias for clarity
+        return run_batch(
+            items=items,
+            out_root=out_dir,
+            skill_content=skill_content,
+            max_turns=self.max_turns,
+            exec_timeout=self.exec_timeout,
+            workers=self.workers,
+            diagnostic_mode=kwargs.get("diagnostic_mode", False),
+            diagnostic_instruction=kwargs.get("diagnostic_instruction", ""),
+            diagnostic_trace_context_by_id=kwargs.get("diagnostic_trace_context_by_id"),
+            task_timeout=self.exec_timeout,
+        )
+
+    def reflect(
+        self,
+        results: list[dict],
+        skill_content: str,
+        out_dir: str,
+        **kwargs,
+    ) -> list[dict | None]:
+        prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
+        patches_dir = kwargs.get("patches_dir", os.path.join(out_dir, "patches"))
+        random_seed = kwargs.get("random_seed")
+        step_buffer_context = kwargs.get("step_buffer_context", "")
+        meta_skill_context = kwargs.get("meta_skill_context", "")
+
+        return run_minibatch_reflect(
+            results=results,
+            skill_content=skill_content,
+            prediction_dir=prediction_dir,
+            patches_dir=patches_dir,
+            workers=self.analyst_workers,
+            failure_only=self.failure_only,
+            minibatch_size=self.minibatch_size,
+            edit_budget=self.edit_budget,
+            random_seed=random_seed,
+            error_system=self.get_error_minibatch_prompt(),
+            success_system=self.get_success_minibatch_prompt(),
+            step_buffer_context=step_buffer_context,
+            meta_skill_context=meta_skill_context,
+            update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
+        )
+
+    def deep_reflect(
+        self,
+        results: list[dict],
+        skill_content: str,
+        out_dir: str,
+        **kwargs,
+    ) -> list[dict | None]:
+        if not self.use_deep_reflect:
+            return []
+
+        env_manager = kwargs.get("env_manager")
+        if not isinstance(env_manager, list):
+            return []
+
+        prediction_dir = kwargs.get("prediction_dir", os.path.join(out_dir, "predictions"))
+        random_seed = kwargs.get("random_seed")
+        step_buffer_context = kwargs.get("step_buffer_context", "")
+        meta_skill_context = kwargs.get("meta_skill_context", "")
+        codex_backend = get_student_backend() == "codex_exec"
+        selected_items = self.select_representative_items(
+            results,
+            env_manager,
+            n_failures=self.deep_reflect_failures,
+            n_successes=self.deep_reflect_successes,
+            seed=random_seed,
+        )
+        if not selected_items:
+            return []
+
+        selected_ids = {str(item["id"]) for item in selected_items}
+        selected_results = [row for row in results if str(row.get("id")) in selected_ids]
+        selected_examples = (
+            self.attach_codex_probe_context(selected_results, prediction_dir)
+            if codex_backend
+            else selected_results
+        )
+        selected_metadata = [
+            {
+                "id": str(item["id"]),
+                "question_preview": str(item.get("question") or "")[:200],
+                "has_context": bool(str(item.get("context") or "").strip()),
+                "n_gold_answers": len(item.get("answers") or []),
+            }
+            for item in selected_items
+        ]
+
+        deep_dir = os.path.join(out_dir, "deep_reflect")
+        rollout_dir = os.path.join(deep_dir, "rollout")
+        patches_dir = os.path.join(deep_dir, "patches")
+        os.makedirs(deep_dir, exist_ok=True)
+        print(
+            f"    [2b/6 DEEP REFLECT setup] selected={len(selected_items)} "
+            f"mode=no_reference_probe"
+        )
+        probe = generate_deep_probe_instruction(
+            skill_content=skill_content,
+            items=selected_examples,
+            prediction_dir=prediction_dir,
+            system_prompt=self.get_codex_deep_probe_prompt() if codex_backend else self.get_deep_probe_prompt(),
+            step_buffer_context=step_buffer_context,
+            meta_skill_context=meta_skill_context,
+            output_requirements=[
+                "- There is no hidden reference block. Use only the question, provided context, the student's output, and the evaluation result to infer what intermediate state is worth probing.",
+                "- The instruction must explicitly request a short <analysis>...</analysis> block before the final <answer>...</answer>.",
+                "- The readout should focus on likely evidence span, top candidate and runner-up, decisive clue, or a few short intermediate conclusions.",
+                "- Do not ask for exhaustive copying of the context or a full chain-of-thought.",
+                "- The instruction text should be ready to append directly to the student's prompt.",
+            ],
+        )
+        if not probe:
+            return []
+        diagnostic_trace_context_by_id = None
+        if codex_backend:
+            selected_items, diagnostic_trace_context_by_id, probe = self.resolve_codex_probe_target(
+                selected_items=selected_items,
+                selected_examples=selected_examples,
+                prediction_dir=prediction_dir,
+                probe=probe,
+            )
+
+        with open(os.path.join(deep_dir, "probe.json"), "w", encoding="utf-8") as f:
+            json.dump(
+                {
+                    **probe,
+                    "selected_examples": selected_metadata,
+                },
+                f,
+                ensure_ascii=False,
+                indent=2,
+            )
+
+        deep_results = self.rollout(
+            selected_items,
+            skill_content,
+            rollout_dir,
+            diagnostic_mode=True,
+            diagnostic_instruction=probe["probe_instruction"],
+            diagnostic_trace_context_by_id=diagnostic_trace_context_by_id,
+        )
+        return run_minibatch_reflect(
+            results=deep_results,
+            skill_content=skill_content,
+            prediction_dir=os.path.join(rollout_dir, "predictions"),
+            patches_dir=patches_dir,
+            workers=self.analyst_workers,
+            failure_only=self.failure_only,
+            minibatch_size=self.minibatch_size,
+            edit_budget=self.edit_budget,
+            random_seed=random_seed,
+            error_system=self.get_error_minibatch_prompt(),
+            success_system=self.get_success_minibatch_prompt(),
+            step_buffer_context=step_buffer_context,
+            meta_skill_context=meta_skill_context,
+            update_mode=getattr(self, "_cfg", {}).get("skill_update_mode", "patch"),
+        )
+
+    def get_task_types(self) -> list[str]:
+        return ["qa"]

skillopt/envs/searchqa/dataloader.py

@@ -0,0 +1,42 @@
+"""SearchQA task dataloader."""
+from __future__ import annotations
+
+import json
+
+from skillopt.datasets.base import SplitDataLoader
+
+
+# ── Raw data loading utilities (for preprocessing / standalone eval) ─────
+
+def _load_items(path: str) -> list[dict]:
+    """Load items from JSON or JSONL file."""
+    with open(path) as f:
+        content = f.read().strip()
+    try:
+        data = json.loads(content)
+        if isinstance(data, list):
+            return data
+        if isinstance(data, dict):
+            return data.get("data") or list(data.values())
+    except json.JSONDecodeError:
+        pass
+
+    items = []
+    for line in content.splitlines():
+        line = line.strip()
+        if line:
+            items.append(json.loads(line))
+    return items
+
+
+# ── Dataloader ───────────────────────────────────────────────────────────
+
+class SearchQADataLoader(SplitDataLoader):
+    """SearchQA dataloader.
+
+    Each split directory (train/, val/, test/) contains a .json file —
+    a JSON array of question items.
+    """
+
+    def load_raw_items(self, data_path: str) -> list[dict]:
+        return _load_items(data_path)

skillopt/envs/searchqa/evaluator.py

@@ -0,0 +1,100 @@
+"""SearchQA evaluation metrics: Exact Match, F1, and Substring Match.
+
+Normalization follows the SQuAD convention:
+  - lowercase
+  - remove punctuation
+  - remove articles (a, an, the)
+  - collapse whitespace
+
+Answer extraction looks for <answer>...</answer> XML tags,
+falling back to the last non-empty line of the response.
+"""
+from __future__ import annotations
+
+import re
+import string
+from collections import Counter
+
+
+def normalize_answer(s: str) -> str:
+    """Normalize answer string (SQuAD convention)."""
+    s = s.lower()
+    s = "".join(ch for ch in s if ch not in string.punctuation)
+    s = re.sub(r"\b(a|an|the)\b", " ", s)
+    s = " ".join(s.split())
+    return s.strip()
+
+
+def extract_answer(text: str) -> str:
+    """Extract answer from <answer>...</answer> tags.
+
+    Fallback: last non-empty line, then full response stripped.
+    """
+    matches = re.findall(r"<answer>(.*?)</answer>", text, re.DOTALL | re.IGNORECASE)
+    if matches:
+        return matches[-1].strip()
+    lines = [ln.strip() for ln in text.strip().splitlines() if ln.strip()]
+    if lines:
+        return lines[-1]
+    return text.strip()
+
+
+def exact_match(prediction: str, gold_answers: list[str]) -> float:
+    norm_pred = normalize_answer(prediction)
+    for gold in gold_answers:
+        if normalize_answer(gold) == norm_pred:
+            return 1.0
+    return 0.0
+
+
+def f1_score(prediction: str, gold_answers: list[str]) -> float:
+    """Token-level F1 (SQuAD-style), max across all gold answers."""
+    norm_pred = normalize_answer(prediction)
+    pred_tokens = norm_pred.split()
+
+    if not pred_tokens:
+        for gold in gold_answers:
+            if not normalize_answer(gold).split():
+                return 1.0
+        return 0.0
+
+    best_f1 = 0.0
+    for gold in gold_answers:
+        gold_tokens = normalize_answer(gold).split()
+        if not gold_tokens:
+            continue
+        common = Counter(pred_tokens) & Counter(gold_tokens)
+        n_common = sum(common.values())
+        if n_common == 0:
+            continue
+        precision = n_common / len(pred_tokens)
+        recall = n_common / len(gold_tokens)
+        f1 = 2 * precision * recall / (precision + recall)
+        best_f1 = max(best_f1, f1)
+
+    return best_f1
+
+
+def sub_em(prediction: str, gold_answers: list[str]) -> float:
+    """1.0 if any normalized gold is a substring of prediction, or vice versa."""
+    norm_pred = normalize_answer(prediction)
+    for gold in gold_answers:
+        norm_gold = normalize_answer(gold)
+        if norm_gold in norm_pred or norm_pred in norm_gold:
+            return 1.0
+    return 0.0
+
+
+def evaluate(prediction_text: str, gold_answers: list[str]) -> dict:
+    """Evaluate a single QA prediction against gold answers.
+
+    Returns dict with: em, f1, sub_em, predicted_answer, gold_answers.
+    """
+    answer = extract_answer(prediction_text)
+    return {
+        "em": exact_match(answer, gold_answers),
+        "f1": f1_score(answer, gold_answers),
+        "sub_em": sub_em(answer, gold_answers),
+        "predicted_answer": answer,
+        "gold_answers": gold_answers,
+    }

skillopt/envs/searchqa/prompts/analyst_error.md

@@ -0,0 +1,46 @@
+You are an expert failure-analysis agent for question answering tasks.
+
+You will be given MULTIPLE failed QA agent responses from a single minibatch
+and the current skill document. Each trajectory includes the agent's response
+and an evaluation result showing the predicted answer vs. the gold answer(s).
+
+Your job is to identify the most important COMMON failure patterns across
+the batch and propose a concise set of skill edits.
+
+## Failure Type Categories
+- **rule_missing**: the skill lacks a relevant rule for this type of question
+- **rule_wrong**: an existing skill rule is misleading or incorrect
+- **rule_ignored**: the skill has the right rule but the agent did not follow it
+- **answer_format**: the agent found the right information but formatted it incorrectly
+- **other**: none of the above
+
+## Analysis Process
+1. Read ALL failed trajectories in the minibatch.
+2. Carefully compare each predicted answer against the gold answer(s) —
+   understand exactly WHY the Exact Match failed.
+3. Identify the most prevalent, systematic failure patterns across them.
+4. For each pattern, classify its failure type.
+5. Propose skill edits that address the COMMON patterns — not individual edge cases.
+6. Edits must be generalizable; do not hardcode question-specific values.
+7. Only patch gaps in the skill — do not duplicate existing content.
+
+You will be told the maximum number of edits (the budget L). Produce AT MOST L edits,
+focusing on the highest-impact patterns. You may produce fewer if warranted.
+
+Respond ONLY with a valid JSON object (no markdown fences, no extra text):
+{
+  "batch_size": <number of trajectories analysed>,
+  "failure_summary": [
+    {"failure_type": "<type>", "count": <int>, "description": "<one-line>"}
+  ],
+  "patch": {
+    "reasoning": "<why these edits address the batch's common failures>",
+    "edits": [
+      {"op": "append",       "content": "<markdown to add at end of skill>"},
+      {"op": "insert_after", "target": "<exact heading/text to insert after>", "content": "<markdown>"},
+      {"op": "replace",      "target": "<exact text to replace>",              "content": "<replacement>"},
+      {"op": "delete",       "target": "<exact text to remove>"}
+    ]
+  }
+}
+Only include edits that are needed. "edits" can be an empty list if no patch is warranted.

skillopt/envs/searchqa/prompts/analyst_success.md

@@ -0,0 +1,32 @@
+You are an expert success-pattern analyst for AI question answering agents.
+
+You will be given MULTIPLE successful QA agent responses from a single minibatch
+and the current skill document. Your job is to identify generalizable behavior
+patterns that are COMMON across the batch and worth encoding in the skill.
+
+## Rules
+- Only propose patches for patterns NOT already covered in the skill.
+- Focus on patterns that appear across MULTIPLE trajectories in the batch.
+- Be concise. Patterns must generalize beyond specific questions.
+- Prefer reinforcing existing sections over adding new top-level sections.
+- If the agents' success involved a smart reading strategy or disambiguation
+  approach, consider reinforcing that in the patch.
+
+You will be told the maximum number of edits (the budget L). Produce AT MOST L edits,
+focusing on the most broadly applicable patterns. You may produce fewer if warranted.
+
+Respond ONLY with a valid JSON object:
+{
+  "batch_size": <number of trajectories analysed>,
+  "success_patterns": ["<pattern 1>", "<pattern 2>"],
+  "patch": {
+    "reasoning": "<why these patterns are worth encoding>",
+    "edits": [
+      {"op": "append",       "content": "<markdown>"},
+      {"op": "insert_after", "target": "<heading/text>", "content": "<markdown>"},
+      {"op": "replace",      "target": "<old text>",     "content": "<new text>"},
+      {"op": "delete",       "target": "<exact text to remove>"}
+    ]
+  }
+}
+"edits" may be empty if the skill already covers all observed patterns.

skillopt/envs/searchqa/prompts/deep_probe.md

@@ -0,0 +1,27 @@
+You are an expert diagnostic-probe designer for retrieval-style question answering tasks.
+
+You will be shown representative trajectories, the current student skill, the student's prompt context,
+and the evaluation result including the gold answer. There is NO hidden chain-of-thought reference.
+Design one SMALL diagnostic instruction that exposes the student's intermediate reading or evidence-selection state
+without materially changing the original scaffold.
+
+## Hard Constraints
+1. Do NOT substantially change the original scaffold.
+2. Do NOT prescribe a brand-new multi-step solving procedure.
+3. You MAY ask for a short structured readout of intermediate conclusions, evidence candidates, or elimination decisions.
+4. Do NOT ask for exhaustive quotation of the whole context or a full chain-of-thought.
+5. Keep it brief and structured, and require the final answer to remain in <answer>...</answer>.
+6. Use the gold answer only to target a useful probe; do not simply force the student to restate the gold answer.
+
+## Good Probe Targets
+- the most likely supporting span or document cue
+- top answer candidate and runner-up
+- decisive lexical clue / entity / date / title
+- why a tempting alternative was rejected
+- 2-4 short intermediate conclusions that directly support the final answer
+
+Respond ONLY with a valid JSON object:
+{
+  "reasoning": "<why this probe is informative>",
+  "probe_instruction": "<the exact instruction text to append to the student prompt>"
+}

skillopt/envs/searchqa/prompts/rollout_system.md

@@ -0,0 +1,13 @@
+You are an expert question answering agent.
+
+{skill_section}## Task Format
+You will receive a CONTEXT containing document passages and a QUESTION.
+Read the context carefully and answer the question based on the information provided.
+
+## Answer Format
+Think step by step, then provide your final answer inside <answer>...</answer> tags.
+Keep your answer concise — typically a few words or a short phrase.
+Do not repeat the question. Do not include unnecessary explanation in the answer tags.
+
+Example:
+<answer>Abraham Lincoln</answer>

skillopt/envs/searchqa/reflect.py

@@ -0,0 +1,4 @@
+"""SearchQA Reflect stage.
+
+Prompts are now loaded from .md files by the base adapter.
+"""

skillopt/envs/searchqa/rollout.py

@@ -0,0 +1,455 @@
+"""SearchQA rollout — single-turn QA agent + batch execution.
+
+The QA agent receives a skill document, question, and context passages,
+then produces an answer in <answer>...</answer> tags.
+
+Public API
+----------
+- :func:`process_one`  — run + evaluate one QA item
+- :func:`run_batch`    — parallel execution of a list of items
+"""
+from __future__ import annotations
+
+import json
+import os
+import time
+import traceback
+from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait
+
+from skillopt.model import chat_student, get_student_backend, is_student_exec_backend
+from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_student_exec
+from skillopt.prompts import load_prompt
+from skillopt.envs.searchqa.evaluator import evaluate
+
+
+# ── Prompt templates ─────────────────────────────────────────────────────────
+
+_MAX_CONTEXT_CHARS = 6000
+
+
+def _truncate_context(context: str, max_chars: int = _MAX_CONTEXT_CHARS) -> str:
+    """Truncate context at [DOC] boundaries to stay within budget."""
+    if len(context) <= max_chars:
+        return context
+    docs = context.split("[DOC]")
+    result = ""
+    for doc in docs:
+        candidate = result + "[DOC]" + doc if result else doc
+        if len(candidate) > max_chars:
+            break
+        result = candidate
+    if not result:
+        result = context[:max_chars] + "\n...[truncated]"
+    return result
+
+
+def _build_system(skill_content: str) -> str:
+    if skill_content.strip():
+        skill_section = f"## Skill\n{skill_content.strip()}\n\n"
+    else:
+        skill_section = ""
+    return load_prompt("rollout_system", env="searchqa").format(skill_section=skill_section)
+
+
+def _build_user(
+    question: str,
+    context: str,
+    *,
+    diagnostic_mode: bool = False,
+    diagnostic_instruction: str = "",
+    diagnostic_trace_context: str = "",
+) -> str:
+    context = _truncate_context(context)
+    parts = [
+        f"## Context\n{context}",
+        f"## Question\n{question}",
+    ]
+    if diagnostic_trace_context.strip():
+        parts.append(
+            "## Previous Codex Trace Snapshot\n"
+            "This is a partial transcript from an earlier attempt. Use it as your current reasoning context.\n\n"
+            f"{diagnostic_trace_context.strip()}"
+        )
+    if diagnostic_mode and diagnostic_instruction.strip():
+        parts.append(f"## Training Readout\n{diagnostic_instruction.strip()}")
+    return "\n\n".join(parts)
+
+
+def _build_codex_skill(skill_content: str) -> str:
+    return render_skill_md(
+        skill_content,
+        description="Dynamic ReflACT skill for solving the current SearchQA example.",
+        preamble=(
+            "Use this skill when solving the current SearchQA task.\n"
+            "Read the provided context carefully, ground the answer in that context,\n"
+            "and return the final answer inside <answer>...</answer>."
+        ),
+    )
+
+
+def _run_codex_once(
+    *,
+    pred_dir: str,
+    skill_content: str,
+    question: str,
+    context: str,
+    model: str,
+    timeout: int,
+    diagnostic_mode: bool = False,
+    diagnostic_instruction: str = "",
+    diagnostic_trace_context: str = "",
+    previous_response: str = "",
+) -> tuple[str, str, str, str]:
+    user = _build_user(
+        question,
+        context,
+        diagnostic_mode=diagnostic_mode,
+        diagnostic_instruction=diagnostic_instruction,
+        diagnostic_trace_context=diagnostic_trace_context,
+    )
+    task_parts = [user]
+    if previous_response:
+        task_parts.append(
+            "## Previous Attempt\n"
+            f"{previous_response}\n\n"
+            "Review it against the same context and question. If needed, correct it."
+        )
+    task_text = "\n\n".join(task_parts)
+    skill_md = _build_codex_skill(skill_content)
+    work_dir = os.path.join(pred_dir, "codex_exec")
+    prepare_workspace(
+        work_dir=work_dir,
+        skill_md=skill_md,
+        task_text=task_text,
+    )
+    prompt = (
+        "Use the `skillopt-student` skill available in this workspace.\n"
+        "Read `task.md` and answer the SearchQA question.\n"
+        "Return the final answer inside <answer>...</answer>."
+    )
+    final_message, raw = run_student_exec(
+        work_dir=work_dir,
+        prompt=prompt,
+        model=model,
+        timeout=timeout,
+    )
+    return final_message or raw, raw, skill_md, task_text
+
+
+# ── Single-item execution ───────────────────────────────────────────────────
+
+
+def process_one(
+    item: dict,
+    out_root: str,
+    skill_content: str,
+    max_turns: int = 1,
+    diagnostic_mode: bool = False,
+    diagnostic_instruction: str = "",
+    diagnostic_trace_context: str = "",
+    exec_timeout: int = 120,
+) -> dict:
+    """Process a single QA item: run agent + evaluate.
+
+    Parameters
+    ----------
+    item : dict
+        Must have keys: ``id``, ``question``, ``context``, ``answers``.
+    out_root : str
+        Output directory (predictions saved under ``predictions/<id>/``).
+    skill_content : str
+        Current skill document text.
+    max_turns : int
+        Max reasoning turns (1 = single-turn QA).
+
+    Returns
+    -------
+    dict
+        Result with ``hard`` (EM as int), ``soft`` (F1), etc.
+    """
+    item_id = str(item["id"])
+    question = item["question"]
+    context = item.get("context", "")
+    gold_answers = item.get("answers", [])
+
+    result = {
+        "id": item_id,
+        "question": question,
+        "em": 0.0,
+        "f1": 0.0,
+        "sub_em": 0.0,
+        "hard": 0,
+        "soft": 0.0,
+        "predicted_answer": "",
+        "gold_answers": gold_answers,
+        "response": "",
+        "fail_reason": "",
+        "agent_ok": False,
+        "n_turns": 0,
+    }
+
+    try:
+        pred_dir = os.path.join(out_root, "predictions", item_id)
+        os.makedirs(pred_dir, exist_ok=True)
+
+        if is_student_exec_backend():
+            from skillopt.model import azure_openai as _llm
+
+            conversation: list[dict] = []
+            response = ""
+            system = ""
+            user = ""
+            for turn in range(max_turns):
+                response, raw, system, user = _run_codex_once(
+                    pred_dir=pred_dir,
+                    skill_content=skill_content,
+                    question=question,
+                    context=context,
+                    model=_llm.STUDENT_DEPLOYMENT,
+                    timeout=exec_timeout,
+                    diagnostic_mode=diagnostic_mode if turn == 0 else False,
+                    diagnostic_instruction=diagnostic_instruction if turn == 0 else "",
+                    diagnostic_trace_context=diagnostic_trace_context if turn == 0 else "",
+                    previous_response=response if turn > 0 else "",
+                )
+                conversation.append({"type": "message", "turn": turn + 1, "content": response})
+                if turn > 0 and "<answer>" in response.lower():
+                    break
+
+            result["response"] = response
+            result["agent_ok"] = True
+            result["n_turns"] = len(conversation)
+
+            with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w") as f:
+                f.write(system)
+            with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w") as f:
+                f.write(user)
+            with open(os.path.join(pred_dir, "conversation.json"), "w") as f:
+                json.dump(conversation, f, ensure_ascii=False, indent=2)
+
+            eval_result = evaluate(response, gold_answers)
+            result["em"] = eval_result["em"]
+            result["f1"] = eval_result["f1"]
+            result["sub_em"] = eval_result["sub_em"]
+            result["predicted_answer"] = eval_result["predicted_answer"]
+            result["hard"] = int(eval_result["em"])
+            result["soft"] = eval_result["f1"]
+            if eval_result["em"] < 1.0:
+                result["fail_reason"] = (
+                    f"EM=0: predicted '{eval_result['predicted_answer']}' "
+                    f"but expected {gold_answers}"
+                )
+            eval_detail = (
+                f"[EVALUATION RESULT]\n"
+                f"Question: {question}\n"
+                f"Predicted answer: {eval_result['predicted_answer']!r}\n"
+                f"Gold answers: {gold_answers!r}\n"
+                f"Exact Match: {eval_result['em']}\n"
+                f"F1: {eval_result['f1']:.4f}"
+            )
+            conversation.append({"role": "system", "content": eval_detail})
+            with open(os.path.join(pred_dir, "conversation.json"), "w") as f:
+                json.dump(conversation, f, ensure_ascii=False, indent=2)
+            return result
+
+        system = _build_system(skill_content)
+        user = _build_user(
+            question,
+            context,
+            diagnostic_mode=diagnostic_mode,
+            diagnostic_instruction=diagnostic_instruction,
+            diagnostic_trace_context=diagnostic_trace_context,
+        )
+
+        conversation: list[dict] = []
+        response = ""
+
+        for turn in range(max_turns):
+            if turn == 0:
+                resp_text, _ = chat_student(
+                    system=system, user=user,
+                    max_completion_tokens=512,
+                    retries=5, stage="rollout",
+                    timeout=exec_timeout,
+                )
+            else:
+                refinement = (
+                    f"Your previous answer was:\n{response}\n\n"
+                    f"Review it against the context and question. "
+                    f"If correct, repeat it. If wrong, provide a corrected answer.\n"
+                    f"Use <answer>...</answer> tags for your final answer."
+                )
+                resp_text, _ = chat_student(
+                    system=system, user=refinement,
+                    max_completion_tokens=512,
+                    retries=5, stage="rollout",
+                    timeout=exec_timeout,
+                )
+
+            response = resp_text
+            conversation.append({"type": "message", "turn": turn + 1, "content": resp_text})
+
+            if turn > 0 and "<answer>" in resp_text.lower():
+                break
+
+        result["response"] = response
+        result["agent_ok"] = True
+        result["n_turns"] = len(conversation)
+
+        # Save conversation
+        with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w") as f:
+            f.write(system)
+        with open(os.path.join(pred_dir, "student_user_prompt.txt"), "w") as f:
+            f.write(user)
+        with open(os.path.join(pred_dir, "conversation.json"), "w") as f:
+            json.dump(conversation, f, ensure_ascii=False, indent=2)
+
+        # Evaluate
+        eval_result = evaluate(response, gold_answers)
+        result["em"] = eval_result["em"]
+        result["f1"] = eval_result["f1"]
+        result["sub_em"] = eval_result["sub_em"]
+        result["predicted_answer"] = eval_result["predicted_answer"]
+        result["hard"] = int(eval_result["em"])
+        result["soft"] = eval_result["f1"]
+
+        if eval_result["em"] < 1.0:
+            result["fail_reason"] = (
+                f"EM=0: predicted '{eval_result['predicted_answer']}' "
+                f"but expected {gold_answers}"
+            )
+
+        # Append eval details to conversation for the analyst
+        eval_detail = (
+            f"[EVALUATION RESULT]\n"
+            f"Question: {question}\n"
+            f"Predicted answer: {eval_result['predicted_answer']!r}\n"
+            f"Gold answers: {gold_answers!r}\n"
+            f"Exact Match: {eval_result['em']}\n"
+            f"F1: {eval_result['f1']:.4f}"
+        )
+        conversation.append({
+            "role": "system",
+            "content": eval_detail,
+        })
+        # Re-save enriched conversation
+        with open(os.path.join(pred_dir, "conversation.json"), "w") as f:
+            json.dump(conversation, f, ensure_ascii=False, indent=2)
+
+    except Exception as e:  # noqa: BLE001
+        result["fail_reason"] = f"error: {e}"
+
+    return result
+
+
+# ── Batch execution ──────────────────────────────────────────────────────────
+
+
+def run_batch(
+    items: list[dict],
+    out_root: str,
+    skill_content: str,
+    max_turns: int = 1,
+    exec_timeout: int = 120,
+    workers: int = 64,
+    diagnostic_mode: bool = False,
+    diagnostic_instruction: str = "",
+    diagnostic_trace_context_by_id: dict[str, str] | None = None,
+    task_timeout: int = 600,
+) -> list[dict]:
+    """Run QA agent on all items with ThreadPoolExecutor. Resume-aware."""
+    task_timeout = max(int(task_timeout), int(exec_timeout) + 60)
+    results_path = os.path.join(out_root, "results.jsonl")
+    os.makedirs(out_root, exist_ok=True)
+
+    # Resume: load already-done
+    done_ids: set[str] = set()
+    existing: list[dict] = []
+    if os.path.exists(results_path):
+        with open(results_path) as f:
+            for line in f:
+                try:
+                    r = json.loads(line)
+                    done_ids.add(str(r["id"]))
+                    existing.append(r)
+                except Exception:
+                    pass
+
+    pending = [it for it in items if str(it["id"]) not in done_ids]
+    if not pending:
+        return existing
+
+    results = list(existing)
+
+    def _timeout_result(item: dict) -> dict:
+        return {
+            "id": str(item["id"]),
+            "question": item.get("question", ""),
+            "task_description": item.get("question", ""),
+            "task_type": item.get("task_type") or "searchqa",
+            "hard": 0,
+            "soft": 0.0,
+            "predicted_answer": "",
+            "response": "",
+            "fail_reason": f"task-timeout-{task_timeout}s",
+            "agent_ok": False,
+            "n_turns": 0,
+            "gold_answer": item.get("answers", []),
+            "phase": "timeout",
+        }
+
+    def _error_result(item: dict, exc: Exception) -> dict:
+        row = _timeout_result(item)
+        row["phase"] = "error"
+        row["fail_reason"] = f"unexpected: {type(exc).__name__}: {exc}"
+        return row
+
+    started_at: dict[str, float] = {}
+
+    def _run_one(item: dict) -> dict:
+        started_at[str(item["id"])] = time.time()
+        return process_one(
+            item,
+            out_root,
+            skill_content,
+            max_turns,
+            diagnostic_mode,
+            diagnostic_instruction,
+            (diagnostic_trace_context_by_id or {}).get(str(item["id"]), ""),
+            exec_timeout,
+        )
+
+    with open(results_path, "a") as outf:
+        ex = ThreadPoolExecutor(max_workers=workers)
+        try:
+            futs = {ex.submit(_run_one, it): it for it in pending}
+            pending_futs = set(futs)
+            while pending_futs:
+                done, _ = wait(pending_futs, timeout=5, return_when=FIRST_COMPLETED)
+                now = time.time()
+                timed_out = [
+                    fut for fut in pending_futs - done
+                    if str(futs[fut]["id"]) in started_at
+                    and now - started_at[str(futs[fut]["id"])] >= task_timeout
+                ]
+                for fut in done:
+                    pending_futs.remove(fut)
+                    item = futs[fut]
+                    try:
+                        res = fut.result()
+                    except Exception as exc:  # noqa: BLE001
+                        res = _error_result(item, exc)
+                    results.append(res)
+                    outf.write(json.dumps(res, ensure_ascii=False) + "\n")
+                    outf.flush()
+                for fut in timed_out:
+                    pending_futs.remove(fut)
+                    fut.cancel()
+                    res = _timeout_result(futs[fut])
+                    results.append(res)
+                    outf.write(json.dumps(res, ensure_ascii=False) + "\n")
+                    outf.flush()
+        finally:
+            ex.shutdown(wait=False, cancel_futures=True)
+
+    return results

skillopt/envs/searchqa/skills/initial.md

@@ -0,0 +1,3 @@
+# Question Answering Skill
+
+(No learned rules yet. Rules will be added through the reflection process.)