SkillOpt/reflact/envs/officeqa/rollout.py

from __future__ import annotations

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

from reflact.envs.officeqa.evaluator import evaluate
from reflact.envs.officeqa.tool_runtime import resolve_candidate_files, resolve_docs_roots, run_tool
from reflact.model import chat_student_messages, get_student_backend, is_student_exec_backend
from reflact.model.codex_harness import prepare_workspace, render_skill_md, run_student_exec
from reflact.prompts import load_prompt

_TOOL_SCHEMAS = [
    {
        "type": "function",
        "function": {
            "name": "glob",
            "description": "Find candidate local document files by filename or relative-path glob pattern.",
            "parameters": {
                "type": "object",
                "properties": {"pattern": {"type": "string"}},
                "required": ["pattern"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "read",
            "description": "Read a local text document excerpt by path and line window.",
            "parameters": {
                "type": "object",
                "properties": {
                    "path": {"type": "string"},
                    "start": {"type": "integer"},
                    "limit": {"type": "integer"},
                },
                "required": ["path"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "grep",
            "description": "Search a local text document for a literal pattern and return matching lines.",
            "parameters": {
                "type": "object",
                "properties": {
                    "pattern": {"type": "string"},
                    "path": {"type": "string"},
                },
                "required": ["pattern", "path"],
            },
        },
    },
]

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


def _build_user(
    item: dict,
    candidate_files: list[str],
    *,
    diagnostic_mode: bool = False,
    diagnostic_instruction: str = "",
    corpus_note: str = "",
) -> str:
    file_block = "\n".join(f"- {path}" for path in candidate_files[:20]) or "- none resolved"
    parts = [f"## Question\n{item['question']}"]
    if corpus_note.strip():
        parts.append(f"## Document Corpus\n{corpus_note.strip()}")
    parts.append(f"## Candidate Files\n{file_block}")
    if item.get("source_docs"):
        parts.append("## Source Hints\n" + "\n".join(f"- {hint}" for hint in item["source_docs"]))
    if diagnostic_mode and diagnostic_instruction.strip():
        parts.append(f"## Training Readout\n{diagnostic_instruction.strip()}")
    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 _docs_link_targets(docs_roots: list[str]) -> list[tuple[str, str]]:
    return [(root, os.path.join("docs", f"root_{idx}")) for idx, root in enumerate(docs_roots, start=1)]


def _workspace_doc_path(path: str, docs_roots: list[str]) -> str:
    resolved_path = os.path.realpath(path)
    for idx, root in enumerate(docs_roots, start=1):
        resolved_root = os.path.realpath(root)
        if resolved_path == resolved_root or resolved_path.startswith(resolved_root + os.sep):
            rel_path = os.path.relpath(resolved_path, resolved_root)
            return os.path.join("docs", f"root_{idx}", rel_path)
    return path


def _build_codex_skill(skill_content: str) -> str:
    return render_skill_md(
        skill_content,
        description="Dynamic ReflACT skill for solving the current OfficeQA local-document question.",
        preamble=(
            "Use this skill when answering the current OfficeQA question.\n"
            "Inspect the provided local document excerpts or files, ground the answer in the evidence,\n"
            "and return the final answer inside <answer>...</answer>."
        ),
    )


def _run_codex_once(
    *,
    pred_dir: str,
    item: dict,
    skill_content: str,
    candidate_files: list[str],
    docs_roots: list[str],
    model: str,
    timeout: int,
    diagnostic_mode: bool = False,
    diagnostic_instruction: str = "",
    previous_response: str = "",
) -> tuple[str, str, str, str]:
    rel_files = [_workspace_doc_path(path, docs_roots) for path in candidate_files[:20]]
    corpus_note = (
        "The full OfficeQA document corpus is available under `docs/`. "
        "The candidate files below are source hints or likely starting points; search the full corpus if needed."
    )
    user = _build_user(
        item,
        rel_files,
        diagnostic_mode=diagnostic_mode,
        diagnostic_instruction=diagnostic_instruction,
        corpus_note=corpus_note,
    )
    task_parts = [user]
    if previous_response:
        task_parts.append(
            "## Previous Attempt\n"
            f"{previous_response}\n\n"
            "Review the local documents again and correct the answer if needed."
        )
    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,
        link_dirs=_docs_link_targets(docs_roots),
    )
    prompt = (
        "Use the `reflact-student` skill available in this workspace.\n"
        "Read `task.md`, inspect or search the full OfficeQA corpus under `docs/`, and answer the question.\n"
        "Treat candidate files in `task.md` as hints, not an access limit.\n"
        "Return the final answer inside <answer>...</answer>."
    )
    final_message, raw = run_student_exec(
        work_dir=work_dir,
        prompt=prompt,
        model=model,
        timeout=timeout,
        data_dirs=docs_roots,
    )
    return final_message or raw, raw, skill_md, task_text


def process_one(
    item: dict,
    out_root: str,
    skill_content: str,
    *,
    max_tool_turns: int = 12,
    data_dirs: list[str] | str | None = None,
    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)

    docs_roots = resolve_docs_roots(data_dirs)
    candidate_files = resolve_candidate_files(item.get("source_files", []), docs_roots)
    system = _build_system(skill_content)
    user = _build_user(item, candidate_files, diagnostic_mode=diagnostic_mode, diagnostic_instruction=diagnostic_instruction)

    messages: list[dict] = [
        {"role": "system", "content": system},
        {"role": "user", "content": user},
    ]
    conversation: list[dict] = [{"role": "user", "content": user}]
    final_response = ""
    final_answer = ""
    fail_reason = ""

    allowed_files = [os.path.basename(path) for path in candidate_files]

    try:
        if is_student_exec_backend():
            from reflact.model import azure_openai as _llm

            response = ""
            system = ""
            user = ""
            for turn in range(1, max_tool_turns + 1):
                response, _raw, system, user = _run_codex_once(
                    pred_dir=pred_dir,
                    item=item,
                    skill_content=skill_content,
                    candidate_files=candidate_files,
                    docs_roots=docs_roots,
                    model=_llm.STUDENT_DEPLOYMENT,
                    timeout=180,
                    diagnostic_mode=diagnostic_mode if turn == 1 else False,
                    diagnostic_instruction=diagnostic_instruction if turn == 1 else "",
                    previous_response=response if turn > 1 else "",
                )
                final_response = response
                conversation.append({"type": "message", "turn": turn, "content": response})
                if "<answer>" in response.lower():
                    final_answer = _extract_answer(response)
                    break
            if not final_answer:
                fail_reason = f"Exceeded codex turn budget ({max_tool_turns})"
            system = system or _build_codex_skill(skill_content)
            user = user or _build_user(item, [_workspace_doc_path(path, docs_roots) for path in candidate_files])
        else:
            for turn in range(1, max_tool_turns + 1):
                message, _ = chat_student_messages(
                    messages=messages,
                    max_completion_tokens=768,
                    retries=5,
                    stage="rollout",
                    tools=_TOOL_SCHEMAS,
                    tool_choice="auto",
                    return_message=True,
                )
                response = message.content or ""
                final_response = response
                assistant_message = {"role": "assistant", "content": response}
                if getattr(message, "tool_calls", None):
                    assistant_message["tool_calls"] = [tool_call.model_dump(mode="json") for tool_call in message.tool_calls]
                messages.append(assistant_message)
                conversation.append({"type": "message", "content": response})

                if getattr(message, "tool_calls", None):
                    for tool_call in message.tool_calls:
                        tool_name = tool_call.function.name
                        arguments = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {}
                        cmd, obs = run_tool(tool_name, arguments, allowed_roots=docs_roots, allowed_files=allowed_files)
                        conversation.append({"type": "tool_call", "cmd": cmd, "obs": obs})
                        messages.append({
                            "role": "tool",
                            "tool_call_id": tool_call.id,
                            "content": obs,
                        })
                    continue

                if "<answer>" in response.lower():
                    final_answer = _extract_answer(response)
                    break
                if turn == max_tool_turns:
                    fail_reason = f"Exceeded tool-turn budget ({max_tool_turns})"
                else:
                    fail_reason = "Model neither produced a tool request nor a final answer"
                    break
    except Exception as e:  # noqa: BLE001
        fail_reason = f"error: {e}"

    with open(os.path.join(pred_dir, "student_system_prompt.txt"), "w", encoding="utf-8") as f:
        f.write(system)
    with open(os.path.join(pred_dir, "student_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)

    eval_result = evaluate(final_answer, item.get("ground_truth", "")) if final_answer else {"em": 0.0, "f1": 0.0, "predicted_answer": "", "gold_answer": item.get("ground_truth", "")}
    result = {
        "id": item_id,
        "question": item.get("question", ""),
        "task_type": item.get("task_type", "officeqa"),
        "task_description": item.get("question", ""),
        "predicted_answer": eval_result["predicted_answer"],
        "response": final_response,
        "ground_truth": item.get("ground_truth", ""),
        "source_files": item.get("source_files", []),
        "resolved_source_paths": candidate_files,
        "hard": int(eval_result["em"]),
        "soft": eval_result["f1"],
        "fail_reason": fail_reason or ("" if eval_result["em"] else f"predicted '{eval_result['predicted_answer']}' but expected '{item.get('ground_truth', '')}'"),
        "agent_ok": not fail_reason,
        "n_turns": len(conversation),
        "student_system_prompt": system,
        "student_user_prompt": user,
    }
    return result


def run_batch(
    items: list[dict],
    out_root: str,
    skill_content: str,
    *,
    workers: int = 8,
    max_tool_turns: int = 12,
    data_dirs: list[str] | str | None = None,
    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

    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,
                data_dirs=data_dirs,
                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)
            outf.write(json.dumps(res, ensure_ascii=False) + "\n")
            outf.flush()
    return results