Agent/examples/process_pipeline/script/search_eval/procedure-dsl/spec/tools/quote-source.py

#!/usr/bin/env python3
"""quote-source.py — 从原文 case json 按字符匹配查找并**原样返回**片段, 供填 value / directive。

为什么需要: value / directive 要填原文里的真实内容(那段 JSON 风格分析、那句完整提示词),
不能写 "<JSON 风格分析数据>" 这种空壳, 也不能凭记忆缩写。用本工具把原文那段逐字捞出来,
再粘进 workflow.json 的 value / directive。

匹配是**空白无关**的: 原文 body_text 常有乱换行 / 前导空格("交给\\n AI"), 直接子串匹配
多词查询会落空。本工具匹配时忽略所有空白, 但返回的是**原文逐字片段**(连原排版一起)。

用法 (Agent 走 bash_command 调):
    # 按几个关键词捞出周边片段 (默认前后各 300 字)
    python spec/tools/quote-source.py --source input/case-N.json --query "视觉风格分析"
    # 捞大块时放宽窗口
    python spec/tools/quote-source.py --source input/case-N.json --query "主色调" --window 800
    # 命中点落在 {...} 里时, 返回整个 JSON 块 (适合捞那段结构化分析)
    python spec/tools/quote-source.py --source input/case-N.json --query "视觉风格分析" --json-block
    # 捞一句完整提示词当 directive
    python spec/tools/quote-source.py --source input/case-N.json --query "请以 JSON 结构化"
    # 范围引用: 用首尾两个短锚点引出之间的整段长原文 (比 --window 猜长度精确)
    python spec/tools/quote-source.py --source input/case-N.json --from "请以 JSON" --to "500 字以内"
    python spec/tools/quote-source.py --source input/case-N.json --from "视觉风格分析" --to "电影感人像"

输出: 命中的原文逐字片段 (多处命中各自分隔)。无精确命中时给出最接近的一段 (标 ~approx)。
退出码: 0 有命中(含 approx) / 1 完全找不到 / 2 CLI/IO 错。
"""
import argparse
import json
import sys
from pathlib import Path

for _s in (sys.stdout, sys.stderr):
    if hasattr(_s, "reconfigure"):
        try:
            _s.reconfigure(encoding="utf-8", errors="replace")
        except Exception:
            pass


