610 lines
20 KiB
Python
610 lines
20 KiB
Python
"""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 .invoice import CACHE_DIR_NAME
|
||
from .prompt import (
|
||
build_invoice_system_prompt,
|
||
build_normal_info_system_prompt,
|
||
build_supplement_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 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"],
|
||
)
|
||
|
||
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("LLM 请求完成，响应总长度: %d 字符", len(text))
|
||
if source_dir:
|
||
_emit_llm_stream(source_dir, "end", label="分析完成")
|
||
return text
|
||
except Exception as e:
|
||
log.error("LLM 请求失败: %s", e)
|
||
if source_dir:
|
||
_emit_llm_stream(source_dir, "error", error=str(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
|
||
|
||
|
||
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),
|
||
)
|
||
|
||
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("LLM 多模态请求完成，响应总长度: %d 字符", len(text))
|
||
log.info("LLM 多模态响应: %s", text)
|
||
if source_dir:
|
||
_emit_llm_stream(source_dir, "end", label="分析完成")
|
||
return text
|
||
except Exception as e:
|
||
log.error("LLM 多模态请求失败: %s", e)
|
||
if source_dir:
|
||
_emit_llm_stream(source_dir, "error", error=str(e))
|
||
raise
|
||
|
||
|
||
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,
|
||
) -> str:
|
||
"""构建提取请求的用户消息（供 Agent 调度使用）。
|
||
|
||
Args:
|
||
cache_map: 缓存数据映射。
|
||
match_result: 匹配结果。
|
||
previous_analysis: 上一轮 LLM 分析结果（可选，补充文件时传入作为历史上下文）。
|
||
|
||
Returns:
|
||
拼接好的用户消息字符串。
|
||
"""
|
||
parts = [
|
||
"以下是本次报销的所有源文件及其提取出的结构化数据。"
|
||
"每个源文件的数据来自 OCR 识别和发票信息提取，已按文件名分组展示。"
|
||
]
|
||
|
||
if previous_analysis:
|
||
parts.append(
|
||
"【上一轮分析结果】"
|
||
"以下是上一轮 LLM 对已有文件的分析结果。"
|
||
"注意：用户可能已补充新文件，请综合所有数据（含新文件）重新分析。"
|
||
"如果新文件填补了之前的信息缺失，请相应更新分析结果。\n"
|
||
+ json.dumps(previous_analysis, ensure_ascii=False, indent=2)
|
||
)
|
||
|
||
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_travel_info(
|
||
source_dir: Path | None = None,
|
||
) -> dict[str, Any]:
|
||
"""根据差旅发票（bot 格式），让 LLM 提取出差相关信息。
|
||
|
||
纯提取，不包含校验逻辑。校验由 Agent 层调度。
|
||
|
||
Args:
|
||
source_dir: 源文件目录（必填，包含 .invoice_cache 子目录）。
|
||
|
||
Returns:
|
||
包含出差事由、地点、交通工具、时间、住宿信息等字段的字典。
|
||
"""
|
||
if not source_dir:
|
||
log.warning("未提供 source_dir，无法加载缓存数据")
|
||
return {}
|
||
|
||
system_prompt = build_travel_info_system_prompt()
|
||
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(f"user_message: {user_message}")
|
||
|
||
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 差旅信息提取成功")
|
||
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 提取报销相关信息。
|
||
|
||
纯提取，不包含校验逻辑。校验由 Agent 层调度。
|
||
|
||
Args:
|
||
source_dir: 源文件目录（必填，包含 .invoice_cache 子目录）。
|
||
|
||
Returns:
|
||
包含报销说明、发票总数、总金额、支付方式、附件清单等字段的字典。
|
||
"""
|
||
if not source_dir:
|
||
log.warning("未提供 source_dir，无法加载缓存数据")
|
||
return {}
|
||
|
||
system_prompt = build_normal_info_system_prompt()
|
||
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(f"user_message: {user_message}")
|
||
|
||
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 普通发票信息提取成功")
|
||
return result
|
||
except Exception as e:
|
||
log.error("LLM 普通发票信息提取失败: %s", e)
|
||
raise
|
||
|
||
|
||
# ------------------------------------------------------------------
|
||
# 用户补充信息处理
|
||
# ------------------------------------------------------------------
|
||
|
||
|
||
def process_user_supplement(
|
||
user_text: str,
|
||
extracted_info: dict[str, Any],
|
||
invoice_type: str,
|
||
source_dir: Path | None = None,
|
||
) -> dict[str, Any]:
|
||
"""让用户补充的文字信息通过 LLM 分析，返回需要更新的字段。
|
||
|
||
Args:
|
||
user_text: 用户输入的文字。
|
||
extracted_info: 当前已提取的报销信息。
|
||
invoice_type: "travel" 或 "normal"。
|
||
source_dir: 会话目录（可选，传入时启用 SSE 流式事件写入）。
|
||
|
||
Returns:
|
||
包含 updated_fields, changes, confidence, unparsed_info 的字典。
|
||
"""
|
||
system_prompt = build_supplement_system_prompt()
|
||
|
||
parts = [
|
||
f"发票类型: {'差旅报销' if invoice_type == 'travel' else '普通报销'}",
|
||
"",
|
||
"以下是当前已提取的报销信息：",
|
||
json.dumps(extracted_info, ensure_ascii=False, indent=2),
|
||
"",
|
||
f"用户补充信息：{user_text}",
|
||
"",
|
||
"=== 请分析用户输入并返回需要更新的字段 ===",
|
||
]
|
||
user_message = "\n".join(parts)
|
||
|
||
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 补充信息分析完成")
|
||
|
||
return result
|
||
except Exception as e:
|
||
log.error("LLM 补充信息分析失败: %s", e)
|
||
return {
|
||
"updated_fields": {},
|
||
"changes": [],
|
||
"confidence": 0.0,
|
||
"unparsed_info": f"分析失败: {e}",
|
||
}
|
||
|
||
|
||
def merge_supplement_into_info(
|
||
extracted_info: dict[str, Any],
|
||
updated_fields: dict[str, Any],
|
||
) -> dict[str, Any]:
|
||
"""将 LLM 返回的更新字段合并到已提取的信息中。
|
||
|
||
支持点号路径（如 basic_info.travel_purpose）表示嵌套更新。
|
||
|
||
Args:
|
||
extracted_info: 当前已提取的报销信息。
|
||
updated_fields: LLM 返回的需要更新的字段。
|
||
|
||
Returns:
|
||
更新后的报销信息。
|
||
"""
|
||
import copy
|
||
|
||
result = copy.deepcopy(extracted_info)
|
||
|
||
for field_path, value in updated_fields.items():
|
||
parts = field_path.split(".")
|
||
current = result
|
||
for part in parts[:-1]:
|
||
if part not in current:
|
||
current[part] = {}
|
||
current = current[part]
|
||
current[parts[-1]] = value
|
||
log.info("更新字段 %s = %s", field_path, value)
|
||
|
||
return result
|