SkillOpt/skillopt/envs/sealqa/rollout.py

from __future__ import annotations

import json
import os
import re
from concurrent.futures import ThreadPoolExecutor, as_completed

from skillopt.envs.sealqa.evaluator import score_sealqa
from skillopt.envs.sealqa.tool_runtime import web_fetch
from skillopt.model import chat_target, get_target_backend, is_target_exec_backend
from skillopt.model.codex_harness import prepare_workspace, render_skill_md, run_target_exec
from skillopt.prompts import load_prompt

_FINAL_RE = re.compile(r"<answer>(.*?)</answer>", re.IGNORECASE | re.DOTALL)


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="sealqa").format(skill_section=skill_section)


def _build_user(item: dict, *, diagnostic_mode: bool = False, diagnostic_instruction: str = '') -> str:
    parts = [f"## Question\n{item['question']}"]
    if item.get('search_results'):
        parts.append(f"## Search Results\n{item['search_results']}")
    if item.get('urls'):
        parts.append(f"## URL Hints\n{item['urls']}")
    if item.get('freshness'):
        parts.append(f"## Freshness\n{item['freshness']}")
    if item.get('question_types'):
        parts.append(f"## Question Types\n{item['question_types']}")
    if diagnostic_mode and diagnostic_instruction.strip():
        parts.append(f"## Training Readout\n{diagnostic_instruction.strip()}")
    parts.append('Use the provided search evidence as your primary context. Do not rely on external tool use.')
    return "\n\n".join(parts)


def _extract_answer(text: str) -> str:
    match = _FINAL_RE.search(text)
    if match:
        return match.group(1).strip()
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    return lines[-1] if lines else text.strip()


def _build_codex_skill(skill_content: str) -> str:
    return render_skill_md(
        skill_content,
        description="Dynamic ReflACT skill for solving the current SealQA evidence-grounded question.",
        preamble=(
            "Use this skill when answering the current SealQA question.\n"
            "Use the provided search evidence first, reconcile conflicts carefully,\n"
            "and return the final answer inside <answer>...</answer>."
        ),
    )


def _run_codex_once(
    *,
    pred_dir: str,
    skill_content: str,
    task_text: str,
    model: str,
    timeout: int,
    previous_response: str = '',
) -> tuple[str, str, str, str]:
    task_parts = [task_text]
    if previous_response:
        task_parts.append(
            "## Previous Attempt\n"
            f"{previous_response}\n\n"
            "Review the evidence again and correct the final answer if needed."
        )
    final_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=final_task_text,
    )
    prompt = (
        "Use the `skillopt-target` skill available in this workspace.\n"
        "Read `task.md`, answer the SealQA question using the provided evidence,\n"
        "and return the final answer inside <answer>...</answer>."
    )
    final_message, raw = run_target_exec(
        work_dir=work_dir,
        prompt=prompt,
        model=model,
        timeout=timeout,
    )
    return final_message or raw, raw, skill_md, final_task_text


