Auto-Finance/src/core/extraction/llm_extractor.py

"""LLM 信息提取

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

## 功能模块

- **统一文档提取**：使用一套提示词，LLM 自行判断文档类型（发票/支付记录/出差事前申请单等），支持 JSON 格式输出。
- **差旅信息提取**：综合多张发票、支付记录和匹配结果，提取出差事由、地点、时间等差旅相关信息。
- **缓存管理**：支持从 `.invoice_cache/` 目录加载已提取的结构化数据和匹配结果，避免重复处理。
- **SSE 流式事件**：`extract_travel_info` 和 `extract_normal_info` 在调用 LLM 时，向 `source_dir/llm_stream.log` 写入流式事件（start/reasoning/chunk/end/error），前端通过 SSE 实时展示 AI 思考过程与正式回答。

## 对外接口

- `extract_document(file_path) -> dict` — 统一入口：从任意图片/PDF 提取信息
- `extract_travel_info(source_dir) -> dict` — 综合发票和匹配结果提取差旅信息
- `extract_normal_info(source_dir) -> dict` — 提取普通发票报销信息
- `load_cache(source_dir) -> dict` — 加载缓存的结构化数据
- `load_match_result(source_dir) -> dict` — 加载发票与支付记录的匹配结果
- `llm_query_text(system_prompt, text, source_dir) -> str` — 纯文本 LLM 查询（供 Agent 调度使用）
- `parse_json_response(text) -> dict` — 从 LLM 响应中提取 JSON（供 Agent 调度使用）
- `build_extraction_user_message(cache_map, match_result) -> str` — 构建提取请求的用户消息（供 Agent 调度使用）

## SSE 流式事件协议

`llm_stream.log` 每行一个 JSON 对象：

- `{"type": "llm_stream", "phase": "start", "label": "..."}` — LLM 调用开始
- `{"type": "llm_stream", "phase": "reasoning", "text": "..."}` — 模型原生推理/思考片段（来自 `thinking_delta`）
- `{"type": "llm_stream", "phase": "chunk", "text": "..."}` — 流式文本片段（正式回答）
- `{"type": "llm_stream", "phase": "end", "label": "..."}` — LLM 调用结束
- `{"type": "llm_stream", "phase": "error", "error": "..."}` — LLM 调用失败
"""

from __future__ import annotations

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

from ... import get_logger
from ...infra.documents.invoice import CACHE_DIR_NAME
from ...infra.llm.prompt import (
    build_invoice_system_prompt,
    build_normal_info_system_prompt,
    build_travel_info_system_prompt,
)

log = get_logger("llm_extractor")

# Re-export CACHE_DIR_NAME for convenience
__all__ = [
    "CACHE_DIR_NAME",
    "extract_document",
    "extract_travel_info",
    "extract_normal_info",
    "load_cache",
    "load_match_result",
    "llm_query_text",
    "parse_json_response",
    "build_extraction_user_message",
]

# SSE LLM 流式事件日志文件名
LLM_STREAM_LOG = "llm_stream.log"


