Auto-Finance/src/doc/llm_extractor.py

"""LLM 信息提取

使用 LLM 从 PDF 文本/图片、支付截图中提取结构化数据。

## 功能模块

- **统一文档提取**：使用一套提示词，LLM 自行判断文档类型（发票/支付记录/出差事前申请单等），支持 JSON 格式输出。
- **差旅信息提取**：综合多张发票、支付记录和匹配结果，提取出差事由、地点、时间等差旅相关信息。
- **缓存管理**：支持从 `.invoice_cache/` 目录加载已提取的结构化数据和匹配结果，避免重复处理。

## 对外接口

- `extract_document(file_path) -> dict` — 统一入口：从任意图片/PDF 提取信息
- `extract_travel_info(source_dir) -> dict` — 综合发票和匹配结果提取差旅信息
- `load_cache(source_dir) -> dict` — 加载缓存的结构化数据
- `load_match_result(source_dir) -> dict` — 加载发票与支付记录的匹配结果
"""

from __future__ import annotations

import base64
import json
from pathlib import Path
from typing import Any, cast

from .. import get_logger
from .prompt import (
    build_invoice_system_prompt,
    build_normal_info_system_prompt,
    build_travel_info_system_prompt,
)

log = get_logger("llm_extractor")


def _create_llm() -> Any:
    """根据配置文件创建 LLM 实例。"""
    try:
        from llama_index.llms.openai_like import OpenAILike
    except ImportError:
        log.error("缺少 llama-index-llms-openai-like，请执行: uv pip install llama-index-llms-openai-like")
        raise

    from ..config import get_llm_config

    llm_config = get_llm_config()
    return OpenAILike(
        model=llm_config["model"],
        api_base=llm_config["api_base"],
        api_key=llm_config.get("api_key", "lm-studio"),
        temperature=0.1,
        max_tokens=65535,
        request_timeout=600.0,
        is_chat_model=True,
    )


def _parse_json_response(text: str) -> dict[str, Any]:
    """从 LLM 响应中提取 JSON，处理可能的 Markdown 包裹。"""
    text = text.strip()

    # 处理 ```json ... ``` 包裹
    if "```" in text:
        # 提取第一个代码块
        start = text.find("```") + 3
        end = text.find("```", start)
        if end > start:
            text = text[start:end].strip()

    # 去掉可能的前缀 (如 "json")
    if text.lower().startswith("json"):
        text = text[4:].strip()

    return cast(dict[str, Any], json.loads(text))


# ------------------------------------------------------------------
# 统一文档提取（多模态，直接传图片给 LLM）
# ------------------------------------------------------------------


def _image_to_base64(image_path: Path) -> str:
    """将图片文件读取为 base64 字符串。"""
    with open(image_path, "rb") as f:
        return base64.b64encode(f.read()).decode("utf-8")


def _llm_query_multimodal(
    system_prompt: str,
    text: str | None = None,
    image_b64s: list[str] | None = None,
    blocks: list[Any] | None = None,
    reasoning_effort: str = "none",
) -> str:
    """发送多模态请求到 LLM。

    Args:
        system_prompt: 系统提示词。
        text: 用户文本（与 image_b64s 配合使用，文本在前、图片在后）。
        image_b64s: base64 编码的图片列表。
        blocks: 预构建的内容块列表（TextBlock/ImageBlock），传入时忽略 text 和 image_b64s。

    Returns:
        LLM 响应文本。
    """
    from llama_index.core.base.llms.types import ImageBlock, TextBlock
    from llama_index.core.llms import ChatMessage

    from ..config import get_llm_config

    if blocks is not None:
        final_blocks = blocks
    else:
        text = text or ""
        image_b64s = image_b64s or []
        final_blocks = [TextBlock(text=text)]
        for img_b64 in image_b64s:
            final_blocks.append(
                ImageBlock(
                    url=f"data:image/jpeg;base64,{img_b64}",
                    detail="high",
                )
            )

    messages = [
        ChatMessage(role="system", content=system_prompt),
        ChatMessage(role="user", blocks=final_blocks),
    ]

    llm_config = get_llm_config()
    llm = _create_llm()
    log.info(
        "开始请求 LLM 多模态 (model=%s, base=%s, blocks=%d)",
        llm_config["model"],
        llm_config["api_base"],
        len(final_blocks),
    )

    try:
        parts = []
        for resp in llm.stream_chat(
            messages,
            temperature=0.1,
            extra_body={"reasoning_effort": reasoning_effort},
        ):
            delta = resp.delta
            if delta:
                parts.append(delta)
        text = "".join(parts)
        log.info("LLM 多模态请求完成，响应总长度: %d 字符", len(text))
        log.info("LLM 多模态响应: %s", text)
        return text
    except Exception as e:
        log.error("LLM 多模态请求失败: %s", e)
        raise