def process_one(
    item: dict,
    out_root: str,
    skill_content: str,
    *,
    max_tool_turns: int = 12,
    diagnostic_mode: bool = False,
    diagnostic_instruction: str = '',
) -> dict:
    item_id = str(item['id'])
    pred_dir = os.path.join(out_root, 'predictions', item_id)
    os.makedirs(pred_dir, exist_ok=True)

    system = _build_system(skill_content)
    user = _build_user(
        item,
        diagnostic_mode=diagnostic_mode,
        diagnostic_instruction=diagnostic_instruction,
    )
    conversation: list[dict] = [{'role': 'user', 'content': user}]
    final_response = ''
    final_answer = ''
    fail_reason = ''

    try:
        if is_target_exec_backend():
            from skillopt.model import azure_openai as _llm

            response, _raw, system, user_for_save = _run_codex_once(
                pred_dir=pred_dir,
                skill_content=skill_content,
                task_text=user,
                model=_llm.TARGET_DEPLOYMENT,
                timeout=120,
            )
            final_response = response
            conversation.append({'type': 'message', 'content': response})
            if '<answer>' in response.lower():
                final_answer = _extract_answer(response)
            else:
                user = user_for_save
        else:
            response, _ = chat_target(
                system=system,
                user=user,
                max_completion_tokens=768,
                retries=5,
                stage='rollout',
            )
            final_response = response
            conversation.append({'type': 'message', 'content': response})
            if '<answer>' in response.lower():
                final_answer = _extract_answer(response)

        if not final_answer:
            urls_text = str(item.get('urls') or '').strip()
            fetched_blocks = []
            for raw_url in re.findall(r'https?://[^\s\]\[\'\",]+', urls_text)[:2]:
                try:
                    fetched = web_fetch(raw_url)
                except Exception as fetch_error:  # noqa: BLE001
                    fetched = f'URL: {raw_url}\n\n[fetch error: {fetch_error}]'
                fetched_blocks.append(fetched)
                conversation.append({'type': 'tool_call', 'cmd': f'web_fetch({raw_url!r})', 'obs': fetched})
            if fetched_blocks:
                retry_user = user + '\n\n## Fetched URL Content\n' + '\n\n'.join(fetched_blocks)
                if is_target_exec_backend():
                    retry_response, _raw, system, retry_user = _run_codex_once(
                        pred_dir=pred_dir,
                        skill_content=skill_content,
                        task_text=retry_user,
                        model=_llm.TARGET_DEPLOYMENT,
                        timeout=120,
                        previous_response=final_response,
                    )
                else:
                    retry_response, _ = chat_target(
                        system=system,
                        user=retry_user,
                        max_completion_tokens=768,
                        retries=5,
                        stage='rollout',
                    )
                final_response = retry_response
                conversation.append({'type': 'message', 'content': retry_response})
                if '<answer>' in retry_response.lower():
                    final_answer = _extract_answer(retry_response)
                else:
                    fail_reason = 'Model did not produce a final answer'
            else:
                fail_reason = 'Model did not produce a final answer'
    except Exception as e:  # noqa: BLE001
        fail_reason = f'error: {e}'

    with open(os.path.join(pred_dir, 'target_system_prompt.txt'), 'w', encoding='utf-8') as f:
        f.write(system)
    with open(os.path.join(pred_dir, 'target_user_prompt.txt'), 'w', encoding='utf-8') as f:
        f.write(user)
    with open(os.path.join(pred_dir, 'conversation.json'), 'w', encoding='utf-8') as f:
        json.dump(conversation, f, ensure_ascii=False, indent=2)

    score = score_sealqa(item.get('question', ''), item.get('ground_truth', ''), final_answer) if final_answer else 0.0
    result = {
        'id': item_id,
        'question': item.get('question', ''),
        'task_type': item.get('task_type', 'sealqa'),
        'task_description': item.get('question', ''),
        'predicted_answer': final_answer,
        'response': final_response,
        'ground_truth': item.get('ground_truth', ''),
        'hard': int(score >= 1.0),
        'soft': float(score),
        'fail_reason': fail_reason or ('' if score >= 1.0 else f"predicted '{final_answer}' but expected '{item.get('ground_truth', '')}'"),
        'agent_ok': not fail_reason,
        'n_turns': len(conversation),
        'target_system_prompt': system,
        'target_user_prompt': user,
    }
    return result


def run_batch(
    items: list[dict],
    out_root: str,
    skill_content: str,
    *,
    workers: int = 4,
    max_tool_turns: int = 12,
    diagnostic_mode: bool = False,
    diagnostic_instruction: str = '',
) -> list[dict]:
    results_path = os.path.join(out_root, 'results.jsonl')
    os.makedirs(out_root, exist_ok=True)

    done_ids: set[str] = set()
    existing: list[dict] = []
    if os.path.exists(results_path):
        with open(results_path, encoding='utf-8') as f:
            for line in f:
                try:
                    row = json.loads(line)
                except json.JSONDecodeError:
                    continue
                done_ids.add(str(row.get('id')))
                existing.append(row)

    pending = [item for item in items if str(item['id']) not in done_ids]
    if not pending:
        return existing

    total = len(existing) + len(pending)
    completed = len(existing)
    correct_count = sum(1 for r in existing if r.get("hard", 0))
    if existing:
        print(f"    [rollout] resuming: {completed}/{total} already done", flush=True)

    results = list(existing)
    with open(results_path, 'a', encoding='utf-8') as outf, ThreadPoolExecutor(max_workers=workers) as ex:
        futs = {
            ex.submit(
                process_one,
                item,
                out_root,
                skill_content,
                max_tool_turns=max_tool_turns,
                diagnostic_mode=diagnostic_mode,
                diagnostic_instruction=diagnostic_instruction,
            ): item
            for item in pending
        }
        for fut in as_completed(futs):
            res = fut.result()
            results.append(res)
            completed += 1
            if res.get("hard", 0):
                correct_count += 1
            acc = correct_count / completed if completed else 0
            print(
                f"    [rollout] {completed}/{total} "
                f"(acc={acc:.3f}) id={res.get('id', '?')} "
                f"hard={res.get('hard', '?')}",
                flush=True,
            )
            outf.write(json.dumps(res, ensure_ascii=False) + '\n')
            outf.flush()
    return results