wandering
370 lines
14 KiB
Python
370 lines
14 KiB
Python
"""发票提取编排
|
||
|
||
统一扫描目录下所有文件（PDF + 图片），通过 LLM 提取结构化数据，
|
||
根据 LLM 返回的「invoice_type」字段自动分类为发票/支付记录/出差事前申请单。
|
||
|
||
对外接口:
|
||
extract_invoices(directory) -> tuple[list[dict], list[dict], dict]
|
||
|
||
错误传播规则:
|
||
- 单个文件提取失败: 记录日志 + SSE error 事件，继续处理下一个文件
|
||
- 全部文件提取失败: raise ExtractionError，携带失败文件列表和原始错误
|
||
- 缓存读取失败: 静默降级，尝试重新提取
|
||
|
||
SSE 文件进度事件:
|
||
在 source_dir 下写入 file_events.log，每行一个 JSON 对象：
|
||
- {"type": "file_progress", "file": "...", "status": "processing"}
|
||
- {"type": "file_progress", "file": "...", "status": "done", "summary": {...}}
|
||
- {"type": "file_progress", "file": "...", "status": "cached"}
|
||
- {"type": "file_progress", "file": "...", "status": "error", "error": "..."}
|
||
"""
|
||
|
||
import json
|
||
import os
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from ... import get_logger
|
||
from ...core.matching import match_invoices_to_cards
|
||
from ...exceptions import ExtractionError
|
||
from ...infra.documents.invoice import CACHE_DIR_NAME, classify_invoice_batch
|
||
from .llm_extractor import extract_document
|
||
|
||
log = get_logger("extractor")
|
||
|
||
# SSE 文件进度事件日志文件名
|
||
FILE_EVENTS_LOG = "file_events.log"
|
||
|
||
# 并行提取文件数，可通过环境变量 EXTRACTION_PARALLEL_COUNT 配置
|
||
EXTRACTION_PARALLEL_COUNT = int(os.environ.get("EXTRACTION_PARALLEL_COUNT", "3"))
|
||
|
||
# 支持的文件扩展名
|
||
SUPPORTED_EXTENSIONS = {".pdf", ".png", ".jpg", ".jpeg", ".bmp", ".webp"}
|
||
|
||
|
||
def _emit_file_event(source_dir: Path, file_name: str, status: str, **kwargs: Any) -> None:
|
||
"""向 file_events.log 追加一行 JSON 事件（线程安全，失败时静默忽略）"""
|
||
event = {
|
||
"type": "file_progress",
|
||
"file": file_name,
|
||
"status": status,
|
||
**kwargs,
|
||
}
|
||
try:
|
||
event_path = source_dir / FILE_EVENTS_LOG
|
||
with open(event_path, "a", encoding="utf-8") as f:
|
||
f.write(json.dumps(event, ensure_ascii=False) + "\n")
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _get_cache_dir(source_dir: Path) -> Path:
|
||
"""获取缓存目录路径"""
|
||
cache_dir = source_dir / CACHE_DIR_NAME
|
||
cache_dir.mkdir(exist_ok=True)
|
||
return cache_dir
|
||
|
||
|
||
def _get_json_path(file_path: Path, cache_dir: Path) -> Path:
|
||
"""根据文件路径生成对应的 JSON 缓存路径（包含后缀名以区分同名的 PDF/图片）"""
|
||
return cache_dir / f"{file_path.stem}{file_path.suffix}.json"
|
||
|
||
|
||
def _save_to_cache(file_path: Path, extracted_data: dict[str, Any], cache_dir: Path) -> Path:
|
||
"""将提取结果保存到 JSON 缓存文件，并记录源文件路径和后缀名"""
|
||
json_path = _get_json_path(file_path, cache_dir)
|
||
cache_data = {
|
||
"source_file": str(file_path),
|
||
"source_filename": file_path.name,
|
||
"source_extension": file_path.suffix.lower(),
|
||
"extracted_data": extracted_data,
|
||
}
|
||
with open(json_path, "w", encoding="utf-8") as f:
|
||
json.dump(cache_data, f, ensure_ascii=False, indent=2)
|
||
log.info(f"提取结果已缓存: {json_path.name}")
|
||
return json_path
|
||
|
||
|
||
def _load_from_cache(json_path: Path, expected_extension: str | None = None) -> dict[str, Any] | None:
|
||
"""从 JSON 缓存文件加载提取结果，可选校验后缀名一致性"""
|
||
if not json_path.exists():
|
||
return None
|
||
try:
|
||
with open(json_path, encoding="utf-8") as f:
|
||
cache_data: dict[str, Any] = json.load(f)
|
||
# 校验后缀名是否一致，防止同名不同后缀的文件误命中缓存
|
||
if expected_extension and cache_data.get("source_extension", "").lower() != expected_extension.lower():
|
||
return None
|
||
result: dict[str, Any] | None = cache_data.get("extracted_data")
|
||
return result
|
||
except Exception as e:
|
||
log.warning(f"读取缓存失败 {json_path.name}: {e}")
|
||
return None
|
||
|
||
|
||
def _build_summary(data: dict[str, Any]) -> dict[str, str]:
|
||
"""从提取结果构建前端展示摘要（人类可读格式）。
|
||
|
||
返回所有非内部字段（排除 _source_file 等下划线前缀字段），
|
||
并将 invoice_type 转换为中文标签。列表/字典类型的值会展开为可读文本。
|
||
"""
|
||
invoice_type_map = {
|
||
"train": "高铁票",
|
||
"hotel": "酒店住宿",
|
||
"general": "普通发票",
|
||
"payment": "支付记录",
|
||
"application": "出差申请单",
|
||
"attachments": "其他附件",
|
||
"note": "易耗品单",
|
||
}
|
||
|
||
summary = {}
|
||
for key, value in data.items():
|
||
# 跳过内部字段
|
||
if key.startswith("_"):
|
||
continue
|
||
# 跳过空值
|
||
if value is None or value == "":
|
||
continue
|
||
# invoice_type 转为中文标签
|
||
if key == "invoice_type":
|
||
summary["invoice_type_label"] = invoice_type_map.get(str(value), str(value))
|
||
elif isinstance(value, list):
|
||
# 列表展开为多行可读文本
|
||
if len(value) == 0:
|
||
continue
|
||
parts = []
|
||
for item in value:
|
||
if isinstance(item, dict):
|
||
# 字典项用 "key: value" 格式拼接
|
||
pair_parts = [f"{k}: {v}" for k, v in item.items()]
|
||
parts.append(" | ".join(pair_parts))
|
||
else:
|
||
parts.append(str(item))
|
||
summary[key] = "\n".join(parts)
|
||
elif isinstance(value, dict):
|
||
# 字典展开为 "key: value" 格式
|
||
pair_parts = [f"{k}: {v}" for k, v in value.items()]
|
||
summary[key] = " | ".join(pair_parts)
|
||
else:
|
||
summary[key] = str(value)
|
||
|
||
return summary
|
||
|
||
|
||
def _extract_document(
|
||
file_path: Path,
|
||
cache_dir: Path,
|
||
source_dir: Path,
|
||
) -> tuple[dict[str, str] | None, str | None]:
|
||
"""提取单个文件的结构化信息，优先使用缓存。
|
||
|
||
根据 LLM 返回的「invoice_type」字段自动分类：
|
||
- "payment" -> 支付记录
|
||
- "application" -> 申请单
|
||
- 有 "invoice_number" -> 发票
|
||
- 其他 -> 无法识别
|
||
|
||
Args:
|
||
file_path: 文件路径（PDF 或图片）。
|
||
cache_dir: JSON 缓存目录。
|
||
source_dir: 源目录（用于写入 SSE 进度事件）。
|
||
|
||
Returns:
|
||
(提取结果字典或 None, 错误信息或 None)。
|
||
"""
|
||
file_name = file_path.name
|
||
|
||
json_path = _get_json_path(file_path, cache_dir)
|
||
cached = _load_from_cache(json_path, expected_extension=file_path.suffix.lower())
|
||
if cached:
|
||
cached["_source_file"] = file_name
|
||
log.info(f"使用缓存: {file_name}")
|
||
_emit_file_event(source_dir, file_name, "cached")
|
||
return cached, None
|
||
|
||
# 发送处理中事件
|
||
_emit_file_event(source_dir, file_name, "processing")
|
||
|
||
log.info(f"使用多模态提取: {file_name}")
|
||
try:
|
||
result = extract_document(file_path)
|
||
if result:
|
||
result["_source_file"] = file_name
|
||
_save_to_cache(file_path, result, cache_dir)
|
||
|
||
# 发送完成事件（含摘要）
|
||
summary = _build_summary(result)
|
||
_emit_file_event(source_dir, file_name, "done", summary=summary)
|
||
|
||
return result, None
|
||
except Exception as e:
|
||
err_msg = str(e)
|
||
log.warning(f"多模态提取失败: {file_name} ({err_msg})")
|
||
_emit_file_event(source_dir, file_name, "error", error=err_msg)
|
||
return None, err_msg
|
||
|
||
return None, None
|
||
|
||
|
||
def _find_all_files(directory: str) -> list[Path]:
|
||
"""扫描目录下所有支持的文件（PDF + 图片）"""
|
||
dir_path = Path(directory)
|
||
files = [f for f in dir_path.iterdir() if f.is_file() and f.suffix.lower() in SUPPORTED_EXTENSIONS]
|
||
return sorted(files)
|
||
|
||
|
||
def _save_match_result(payment_records: list[dict[str, Any]], cache_dir: Path) -> None:
|
||
"""将发票与支付记录的匹配结果保存到缓存。"""
|
||
match_data: dict[str, list[dict[str, Any]]] = {}
|
||
for record in payment_records:
|
||
card_source = record.get("_source_file", "")
|
||
matched = record.get("_matched_invoices", [])
|
||
card_amount = record.get("card_amount", "")
|
||
|
||
if matched:
|
||
invoice_list = [
|
||
{
|
||
"file": inv.get("_source_file", ""),
|
||
"type": inv.get("invoice_type", ""),
|
||
"amount": inv.get("total_amount", inv.get("total_amount", "")),
|
||
}
|
||
for inv in matched
|
||
]
|
||
|
||
if card_source:
|
||
match_data[f"{card_source} (¥{card_amount})"] = invoice_list
|
||
else:
|
||
key = "__unmatched__"
|
||
if key not in match_data:
|
||
match_data[key] = []
|
||
match_data[key].extend(invoice_list)
|
||
|
||
if not match_data:
|
||
return
|
||
|
||
match_path = cache_dir / "match_result.json"
|
||
with open(match_path, "w", encoding="utf-8") as f:
|
||
json.dump(match_data, f, ensure_ascii=False, indent=2)
|
||
log.info(f"匹配结果已缓存: {match_path.name}")
|
||
|
||
|
||
def extract_invoices(
|
||
directory: str = ".",
|
||
) -> tuple[list[dict[str, str]], list[dict[str, str]], dict[str, list[dict[str, str]]]]:
|
||
"""扫描目录下所有文件，提取信息并匹配支付记录
|
||
|
||
统一使用 LLM 提取，根据返回的「invoice_type」自动分类：
|
||
- 发票（有 invoice_number）-> 参与金额匹配
|
||
- 支付记录（invoice_type="payment"）-> 参与金额匹配
|
||
- 出差事前申请单 -> 单独存储，不参与匹配
|
||
|
||
Returns:
|
||
元组 (payment_records, applications, groups):
|
||
- payment_records: 支付记录列表（每项含关联发票 _matched_invoices）
|
||
- applications: 出差事前申请单列表
|
||
- groups: 按文档类型分组的字典
|
||
{'travel': [差旅发票], 'general': [普通发票], 'application': [出差事前申请单],
|
||
'attachments': [其他附件]}
|
||
|
||
Raises:
|
||
ExtractionError: 当所有文件提取均失败时抛出，携带失败文件列表和原始错误。
|
||
"""
|
||
source_dir = Path(directory)
|
||
cache_dir = _get_cache_dir(source_dir)
|
||
|
||
# 清空上次的文件进度事件
|
||
try:
|
||
(source_dir / FILE_EVENTS_LOG).unlink(missing_ok=True)
|
||
except Exception:
|
||
pass
|
||
|
||
all_files = _find_all_files(directory)
|
||
if not all_files:
|
||
log.warning("未找到支持的文件")
|
||
return [], [], {"travel": [], "general": [], "application": [], "attachments": []}
|
||
|
||
log.info(f"发现 {len(all_files)} 个文件")
|
||
|
||
all_invoices = []
|
||
all_cards = []
|
||
applications = []
|
||
attachments = []
|
||
# 记录失败文件及其错误信息
|
||
failed_files: list[tuple[str, str]] = []
|
||
|
||
max_workers = max(1, EXTRACTION_PARALLEL_COUNT)
|
||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||
future_to_file = {executor.submit(_extract_document, fp, cache_dir, source_dir): fp for fp in all_files}
|
||
|
||
for future in as_completed(future_to_file):
|
||
file_path = future_to_file[future]
|
||
try:
|
||
result, err = future.result()
|
||
except Exception as e:
|
||
err_msg = str(e)
|
||
failed_files.append((file_path.name, err_msg))
|
||
log.warning(f"未能解析: {file_path.name} ({err_msg})")
|
||
continue
|
||
|
||
if not result:
|
||
if err:
|
||
failed_files.append((file_path.name, err))
|
||
log.warning(f"未能解析: {file_path.name}")
|
||
continue
|
||
|
||
inv_type = result.get("invoice_type", "")
|
||
|
||
if inv_type == "application":
|
||
applications.append(result)
|
||
log.info(f"[{inv_type}] 已解析: {file_path.name}")
|
||
elif inv_type == "payment":
|
||
all_cards.append(result)
|
||
log.info(f"[{inv_type}] 已解析: {file_path.name}")
|
||
elif inv_type in ("attachments", "note"):
|
||
attachments.append(result)
|
||
log.info(f"[{inv_type}] 已解析: {file_path.name}")
|
||
elif result.get("invoice_number"):
|
||
all_invoices.append(result)
|
||
log.info(f"[{inv_type}] 已解析: {file_path.name}")
|
||
else:
|
||
log.warning(f"无法分类: {file_path.name} (invoice_type={inv_type})")
|
||
|
||
# 全部文件提取失败时抛出异常，携带原始错误信息
|
||
if failed_files and not all_invoices and not all_cards and not applications:
|
||
failed_names = [name for name, _ in failed_files]
|
||
error_details = {name: err for name, err in failed_files}
|
||
raise ExtractionError(
|
||
f"所有 {len(all_files)} 个文件提取均失败",
|
||
failed_files=failed_names,
|
||
details=error_details,
|
||
)
|
||
|
||
log.info(
|
||
f"分类结果: 发票 {len(all_invoices)} 张, 支付记录 {len(all_cards)} 条, 申请单 {len(applications)} 份, 附件 {len(attachments)} 个"
|
||
)
|
||
|
||
if not all_invoices:
|
||
log.warning("未成功解析任何发票")
|
||
|
||
payment_records = match_invoices_to_cards(all_invoices, all_cards)
|
||
_save_match_result(payment_records, cache_dir)
|
||
|
||
# 构建分类
|
||
all_documents: list[dict[str, str]] = []
|
||
for record in payment_records:
|
||
all_documents.extend(record.get("_matched_invoices", []))
|
||
all_documents.extend(applications)
|
||
|
||
groups = classify_invoice_batch(all_documents)
|
||
groups["attachments"] = attachments
|
||
log.info(
|
||
f"文档分类: 差旅发票 {len(groups['travel'])} 张, "
|
||
f"普通发票 {len(groups['general'])} 张, "
|
||
f"出差事前申请单 {len(groups['application'])} 份, "
|
||
f"附件 {len(groups['attachments'])} 个"
|
||
)
|
||
|
||
return payment_records, applications, groups
|