def extract_document(file_path: Path) -> dict[str, Any]:
    """统一文档提取入口：从任意图片/PDF 中提取结构化信息。

    LLM 会根据统一提示词自行判断文档类型（发票/支付记录/出差事前申请单等）。

    Args:
        file_path: 文件路径（支持 PDF 和图片格式）。

    Returns:
        包含提取字段的字典。
    """
    from .pdf import render_pdf_to_images

    system_prompt = build_invoice_system_prompt()
    user_text = f"请分析以下财务文档并提取信息：\n\n文件名: {file_path.name}"

    # PDF 先渲染为图片
    suffix = file_path.suffix.lower()
    if suffix == ".pdf":
        image_b64s = render_pdf_to_images(file_path)
    else:
        image_b64s = [_image_to_base64(file_path)]

    if not image_b64s:
        log.warning(f"文件渲染为空: {file_path.name}")
        return {}

    try:
        response = _llm_query_multimodal(system_prompt, user_text, image_b64s)
        result = _parse_json_response(response)
        log.info("LLM 文档提取成功: %s", file_path.name)
        return result
    except Exception as e:
        log.error("LLM 文档提取失败: %s (%s)", file_path.name, e)
        raise


# ------------------------------------------------------------------
# 差旅信息提取
# ------------------------------------------------------------------

CACHE_DIR_NAME = ".invoice_cache"


def load_cache(source_dir: Path) -> dict[str, Any]:
    """从 JSON 缓存目录加载结构化数据，构建 source filename -> 缓存数据的映射。

    Args:
        source_dir: 源文件目录（包含 .invoice_cache 子目录）。

    Returns:
        {source_filename: extracted_data} 字典。
        额外包含 "travel_info" 键（如果 travel_info.json 存在）。
    """
    cache_map: dict[str, Any] = {}
    cache_dir = source_dir / CACHE_DIR_NAME
    if not cache_dir.exists():
        return cache_map

    for json_path in sorted(cache_dir.glob("*.json")):
        try:
            with open(json_path, encoding="utf-8") as f:
                cache_data = json.load(f)

            # travel_info.json / normal_info.json 结构不同，直接存储
            if json_path.name in ("travel_info.json", "normal_info.json"):
                cache_map[json_path.name.replace(".json", "")] = cache_data
                continue

            extracted = cache_data.get("extracted_data", {})
            src_file = extracted.get("_source_file", "")
            if src_file:
                cache_map[src_file] = extracted
        except Exception as e:
            log.warning(f"读取缓存失败 {json_path.name}: {e}")

    return cache_map


def load_match_result(source_dir: Path) -> dict[str, list[dict[str, Any]]]:
    """从 JSON 缓存目录加载发票与支付记录的匹配结果。

    Args:
        source_dir: 源文件目录（包含 .invoice_cache 子目录）。

    Returns:
        {支付记录源文件 (含金额): [发票信息列表]} 字典。
        每个发票信息包含 file, type, amount 字段。
    """
    cache_dir = source_dir / CACHE_DIR_NAME
    match_path = cache_dir / "match_result.json"
    if not match_path.exists():
        return {}

    try:
        with open(match_path, encoding="utf-8") as f:
            result: dict[str, list[dict[str, Any]]] = json.load(f)
            return result
    except Exception as e:
        log.warning(f"读取匹配结果缓存失败: {e}")
        return {}


