odysseus/scripts/curate_public_search_seed_cases.py

#!/usr/bin/env python3
from __future__ import annotations

import argparse
import hashlib
import json
import random
import re
import time
from pathlib import Path
from typing import Any

from datasets import load_dataset

from run_odysseus_search_teacher_pipeline import call_deepseek_json, db_deepseek_endpoint


REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_OUT = REPO_ROOT / "data/evals/ody_public_search_seed_20260825/cases.json"
DEFAULT_LOCAL_SEEDS = [
    Path("/home/pewds/deep_research_task_ui_seeds.jsonl"),
]


QUESTION_RE = re.compile(r"\?$|^(?:who|what|when|where|why|how|which|can|does|do|is|are|was|were)\b", re.I)
PRIVATE_RE = re.compile(
    r"\b(my|our)\s+(?:email|inbox|calendar|notes?|documents?|files?|computer|desktop|downloads?|contacts?)\b|"
    r"\b(?:send|delete|archive|mark|reply to|draft|schedule|remind me|open my)\b",
    re.I,
)
TOO_CURRENT_RE = re.compile(r"\b(?:today|right now|current|latest|this week|this month|2026|2025)\b", re.I)


def stable_id(prefix: str, value: Any) -> str:
    text = json.dumps(value, sort_keys=True, ensure_ascii=True)
    return f"{prefix}_{hashlib.sha256(text.encode('utf-8')).hexdigest()[:16]}"


def clean_text(value: Any) -> str:
    return re.sub(r"\s+", " ", str(value or "")).strip()


def useful_question(text: str) -> bool:
    q = clean_text(text)
    if len(q) < 18 or len(q) > 240:
        return False
    if PRIVATE_RE.search(q):
        return False
    if not QUESTION_RE.search(q):
        return False
    if len(q.split()) < 5:
        return False
    return True


def prompt_variant(question: str, source: str, index: int) -> str:
    q = clean_text(question).rstrip("?")
    variants = [
        f"Search the web and answer this: {q}?",
        f"Can you look up {q} and give me the answer?",
        f"Find a reliable source for this and answer briefly: {q}?",
        f"Use search to verify: {q}?",
        f"I need a quick sourced answer: {q}?",
    ]
    if source == "hotpot_qa":
        variants.extend([
            f"Search for the two facts needed to answer this: {q}?",
            f"Look this up and combine the evidence: {q}?",
        ])
    return variants[index % len(variants)]


def add_candidate(out: list[dict[str, Any]], seen: set[str], *, source: str, question: str, answer: Any = "", family: str = "") -> None:
    question = clean_text(question)
    if not useful_question(question):
        return
    key = question.lower()
    if key in seen:
        return
    seen.add(key)
    idx = len(out)
    out.append({
        "source": source,
        "source_id": stable_id(source, question),
        "question": question,
        "answer_hint": clean_text(answer)[:220],
        "family": family or ("fresh_or_date_sensitive" if TOO_CURRENT_RE.search(question) else "public_fact_search"),
        "user": prompt_variant(question, source, idx),
    })


def sample_nq_open(out: list[dict[str, Any]], seen: set[str], target: int, seed: int) -> None:
    ds = load_dataset("nq_open", split="train", streaming=True)
    rng = random.Random(seed)
    for i, row in enumerate(ds):
        if i > 250_000 or len(out) >= target:
            break
        if rng.random() > 0.045:
            continue
        add_candidate(
            out,
            seen,
            source="nq_open",
            question=row.get("question"),
            answer=row.get("answer"),
            family="simple_public_fact",
        )


def sample_hotpot(out: list[dict[str, Any]], seen: set[str], target: int, seed: int) -> None:
    ds = load_dataset("hotpot_qa", "distractor", split="train", streaming=True)
    rng = random.Random(seed + 17)
    for i, row in enumerate(ds):
        if i > 180_000 or len(out) >= target:
            break
        if rng.random() > 0.075:
            continue
        add_candidate(
            out,
            seen,
            source="hotpot_qa",
            question=row.get("question"),
            answer=row.get("answer"),
            family=f"multi_hop_{clean_text(row.get('type') or 'qa')}",
        )