def _source_text(path: Path) -> str:
    """从 case json 抽可搜索文本: title + body_text (两套 schema 都有 body_text)。"""
    data = json.loads(path.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        return str(data)
    parts = []
    for k in ("title", "body_text"):
        v = data.get(k)
        if isinstance(v, str) and v:
            parts.append(v)
    return "\n".join(parts)


def _build_norm(text: str):
    """去掉所有空白, 返回 (norm_str, idxmap)。idxmap[i] = norm 第 i 个字符在原文里的下标。"""
    chars, idxmap = [], []
    for i, ch in enumerate(text):
        if ch.isspace():
            continue
        chars.append(ch)
        idxmap.append(i)
    return "".join(chars), idxmap


def _find_all(ns: str, nq: str, max_hits: int):
    hits, start = [], 0
    while len(hits) < max_hits:
        p = ns.find(nq, start)
        if p < 0:
            break
        hits.append(p)
        start = p + max(1, len(nq))
    return hits


def _enclosing_braces(text: str, lo: int, hi: int):
    """找包住 [lo,hi] 的最小 {...} (按花括号配平)。找不到返回 None。"""
    # 向左找候选 '{'
    depth = 0
    open_pos = None
    i = lo
    while i >= 0:
        if text[i] == "}":
            depth += 1
        elif text[i] == "{":
            if depth == 0:
                open_pos = i
                break
            depth -= 1
        i -= 1
    if open_pos is None:
        return None
    # 从 open_pos 向右配平
    depth = 0
    j = open_pos
    while j < len(text):
        if text[j] == "{":
            depth += 1
        elif text[j] == "}":
            depth -= 1
            if depth == 0:
                if j >= hi:
                    return (open_pos, j)
                return None
        j += 1
    return None


def _approx(ns: str, nq: str, idxmap, text: str, window: int):
    """无精确命中时, 用最长公共子串给最接近的一段 (标 approx)。"""
    from difflib import SequenceMatcher
    if not nq:
        return None
    m = SequenceMatcher(None, ns, nq, autojunk=False).find_longest_match(0, len(ns), 0, len(nq))
    if m.size == 0:
        return None
    lo_n, hi_n = m.a, m.a + m.size - 1
    return _slice(idxmap, text, lo_n, hi_n, window, False)


def _slice(idxmap, text: str, lo_n: int, hi_n: int, window: int, json_block: bool):
    o_lo = idxmap[lo_n]
    o_hi = idxmap[hi_n]
    if json_block:
        br = _enclosing_braces(text, o_lo, o_hi)
        if br:
            return text[br[0]: br[1] + 1], br[0], br[1]
    lo = max(0, o_lo - window)
    hi = min(len(text), o_hi + 1 + window)
    return text[lo:hi], lo, hi


def _range_between(text: str, frm: str, to: str):
    """范围引用: 返回 frm 首次出现 → 其后首次 to 结束 之间(含两端)的原文逐字片段。

    匹配空白无关; to 必须落在 frm 之后(取最近的一个), 给出最小包裹区间。找不到返回 None。
    """
    ns, idxmap = _build_norm(text)
    nf, _ = _build_norm(frm)
    nt, _ = _build_norm(to)
    if not nf or not nt:
        return None
    p = ns.find(nf)
    if p < 0:
        return None
    q = ns.find(nt, p + len(nf))      # to 必须在 from 之后
    if q < 0:
        return None
    o_lo = idxmap[p]
    o_hi = idxmap[q + len(nt) - 1]
    return text[o_lo:o_hi + 1], o_lo, o_hi


def main() -> int:
    ap = argparse.ArgumentParser(description="从原文 case json 按字符匹配原样捞片段 (供填 value/directive)")
    ap.add_argument("--source", required=True, help="原文 case json 路径")
    ap.add_argument("--query", default=None, help="关键词模式: 要查找的文字 (空白无关匹配; 配 --window / --json-block)")
    ap.add_argument("--from", dest="frm", default=None, help="范围模式: 起始锚点 (短串)")
    ap.add_argument("--to", dest="to", default=None, help="范围模式: 结束锚点 (短串); 与 --from 同用, 用两个短参数引出之间(含)的整段长原文")
    ap.add_argument("--ocr", default=None, help="配图 OCR 文本文件 (如 <case_dir>/_scratch/ocr.txt); 一并搜图片里的文字")
    ap.add_argument("--window", type=int, default=300, help="命中点前后各保留多少字 (default 300)")
    ap.add_argument("--max", type=int, default=3, help="每个语料最多返回几处命中 (default 3)")
    ap.add_argument("--json-block", action="store_true", help="命中落在 {...} 内时返回整个 JSON 块")
    args = ap.parse_args()

    src = Path(args.source)
    if not src.exists():
        print(f"source not found: {src}", file=sys.stderr)
        return 2
    try:
        corpora = [("原文", _source_text(src))]
    except (json.JSONDecodeError, OSError) as e:
        print(f"read source failed: {e}", file=sys.stderr)
        return 2
    # 第二语料: 配图 OCR 文本
    if args.ocr:
        ocr_p = Path(args.ocr)
        if ocr_p.exists():
            corpora.append(("配图OCR", ocr_p.read_text(encoding="utf-8")))
        else:
            print(f"# (--ocr {ocr_p} 不存在, 仅搜原文)", file=sys.stderr)

    # ── 范围模式: --from / --to 两个短锚点 → 返回之间(含)的整段原文 ──
    if args.frm is not None or args.to is not None:
        if not (args.frm and args.to):
            print("范围模式需同时给 --from 和 --to", file=sys.stderr)
            return 2
        out = []
        for label, text in corpora:
            rng = _range_between(text, args.frm, args.to)
            if rng:
                frag, lo, hi = rng
                warn = "  ⚠ 跨度很大, 确认锚点是否选窄了" if (hi - lo) > 4000 else ""
                out.append(f"# [{label}] 范围命中 (第 {lo}–{hi} 字, 共 {hi - lo + 1} 字{warn}; 下为两锚点间逐字原文)\n{frag}")
        if out:
            print("\n\n".join(out))
            return 0
        print(f"# 找不到从 {args.frm!r} 到 {args.to!r} 的范围 (检查: 两锚点是否各自存在、to 是否在 from 之后)", file=sys.stderr)
        return 1

    # ── 关键词模式: --query ──
    if not args.query:
        print("需要 --query, 或 (--from + --to) 范围模式", file=sys.stderr)
        return 2
    nq, _ = _build_norm(args.query)
    if not nq:
        print("query 去空白后为空", file=sys.stderr)
        return 2

    out, found = [], False
    for label, text in corpora:
        ns, idxmap = _build_norm(text)
        hits = _find_all(ns, nq, args.max)
        if hits:
            found = True
            out.append(f"# [{label}] {len(hits)} 处命中 (query={args.query!r}, 空白无关匹配; 下为逐字片段)")
            for k, p in enumerate(hits, 1):
                frag, lo, hi = _slice(idxmap, text, p, p + len(nq) - 1, args.window, args.json_block)
                out.append(f"\n--- [{label}] 命中 {k} (第 {lo}–{hi} 字) ---\n{frag}")
    if found:
        print("\n".join(out))
        return 0

    # 无精确命中 → 各语料给最接近的一段
    for label, text in corpora:
        ns, idxmap = _build_norm(text)
        ap_res = _approx(ns, nq, idxmap, text, args.window)
        if ap_res:
            frag, lo, hi = ap_res
            out.append(f"# [{label}] 无精确命中, 最接近的一段 (~approx, 第 {lo}–{hi} 字; 请核对):\n\n{frag}")
    if out:
        print("\n".join(out))
        return 0

    print(f"# 原文/OCR 里都找不到与 {args.query!r} 相关的内容", file=sys.stderr)
    return 1


if __name__ == "__main__":
    sys.exit(main())