def extract_travel_info(
    source_dir: Path | None = None,
) -> dict[str, Any]:
    """根据差旅发票（bot 格式），让 LLM 提取出差相关信息。

    仅支持从 JSON 缓存加载数据。

    bot 格式的发票包含以下字段：
        - 发票类型, invoice_no, invoice_date, item_name, spec_model
        - total_amount, seller_name, person_name, person_id
        - card_date, card_no, card_amount, remark

    Args:
        source_dir: 源文件目录（必填，包含 .invoice_cache 子目录）。

    Returns:
        包含出差事由、地点、交通工具、时间、住宿信息等字段的字典。
    """
    # 仅从 JSON 缓存加载结构化数据
    if not source_dir:
        log.warning("未提供 source_dir，无法加载缓存数据")
        return {}

    system_prompt = build_travel_info_system_prompt()

    # 构建 source filename -> 缓存数据的映射
    cache_map = load_cache(source_dir)

    # 加载发票与支付记录的匹配结果
    match_result = load_match_result(source_dir)

    # 拼接纯文本消息
    parts = [
        "以下是本次报销的所有源文件及其提取出的结构化数据。"
        "每个源文件的数据来自 OCR 识别和发票信息提取，已按文件名分组展示。"
    ]

    # 如果有匹配结果，作为额外上下文提供
    if match_result:
        parts.append(
            "【发票与支付记录匹配结果】"
            "以下数据已将发票信息与对应的支付记录进行关联匹配，"
            "用于判断每笔支付对应的发票和商户信息。\n" + json.dumps(match_result, ensure_ascii=False, indent=2)
        )

    # 按源文件名提供结构化数据
    for filename, extracted in cache_map.items():
        parts.append(
            f"【源文件: {filename}】"
            "以下为从该文件提取的结构化发票/支付/申请单数据。\n" + json.dumps(extracted, ensure_ascii=False, indent=2)
        )

    parts.append("\n=== 请返回 JSON 格式结果 ===")
    user_message = "\n".join(parts)
    log.info(f"user_message: {user_message}")
    try:
        response = _llm_query_multimodal(
            system_prompt=system_prompt,
            text=user_message,
            reasoning_effort="low",
        )
        result = _parse_json_response(response)
        log.info("LLM 差旅信息提取成功")
        return result
    except Exception as e:
        log.error("LLM 差旅信息提取失败: %s", e)
        raise


# ------------------------------------------------------------------
# 普通发票信息提取
# ------------------------------------------------------------------


def extract_normal_info(
    source_dir: Path | None = None,
) -> dict[str, Any]:
    """根据普通发票（非差旅），让 LLM 提取报销相关信息。

    仅支持从 JSON 缓存加载数据。

    Args:
        source_dir: 源文件目录（必填，包含 .invoice_cache 子目录）。

    Returns:
        包含报销说明、发票总数、总金额、支付方式、附件清单等字段的字典。
    """
    if not source_dir:
        log.warning("未提供 source_dir，无法加载缓存数据")
        return {}

    system_prompt = build_normal_info_system_prompt()

    # 构建 source filename -> 缓存数据的映射
    cache_map = load_cache(source_dir)

    # 加载发票与支付记录的匹配结果
    match_result = load_match_result(source_dir)

    # 拼接纯文本消息
    parts = [
        "以下是本次报销的所有源文件及其提取出的结构化数据。"
        "每个源文件的数据来自 OCR 识别和发票信息提取，已按文件名分组展示。"
    ]

    # 如果有匹配结果，作为额外上下文提供
    if match_result:
        parts.append(
            "【发票与支付记录匹配结果】"
            "以下数据已将发票信息与对应的支付记录进行关联匹配，"
            "用于判断每笔支付对应的发票和商户信息。\n" + json.dumps(match_result, ensure_ascii=False, indent=2)
        )

    # 按源文件名提供结构化数据
    for filename, extracted in cache_map.items():
        parts.append(
            f"【源文件: {filename}】"
            "以下为从该文件提取的结构化发票/支付/申请单数据。\n" + json.dumps(extracted, ensure_ascii=False, indent=2)
        )

    parts.append("\n=== 请返回 JSON 格式结果 ===")
    user_message = "\n".join(parts)
    log.info(f"user_message: {user_message}")
    try:
        response = _llm_query_multimodal(
            system_prompt=system_prompt,
            text=user_message,
            reasoning_effort="low",
        )
        result = _parse_json_response(response)
        log.info("LLM 普通发票信息提取成功")
        return result
    except Exception as e:
        log.error("LLM 普通发票信息提取失败: %s", e)
        raise