def load_local(out: list[dict[str, Any]], seen: set[str], paths: list[Path], target: int) -> None:
    for path in paths:
        if not path.exists():
            continue
        for line in path.read_text(encoding="utf-8").splitlines():
            if len(out) >= target:
                return
            if not line.strip():
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError:
                continue
            prompt = clean_text(row.get("prompt") or row.get("user") or row.get("question"))
            if not prompt or PRIVATE_RE.search(prompt) or len(prompt) > 1600:
                continue
            key = prompt.lower()
            if key in seen:
                continue
            seen.add(key)
            out.append({
                "source": f"local:{path.name}",
                "source_id": clean_text(row.get("task_id") or row.get("id") or stable_id(path.name, prompt)),
                "question": prompt,
                "answer_hint": clean_text(row.get("reference_solution") or row.get("answer"))[:500],
                "family": clean_text(row.get("task_family") or row.get("family") or "local_web_research"),
                "user": prompt,
            })


def heuristic_rank(item: dict[str, Any]) -> float:
    q = item["question"].lower()
    score = 0.0
    score += 1.0 if item["source"] == "nq_open" else 0.0
    score += 1.4 if item["source"] == "hotpot_qa" else 0.0
    score += 1.0 if item["source"].startswith("local:") else 0.0
    score += 0.4 if 7 <= len(q.split()) <= 22 else 0.0
    score += 0.5 if re.search(r"\b(which|compare|both|between|relationship|part of|head office)\b", q) else 0.0
    score += 0.3 if item.get("answer_hint") else 0.0
    score -= 0.7 if TOO_CURRENT_RE.search(q) else 0.0
    score -= 0.8 if re.search(r"\b(song|lyrics|movie cast|episode)\b", q) else 0.0
    return score


def deepseek_audit(endpoint: dict[str, str], items: list[dict[str, Any]], batch_size: int) -> dict[str, dict[str, Any]]:
    audits: dict[str, dict[str, Any]] = {}
    for start in range(0, len(items), batch_size):
        batch = items[start:start + batch_size]
        payload = {
            "task": "Audit public web-search SFT seed prompts. Pick prompts that are natural, generic, useful for teaching a web_search/web_fetch agent, and not private/user-data tasks.",
            "current_date": "2026-08-25",
            "rating_scale": "0 reject, 1 weak, 2 usable, 3 good, 4 excellent",
            "reject_if": [
                "requires private data, email, calendar, local files, account access, login, or sending/deleting actions",
                "too broad for a 1-3 web tool trace unless it is a small minority of deep research seeds",
                "answer is purely subjective or does not benefit from search",
                "current/date-sensitive but lacks a stable phrasing or source date expectation",
                "unsafe medical/legal/financial advice beyond general sourced information",
            ],
            "items": [
                {
                    "id": item["source_id"],
                    "source": item["source"],
                    "family": item["family"],
                    "user": item["user"],
                    "answer_hint": item.get("answer_hint") or "",
                }
                for item in batch
            ],
            "return_schema": {
                "audits": [
                    {"id": "string", "rating": 0, "keep": False, "family": "string", "reason": "string"}
                ]
            },
        }
        result = call_deepseek_json(endpoint, payload, max_tokens=5000, temperature=0.15, json_mode=True)
        for audit in result.get("audits") or []:
            if not isinstance(audit, dict):
                continue
            item_id = clean_text(audit.get("id"))
            if item_id:
                audits[item_id] = audit
        print(json.dumps({"stage": "deepseek_audit", "start": start, "batch": len(batch), "audited": len(audits)}), flush=True)
    return audits


