"""PDF 提取服务 — 从 PDF 计划书中提取结构化数据。""" import os import re import json import time import hashlib import logging from dataclasses import dataclass, field from typing import Optional logger = logging.getLogger(__name__) CACHE_VERSION = 3 @dataclass class ExtractionResult: pdf_path: str product_name: str plan_type: str # savings/ci/iul status: str # success/cached/error data: Optional[dict] = None usage: Optional[dict] = None error: Optional[str] = None duration_ms: float = 0 def infer_plan_type(raw: dict) -> str: """从 LLM 输出推断产品类型。""" t = str(raw.get("product_type", "")).lower() if "ci" in t or "critical" in t: return "ci" if "iul" in t or "universal" in t: return "iul" rows = raw.get("benefit_illustration", []) if not isinstance(rows, list): rows = [] has_savings = any( r.get("total_surrender_value") is not None or r.get("guaranteed_cash_value") is not None or r.get("reversionary_bonus") is not None for r in rows if isinstance(r, dict) ) has_ci = any( r.get("surrender_value_total") is not None or r.get("death_benefit_total") is not None for r in rows if isinstance(r, dict) ) has_iul = any( r.get("cash_value") is not None or r.get("account_value") is not None for r in rows if isinstance(r, dict) ) if has_iul: return "iul" if has_savings: return "savings" if has_ci: return "ci" policy = raw.get("policy", {}) if isinstance(policy, dict): if policy.get("index_account_rate") is not None or policy.get("capital_partition") is not None: return "iul" if policy.get("sum_insured") is not None: return "ci" return "savings" def _hash_file(file_path: str) -> str: """计算文件 SHA-256 哈希。""" h = hashlib.sha256() with open(file_path, "rb") as f: for chunk in iter(lambda: f.read(8192), b""): h.update(chunk) return h.hexdigest() def _get_cache_path(pdf_path: str, cache_dir: str) -> str: """获取缓存文件路径。""" file_hash = _hash_file(pdf_path) return os.path.join(cache_dir, f"{file_hash}.json") def _extract_pdf_text(pdf_path: str, max_chars: int = 30000) -> str: """使用 PyMuPDF 提取 PDF 文本。""" try: import fitz # PyMuPDF doc = fitz.open(pdf_path) text_parts = [] for page in doc: text_parts.append(page.get_text()) doc.close() full_text = "\n".join(text_parts) if len(full_text) > max_chars: full_text = full_text[:max_chars] return full_text except ImportError: logger.warning("PyMuPDF 未安装,无法提取 PDF 文本") return "" def _looks_corrupted(text: str) -> bool: """检测 PDF 文本是否乱码。""" if not text or len(text) < 50: return True # 计算坏字符比例 bad_chars = sum(1 for c in text if ord(c) < 32 and c not in "\n\r\t") return bad_chars / len(text) > 0.1 class ExtractionOrchestrator: """PDF 提取编排器。""" def __init__(self, use_cache: bool = True, cache_dir: str = ".cache/insurance-ppt"): self.use_cache = use_cache self.cache_dir = cache_dir async def extract_plan(self, pdf_path: str, plan_type: str = "savings") -> ExtractionResult: """从 PDF 提取结构化数据。""" from insurance.ppt.llm_client import llm_client from insurance.ppt.prompts import ( SAVINGS_PLAN_SYSTEM_PROMPT, CI_PLAN_SYSTEM_PROMPT, IUL_SYSTEM_PROMPT, build_savings_prompt, ) start = time.time() abs_path = os.path.abspath(pdf_path) if not os.path.exists(abs_path): return ExtractionResult( pdf_path=abs_path, product_name="unknown", plan_type=plan_type, status="error", error="文件不存在", duration_ms=(time.time() - start) * 1000, ) # 检查缓存 if self.use_cache: cached = self._load_from_cache(abs_path) if cached: cached.duration_ms = (time.time() - start) * 1000 return cached # 提取 PDF 文本 pdf_text = _extract_pdf_text(abs_path) if not pdf_text: return ExtractionResult( pdf_path=abs_path, product_name="unknown", plan_type=plan_type, status="error", error="无法提取 PDF 文本", duration_ms=(time.time() - start) * 1000, ) # 选择 prompt prompts = { "savings": SAVINGS_PLAN_SYSTEM_PROMPT, "ci": CI_PLAN_SYSTEM_PROMPT, "iul": IUL_SYSTEM_PROMPT, } system_prompt = prompts.get(plan_type, SAVINGS_PLAN_SYSTEM_PROMPT) # 调用 LLM try: data, response = await llm_client.structured_output( prompt=f"请从以下PDF文本中提取保险计划书数据:\n\n{pdf_text[:20000]}", system_prompt=system_prompt, ) except Exception as e: return ExtractionResult( pdf_path=abs_path, product_name="unknown", plan_type=plan_type, status="error", error=f"LLM 调用失败: {e}", duration_ms=(time.time() - start) * 1000, ) # 储蓄险数据修复:确保 total >= gcv if plan_type == "savings" and isinstance(data.get("benefit_illustration"), list): for row in data["benefit_illustration"]: if not isinstance(row, dict): continue gcv = float(row.get("guaranteed_cash_value") or 0) rev = float(row.get("reversionary_bonus") or 0) term = float(row.get("terminal_dividend") or 0) total = float(row.get("total_surrender_value") or 0) if total < gcv: row["total_surrender_value"] = gcv + rev + term # 推断产品类型 detected_type = infer_plan_type(data) # 写入缓存 if self.use_cache: self._save_to_cache(abs_path, data) product_name = data.get("product_name", "unknown") return ExtractionResult( pdf_path=abs_path, product_name=product_name, plan_type=detected_type, status="success", data=data, usage={"input": response.tokens.get("input", 0), "output": response.tokens.get("output", 0)} if response.tokens else None, duration_ms=(time.time() - start) * 1000, ) async def extract_multiple(self, pdf_paths: list[str], plan_type: str = "savings") -> list[ExtractionResult]: """顺序提取多个 PDF。""" results = [] for pdf_path in pdf_paths: logger.info(f" 📄 {os.path.basename(pdf_path)}...") results.append(await self.extract_plan(pdf_path, plan_type)) return results def _load_from_cache(self, pdf_path: str) -> Optional[ExtractionResult]: """从缓存加载。""" try: cache_path = _get_cache_path(pdf_path, self.cache_dir) if not os.path.exists(cache_path): return None with open(cache_path, "r", encoding="utf-8") as f: raw = json.load(f) meta = raw.get("_meta", {}) if meta.get("cacheVersion") != CACHE_VERSION: return None data = raw.get("_data", raw) product_name = data.get("product_name", "unknown") plan_type = infer_plan_type(data) return ExtractionResult( pdf_path=pdf_path, product_name=product_name, plan_type=plan_type, status="cached", data=data, ) except Exception: return None def _save_to_cache(self, pdf_path: str, data: dict): """写入缓存。""" try: os.makedirs(self.cache_dir, exist_ok=True) file_hash = _hash_file(pdf_path) cache_data = { "_data": data, "_meta": { "cacheVersion": CACHE_VERSION, "originalFile": os.path.basename(pdf_path), "extractedAt": __import__("datetime").datetime.now().isoformat(), "fileHash": file_hash, }, } cache_path = os.path.join(self.cache_dir, f"{file_hash}.json") with open(cache_path, "w", encoding="utf-8") as f: json.dump(cache_data, f, ensure_ascii=False, indent=2) except Exception as e: logger.warning(f"缓存写入失败: {e}") async def extract_for_poster(self, filepath: str) -> dict: """提取海报所需的关键字段(精简 prompt,降低 LLM 成本)。 返回: {age, gender, currency, sum_assured, premium_term, annual_premium, coverage_period, key_benefits} """ from insurance.ppt.llm_client import llm_client abs_path = os.path.abspath(filepath) text = _extract_pdf_text(abs_path) if not text: raise ValueError("无法提取 PDF 文本") text = text[:6000] system_prompt = ( "你是一位保险计划书解析专家。请从以下计划书内容中提取海报所需的关键字段。\n" "输出 JSON 格式(不要包含 markdown 代码块标记):\n" '{"age": 35, "gender": "男", "currency": "USD", "sum_assured": 500000, ' '"premium_term": 5, "annual_premium": 100000, "coverage_period": "终身", ' '"key_benefits": ["身故赔偿", "全残保障"]}\n' "注意:数值用数字,不要带货币符号。" ) result, _response = await llm_client.structured_output( text, system_prompt, schema={ "type": "object", "properties": { "age": {"type": "number"}, "gender": {"type": "string"}, "currency": {"type": "string"}, "sum_assured": {"type": "number"}, "premium_term": {"type": "number"}, "annual_premium": {"type": "number"}, "coverage_period": {"type": "string"}, "key_benefits": {"type": "array", "items": {"type": "string"}}, }, "required": ["age", "gender", "currency", "sum_assured", "annual_premium"], }, ) return result