Auto-Finance/src/mcp/tools/extraction.py

"""抽取交接工具：渲染页面、读文本层、落盘抽取结果

抽取本身由宿主 Agent 的多模态能力完成（见 skills/reimbursement 的 schema 说明）。
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

from ...cases import CaseStore
from ...core.extraction.render import prepare_image, render_pdf
from ..envelope import fail, ok
from ..schemas import normalize_extraction, schema_doc

_store = CaseStore()


def case_render_pages(
    case_id: str, files: list[str] | None = None, dpi: int = 150, force: bool = False
) -> dict[str, Any]:
    """为材料准备可读取的页面图片与 PDF 文本层。

    返回每项材料的页 PNG 路径、文本层字符数与自动转正角度。
    有文本层（text_chars > 0）时应优先读文本层，比读图更准且省 token。
    """
    d = _store.resolve(case_id)
    if d is None:
        return fail(f"会话不存在: {case_id}", next_action="先 case_create 或 case_add_files")
    materials = _store.materials(d)
    if files:
        wanted = set(files)
        materials = [m for m in materials if m.name in wanted]
        missing_files = sorted(wanted - {m.name for m in materials})
    else:
        missing_files = []
    if not materials:
        return fail("没有可渲染的材料", next_action="先 case_add_files 拷入 PDF 或图片", missing=missing_files)

    rendered: list[dict[str, Any]] = []
    for m in materials:
        out_dir = _store.pages_dir(d, m.stem)
        if m.suffix.lower() == ".pdf":
            info = render_pdf(m, out_dir, dpi=dpi, force=force)
        else:
            info = prepare_image(m, out_dir)
        info["file"] = m.name
        info["needs_extraction"] = not _store.extraction_path(d, m.name).exists()
        rendered.append(info)
    return ok(
        rendered=rendered,
        not_found=missing_files,
        hint="读 pages 里的 PNG 或 text_file，按 extraction_spec 的字段定义产出 JSON，再 arb_case_save_extraction",
    )


def case_read_text(case_id: str, file: str) -> dict[str, Any]:
    """读取某材料渲染时提取的 PDF 文本层；无文本层时返回空。"""
    d = _store.resolve(case_id)
    if d is None:
        return fail(f"会话不存在: {case_id}")
    text_file = _store.pages_dir(d, Path(file).stem) / "text.txt"
    if not text_file.exists():
        return fail("尚未渲染该材料或它不是 PDF", next_action=f"先 case_render_pages(case_id, files=['{file}'])")
    text = text_file.read_text(encoding="utf-8")
    return ok(file=file, text_chars=len(text), text=text)


def case_save_extraction(case_id: str, source_file: str, data: dict[str, Any]) -> dict[str, Any]:
    """把单张票据的抽取结果写入会话缓存（按 <stem><suffix>.json 约定）。

    日期会规范成 YYYY-MM-DD、金额规范成两位小数字符串；
    返回 warnings 提示必填缺失，但不阻断写入（完整性由 arb_validate_info 判定）。
    """
    d = _store.resolve(case_id)
    if d is None:
        return fail(f"会话不存在: {case_id}")
    src = d / source_file
    if not src.exists():
        return fail(f"会话里没有源文件 {source_file}", next_action="先 case_add_files 或检查文件名拼写")
    if not isinstance(data, dict) or not data:
        return fail("data 必须是非空对象")
    try:
        normalized, warnings = normalize_extraction(data)
    except ValueError as exc:
        return fail(str(exc), next_action="修正 invoice_type 后重存；无法判读类型时用 attachments")
    if not normalized.get("_source_file"):
        normalized["_source_file"] = source_file
    path = _store.write_extraction(d, source_file, normalized)
    return ok(cached=str(path), invoice_type=normalized.get("invoice_type"), warnings=warnings)


def case_patch_extraction(case_id: str, source_file: str, patch: dict[str, Any]) -> dict[str, Any]:
    """局部修正已落盘的抽取结果（人工核对后改字段用）。

    只合并传入的键，返回每个被改字段的 from/to，避免静默覆盖；
    未在 schema 中出现的键会告警但仍写入（有些字段是给下游自定义用的）。
    """
    d = _store.resolve(case_id)
    if d is None:
        return fail(f"会话不存在: {case_id}")
    if not isinstance(patch, dict) or not patch:
        return fail("patch 必须是非空对象")
    from ...core.extraction.cache import load_cache as _load
    from ...core.extraction.cache import save_extraction

    current = _load(d).get(source_file)
    if current is None:
        return fail(f"该材料还没有抽取结果: {source_file}", next_action="先 case_save_extraction 落一份完整结果")
    from ..schemas import DOCUMENT_SCHEMAS, INVOICE_TYPES

    doc_type = str(current.get("invoice_type", ""))
    schema = DOCUMENT_SCHEMAS.get(doc_type, {})
    known = set(schema.get("required", [])) | set(schema.get("optional", [])) | {"invoice_type", "_source_file"}
    warnings: list[str] = []
    changed: dict[str, dict[str, Any]] = {}
    merged = dict(current)
    for key, value in patch.items():
        if key.startswith("_"):
            warnings.append(f"忽略内部字段 {key} 的修改")
            continue
        if key not in changed and merged.get(key) != value:
            changed[key] = {"from": merged.get(key), "to": value}
        if doc_type in INVOICE_TYPES and known and key not in known:
            warnings.append(f"{key} 不在 {doc_type} 的字段定义里，确认拼写无误")
        merged[key] = value
    if not changed:
        return ok(source_file=source_file, changed={}, warnings=warnings, note="patch 与现值相同，未做写入")
    normalized, normalize_warnings = normalize_extraction(merged)
    normalized["_source_file"] = source_file
    path = save_extraction(d, source_file, normalized)
    return ok(source_file=source_file, changed=changed, warnings=warnings + normalize_warnings, cached=str(path))


def extraction_spec(kind: str | None = None) -> dict[str, Any]:
    """返回抽取字段定义。kind 省略给全部文档类型；'info' 给 travel_info/normal_info 骨架。"""
    return ok(**schema_doc(kind))