101 lines
3.2 KiB
Python
101 lines
3.2 KiB
Python
"""
|
||
管道核心逻辑
|
||
|
||
抽取 pipeline.py（CLI 管道）和 pipeline_web.py（Web 管道）的公共数据流：
|
||
发票分类判断 -> 差旅/普通信息提取 -> 缓存读写
|
||
|
||
两个入口分别传入不同的目录参数，复用此模块。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from . import get_logger
|
||
from .doc.llm_extractor import (
|
||
CACHE_DIR_NAME,
|
||
extract_normal_info,
|
||
extract_travel_info,
|
||
load_cache,
|
||
)
|
||
|
||
log = get_logger("pipeline_core")
|
||
|
||
|
||
def is_travel_invoice(groups: dict[str, list[dict[str, Any]]]) -> bool:
|
||
"""判断是否为纯差旅发票（有差旅发票且无普通发票）。
|
||
|
||
注意：系统仅支持「纯差旅」和「普通报销」两种模式。
|
||
若同时存在差旅发票和普通发票（混合），则视为普通报销模式处理——差旅发票
|
||
对应的费用仍会在普通报销中按项目填报。若需要严格区分，上游应在发票分类
|
||
后报错提示用户分开提交。
|
||
"""
|
||
return bool(groups.get("travel")) and not bool(groups.get("general"))
|
||
|
||
|
||
def save_cache_info(cache_path: Path, info_key: str, info: dict[str, Any]) -> None:
|
||
"""将提取结果保存到缓存目录
|
||
|
||
Args:
|
||
cache_path: 会话目录路径。
|
||
info_key: 缓存键名（"travel_info" 或 "normal_info"）。
|
||
info: 提取结果字典。
|
||
"""
|
||
cache_dir = cache_path / CACHE_DIR_NAME
|
||
cache_dir.mkdir(parents=True, exist_ok=True)
|
||
with open(cache_dir / f"{info_key}.json", "w", encoding="utf-8") as f:
|
||
json.dump(info, f, ensure_ascii=False, indent=2)
|
||
log.info("%s 已保存到缓存", info_key)
|
||
|
||
|
||
def extract_and_cache_travel_info(
|
||
groups: dict[str, list[dict[str, Any]]],
|
||
cache_path: Path,
|
||
) -> dict[str, Any] | None:
|
||
"""当存在差旅发票时，调用 LLM 提取差旅信息并缓存。
|
||
|
||
Returns:
|
||
差旅信息字典，非差旅时返回 None。
|
||
"""
|
||
if not groups.get("travel"):
|
||
return None
|
||
|
||
# 检查缓存是否已有
|
||
cache_map = load_cache(cache_path)
|
||
travel_info = cache_map.get("travel_info")
|
||
if travel_info:
|
||
log.info("使用已有差旅信息缓存")
|
||
return travel_info # type: ignore[no-any-return]
|
||
|
||
log.info("开始提取差旅信息...")
|
||
travel_info = extract_travel_info(source_dir=cache_path)
|
||
save_cache_info(cache_path, "travel_info", travel_info)
|
||
return travel_info
|
||
|
||
|
||
def extract_and_cache_normal_info(
|
||
groups: dict[str, list[dict[str, Any]]],
|
||
cache_path: Path,
|
||
) -> dict[str, Any] | None:
|
||
"""当存在普通发票时，调用 LLM 提取普通报销信息并缓存。
|
||
|
||
Returns:
|
||
普通报销信息字典，非普通时返回 None。
|
||
"""
|
||
if not groups.get("general"):
|
||
return None
|
||
|
||
# 检查缓存是否已有
|
||
cache_map = load_cache(cache_path)
|
||
normal_info = cache_map.get("normal_info")
|
||
if normal_info:
|
||
log.info("使用已有普通发票信息缓存")
|
||
return normal_info # type: ignore[no-any-return]
|
||
|
||
log.info("开始提取普通发票信息...")
|
||
normal_info = extract_normal_info(source_dir=cache_path)
|
||
save_cache_info(cache_path, "normal_info", normal_info)
|
||
return normal_info
|