def build_cases(items: list[dict[str, Any]], audits: dict[str, dict[str, Any]], count: int) -> list[dict[str, Any]]:
    ranked: list[tuple[float, dict[str, Any], dict[str, Any]]] = []
    for item in items:
        audit = audits.get(item["source_id"]) or {}
        rating = float(audit.get("rating") or 0)
        if audit and not audit.get("keep"):
            continue
        if rating < 2:
            continue
        ranked.append((rating * 10 + heuristic_rank(item), item, audit))
    ranked.sort(key=lambda x: x[0], reverse=True)
    cases = []
    family_counts: dict[str, int] = {}
    source_counts: dict[str, int] = {}
    for _score, item, audit in ranked:
        family = clean_text(audit.get("family") or item.get("family") or "web")
        source = item["source"]
        if family_counts.get(family, 0) >= max(40, count // 5):
            continue
        if source_counts.get(source, 0) >= max(80, int(count * 0.55)):
            continue
        cases.append({
            "id": f"public_search_seed_{len(cases):04d}",
            "kind": "web",
            "family": family,
            "source_dataset": source,
            "source_id": item["source_id"],
            "user": item["user"],
            "expect_first_tool": "web_search",
            "allow_web_search": True,
            "forbidden_final": ["WEB SEARCH RESULTS", "```sources", "Here are links", "Web sources", "from the search results", "snippets"],
            "why_search_needed": clean_text(audit.get("reason") or "public source-backed answer"),
            "answer_hint": item.get("answer_hint") or "",
        })
        family_counts[family] = family_counts.get(family, 0) + 1
        source_counts[source] = source_counts.get(source, 0) + 1
        if len(cases) >= count:
            break
    return cases


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--count", type=int, default=500)
    parser.add_argument("--candidate-count", type=int, default=900)
    parser.add_argument("--out", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--seed", type=int, default=20260825)
    parser.add_argument("--audit-batch-size", type=int, default=35)
    parser.add_argument("--skip-deepseek", action="store_true")
    parser.add_argument("--local-seed", action="append", type=Path, default=[])
    args = parser.parse_args()

    rng = random.Random(args.seed)
    candidates: list[dict[str, Any]] = []
    seen: set[str] = set()
    local_paths = args.local_seed or DEFAULT_LOCAL_SEEDS
    load_local(candidates, seen, local_paths, min(args.candidate_count, 120))
    sample_hotpot(candidates, seen, max(args.candidate_count // 2, 260), args.seed)
    sample_nq_open(candidates, seen, args.candidate_count, args.seed)
    rng.shuffle(candidates)
    candidates.sort(key=heuristic_rank, reverse=True)
    candidates = candidates[: args.candidate_count]

    endpoint = db_deepseek_endpoint()
    endpoint["model"] = args.__dict__.get("teacher_model") or endpoint.get("model") or "deepseek-chat"
    if args.skip_deepseek:
        audits = {
            item["source_id"]: {
                "id": item["source_id"],
                "rating": 3,
                "keep": True,
                "family": item["family"],
                "reason": "heuristic keep",
            }
            for item in candidates
        }
    else:
        audits = deepseek_audit(endpoint, candidates, args.audit_batch_size)

    cases = build_cases(candidates, audits, args.count)
    payload = {
        "generated_at": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
        "generator": Path(__file__).name,
        "current_date": "2026-08-25",
        "source_notes": [
            "nq_open / Natural Questions: CC-BY-SA-3.0 on Hugging Face.",
            "hotpot_qa: CC-BY-SA-4.0 on Hugging Face.",
            "local research seeds are prompt seeds only; inspect before training if exporting outside this workspace.",
        ],
        "candidate_count": len(candidates),
        "audit_count": len(audits),
        "cases": cases,
        "audit_summary": {
            "accepted_cases": len(cases),
            "sources": {source: sum(1 for c in cases if c.get("source_dataset") == source) for source in sorted({c.get("source_dataset") for c in cases})},
            "families": {family: sum(1 for c in cases if c.get("family") == family) for family in sorted({c.get("family") for c in cases})},
        },
    }
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(payload, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    (args.out.parent / "seed_audits.json").write_text(json.dumps({"audits": audits}, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    (args.out.parent / "seed_candidates.jsonl").write_text(
        "".join(json.dumps(item, ensure_ascii=False) + "\n" for item in candidates),
        encoding="utf-8",
    )
    print(json.dumps({"cases": len(cases), "candidates": len(candidates), "out": str(args.out)}, indent=2))
    if len(cases) < args.count:
        raise RuntimeError(f"Only built {len(cases)} cases; requested {args.count}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())