def _emit_llm_stream(source_dir: Path, phase: str, **kwargs: Any) -> None:
    """向 llm_stream.log 追加一行 JSON 事件（线程安全，失败时静默忽略）

    Args:
        source_dir: 会话目录路径。
        phase: 事件阶段 ("start" / "chunk" / "end" / "error")。
        **kwargs: 额外字段 (text, label, error 等)。
    """
    event = {"type": "llm_stream", "phase": phase, **kwargs}
    try:
        event_path = source_dir / LLM_STREAM_LOG
        with open(event_path, "a", encoding="utf-8") as f:
            f.write(json.dumps(event, ensure_ascii=False) + "\n")
    except Exception:
        pass


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 包裹。

    Args:
        text: LLM 响应文本。

    Returns:
        解析后的字典。
    """
    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 _stream_llm_response(
    llm: Any,
    messages: list[Any],
    source_dir: Path | None,
    reasoning_effort: str,
    log_label: str = "LLM",
) -> str:
    """流式调用 LLM 并写入 SSE 事件（供 llm_query_text 和 _llm_query_multimodal 共用）。

    Args:
        llm: LLM 实例。
        messages: 消息列表。
        source_dir: 会话目录（可选，传入时启用 SSE 流式事件写入）。
        reasoning_effort: 推理努力级别。
        log_label: 日志标签（用于区分"纯文本"和"多模态"）。

    Returns:
        LLM 响应文本。
    """
    try:
        if source_dir:
            _emit_llm_stream(source_dir, "start", label="正在分析文件...")

        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)
                if source_dir:
                    _emit_llm_stream(source_dir, "chunk", text=delta)

            thinking = getattr(resp, "additional_kwargs", {}) or {}
            thinking_delta = thinking.get("thinking_delta", "")
            if thinking_delta and source_dir:
                _emit_llm_stream(source_dir, "reasoning", text=thinking_delta)

        text = "".join(parts)
        log.info("%s请求完成，响应总长度: %d 字符", log_label, len(text))
        if source_dir:
            _emit_llm_stream(source_dir, "end", label="分析完成")
        return text
    except Exception as e:
        log.error("%s请求失败: %s", log_label, e)
        if source_dir:
            _emit_llm_stream(source_dir, "error", error=str(e))
        raise


def llm_query_text(
    system_prompt: str,
    text: str,
    reasoning_effort: str = "none",
    source_dir: Path | None = None,
) -> str:
    """发送纯文本请求到 LLM（供 Agent 调度使用）。

    Args:
        system_prompt: 系统提示词。
        text: 用户文本。
        reasoning_effort: 推理努力级别。
        source_dir: 会话目录（可选，传入时启用 SSE 流式事件写入）。

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

    from ...config import get_llm_config

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

    llm_config = get_llm_config()
    llm = _create_llm()
    log.info(
        "开始请求 LLM (model=%s, base=%s)",
        llm_config["model"],
        llm_config["api_base"],
    )

    return _stream_llm_response(llm, messages, source_dir, reasoning_effort, log_label="LLM")


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

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

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

    Returns:
        包含提取字段的字典。
    """
    from ...infra.documents.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


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",
    source_dir: Path | None = None,
) -> str:
    """发送多模态请求到 LLM（内部使用）。

    Args:
        system_prompt: 系统提示词。
        text: 用户文本（与 image_b64s 配合使用，文本在前、图片在后）。
        image_b64s: base64 编码的图片列表。
        blocks: 预构建的内容块列表（TextBlock/ImageBlock），传入时忽略 text 和 image_b64s。
        source_dir: 会话目录（可选，传入时启用 SSE 流式事件写入）。

    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),
    )

    return _stream_llm_response(llm, messages, source_dir, reasoning_effort, log_label="LLM多模态")


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)

            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 build_extraction_user_message(
    cache_map: dict[str, Any],
    match_result: dict[str, list[dict[str, Any]]],
    previous_analysis: dict[str, Any] | None = None,
    user_supplement: str | None = None,
) -> str:
    """构建提取请求的用户消息（供 Agent 调度使用）。

    Args:
        cache_map: 缓存数据映射。
        match_result: 匹配结果。
        previous_analysis: 上一轮 LLM 分析结果（可选，补充文件时传入作为历史上下文）。
        user_supplement: 用户文字补充（可选，传入时作为修正上下文）。

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

    if previous_analysis:
        parts.append(
            "【上一轮分析结果】"
            "以下是上一轮 LLM 对已有文件的分析结果。"
            "注意：用户可能已补充新文件，请综合所有数据（含新文件）重新分析。"
            "如果新文件填补了之前的信息缺失，请相应更新分析结果。\n"
            + json.dumps(previous_analysis, ensure_ascii=False, indent=2)
        )

    if user_supplement:
        parts.append("【用户文字补充】用户补充了以下文字信息，请据此更新你的分析结果。\n" + user_supplement)

    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 格式结果 ===")
    return "\n".join(parts)


def _extract_info(
    source_dir: Path | None,
    system_prompt: str,
    info_type: str,
) -> dict[str, Any]:
    """通用的信息提取函数：加载缓存、构建消息、调用 LLM 并解析 JSON。

    extract_travel_info 和 extract_normal_info 的公共实现。

    Args:
        source_dir: 源文件目录（必填，包含 .invoice_cache 子目录）。
        system_prompt: 系统提示词。
        info_type: 信息类型标签（"差旅" 或 "普通发票"），用于日志。

    Returns:
        LLM 提取的结构化信息字典。
    """
    if not source_dir:
        log.warning("未提供 source_dir，无法加载缓存数据")
        return {}

    cache_map = load_cache(source_dir)
    match_result = load_match_result(source_dir)
    user_message = build_extraction_user_message(cache_map, match_result)

    log.info("开始构建%s信息提取请求，缓存条目: %d, 匹配结果: %d", info_type, len(cache_map), len(match_result))

    try:
        response = llm_query_text(
            system_prompt=system_prompt,
            text=user_message,
            reasoning_effort="low",
            source_dir=source_dir,
        )
        result = parse_json_response(response)
        log.info("LLM %s信息提取成功", info_type)
        return result
    except Exception as e:
        log.error("LLM %s信息提取失败: %s", info_type, e)
        raise


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

    纯提取，不包含校验逻辑。校验由 Agent 层调度。

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

    Returns:
        包含出差事由、地点、交通工具、时间、住宿信息等字段的字典。
    """
    return _extract_info(source_dir, build_travel_info_system_prompt(), "差旅")


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

    纯提取，不包含校验逻辑。校验由 Agent 层调度。

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

    Returns:
        包含报销说明、发票总数、总金额、支付方式、附件清单等字段的字典。
    """
    return _extract_info(source_dir, build_normal_info_system_prompt(), "普通发票")