完成差旅发票录入流程
This commit is contained in:
wandering
@@ -1,12 +1,19 @@
|
||||
"""
|
||||
LLM 信息提取
|
||||
"""LLM 信息提取
|
||||
|
||||
使用 LLM 从 PDF 文本中提取结构化数据，以及从支付截图中提取刷卡信息。
|
||||
支持 JSON 格式输出，字段与 CSV_COLUMNS 对齐。
|
||||
使用 LLM 从 PDF 文本/图片、支付截图中提取结构化数据。
|
||||
|
||||
对外接口:
|
||||
extract_invoice_from_text(text, file_name) -> dict 从 PDF 文本提取发票信息
|
||||
extract_card_info_from_image(image_path) -> dict 从支付截图提取刷卡信息
|
||||
## 功能模块
|
||||
|
||||
- **统一文档提取**：使用一套提示词，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
|
||||
@@ -17,7 +24,10 @@ from pathlib import Path
|
||||
from typing import Any, cast
|
||||
|
||||
from .. import get_logger
|
||||
from .prompt import build_card_info_system_prompt, build_invoice_system_prompt
|
||||
from .prompt import (
|
||||
build_invoice_system_prompt,
|
||||
build_travel_info_system_prompt,
|
||||
)
|
||||
|
||||
log = get_logger("llm_extractor")
|
||||
|
||||
@@ -38,47 +48,12 @@ def _create_llm() -> Any:
|
||||
api_base=llm_config["api_base"],
|
||||
api_key=llm_config.get("api_key", "lm-studio"),
|
||||
temperature=0.1,
|
||||
max_tokens=8192,
|
||||
max_tokens=65535,
|
||||
request_timeout=600.0,
|
||||
is_chat_model=True,
|
||||
)
|
||||
|
||||
|
||||
def _llm_query(system_prompt: str, user_content: str, max_tokens: int = 4096) -> str:
|
||||
"""发送请求到 LLM 并返回完整响应文本。"""
|
||||
from llama_index.core.llms import ChatMessage
|
||||
|
||||
from ..config import get_llm_config
|
||||
|
||||
messages = [
|
||||
ChatMessage(role="system", content=system_prompt),
|
||||
ChatMessage(role="user", content=user_content),
|
||||
]
|
||||
|
||||
llm_config = get_llm_config()
|
||||
llm = _create_llm()
|
||||
log.info("开始请求 LLM (model=%s, base=%s)", llm_config["model"], llm_config["api_base"])
|
||||
|
||||
try:
|
||||
parts = []
|
||||
for resp in llm.stream_chat(
|
||||
messages,
|
||||
temperature=0.1,
|
||||
max_tokens=max_tokens,
|
||||
extra_body={"reasoning_effort": "none"},
|
||||
):
|
||||
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 _parse_json_response(text: str) -> dict[str, Any]:
|
||||
"""从 LLM 响应中提取 JSON，处理可能的 Markdown 包裹。"""
|
||||
text = text.strip()
|
||||
@@ -98,31 +73,8 @@ def _parse_json_response(text: str) -> dict[str, Any]:
|
||||
return cast(dict[str, Any], json.loads(text))
|
||||
|
||||
|
||||
def extract_invoice_from_text(text: str, file_name: str = "") -> dict[str, Any]:
|
||||
"""从 PDF 发票文本中提取结构化数据。
|
||||
|
||||
Args:
|
||||
text: PDF 提取的文本内容。
|
||||
file_name: 原始文件名（用于日志）。
|
||||
|
||||
Returns:
|
||||
包含所有 CSV_COLUMNS 字段的字典。
|
||||
"""
|
||||
system_prompt = build_invoice_system_prompt()
|
||||
user_content = f"请分析以下发票文本并提取信息：\n\n文件名: {file_name}\n\n---\n\n{text}\n\n---"
|
||||
|
||||
try:
|
||||
response = _llm_query(system_prompt, user_content, max_tokens=4096)
|
||||
result = _parse_json_response(response)
|
||||
log.info("LLM 发票提取成功: %s", file_name)
|
||||
return result
|
||||
except Exception as e:
|
||||
log.error("LLM 发票提取失败: %s (%s)", file_name, e)
|
||||
raise
|
||||
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 支付截图信息提取（多模态）
|
||||
# 统一文档提取（多模态，直接传图片给 LLM）
|
||||
# ------------------------------------------------------------------
|
||||
|
||||
|
||||
@@ -134,36 +86,53 @@ def _image_to_base64(image_path: Path) -> str:
|
||||
|
||||
def _llm_query_multimodal(
|
||||
system_prompt: str,
|
||||
text: str,
|
||||
image_b64: str,
|
||||
max_tokens: int = 4096,
|
||||
text: str | None = None,
|
||||
image_b64s: list[str] | None = None,
|
||||
blocks: list[Any] | None = None,
|
||||
reasoning_effort: str = "none",
|
||||
) -> str:
|
||||
"""发送多模态请求（文本 + 图片）到 LLM。"""
|
||||
"""发送多模态请求到 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=[
|
||||
TextBlock(text=text),
|
||||
ImageBlock(
|
||||
url=f"data:image/jpeg;base64,{image_b64}",
|
||||
detail="high",
|
||||
),
|
||||
],
|
||||
),
|
||||
ChatMessage(role="user", blocks=final_blocks),
|
||||
]
|
||||
|
||||
llm_config = get_llm_config()
|
||||
llm = _create_llm()
|
||||
log.info(
|
||||
"开始请求 LLM 多模态 (model=%s, base=%s)",
|
||||
"开始请求 LLM 多模态 (model=%s, base=%s, blocks=%d)",
|
||||
llm_config["model"],
|
||||
llm_config["api_base"],
|
||||
len(final_blocks),
|
||||
)
|
||||
|
||||
try:
|
||||
@@ -171,8 +140,7 @@ def _llm_query_multimodal(
|
||||
for resp in llm.stream_chat(
|
||||
messages,
|
||||
temperature=0.1,
|
||||
max_tokens=max_tokens,
|
||||
extra_body={"reasoning_effort": "none"},
|
||||
extra_body={"reasoning_effort": reasoning_effort},
|
||||
):
|
||||
delta = resp.delta
|
||||
if delta:
|
||||
@@ -186,25 +154,173 @@ def _llm_query_multimodal(
|
||||
raise
|
||||
|
||||
|
||||
def extract_card_info_from_image(image_path: Path) -> dict[str, Any]:
|
||||
"""从支付截图中提取刷卡信息。
|
||||
def extract_document(file_path: Path) -> dict[str, Any]:
|
||||
"""统一文档提取入口：从任意图片/PDF 中提取结构化信息。
|
||||
|
||||
LLM 会根据统一提示词自行判断文档类型（发票/支付记录/出差事前申请单等）。
|
||||
|
||||
Args:
|
||||
image_path: 支付截图图片路径。
|
||||
file_path: 文件路径（支持 PDF 和图片格式）。
|
||||
|
||||
Returns:
|
||||
包含刷卡日期、刷卡金额、公务卡号的字典。
|
||||
包含提取字段的字典。
|
||||
"""
|
||||
system_prompt = build_card_info_system_prompt()
|
||||
user_text = f"请分析以下支付截图并提取信息：\n\n文件名: {image_path.name}"
|
||||
from .pdf import render_pdf_to_images
|
||||
|
||||
image_b64 = _image_to_base64(image_path)
|
||||
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_b64, max_tokens=4096)
|
||||
response = _llm_query_multimodal(system_prompt, user_text, image_b64s)
|
||||
result = _parse_json_response(response)
|
||||
log.info("LLM 支付截图提取成功: %s", image_path.name)
|
||||
log.info("LLM 文档提取成功: %s", file_path.name)
|
||||
return result
|
||||
except Exception as e:
|
||||
log.error("LLM 支付截图提取失败: %s (%s)", image_path.name, 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 结构不同，直接存储
|
||||
if json_path.name == "travel_info.json":
|
||||
cache_map["travel_info"] = 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
|
||||
|
||||
Reference in New Issue
Block a user