重构项目为LLM 驱动

src/doc/llm_extractor.py

@@ -0,0 +1,210 @@
+"""
+LLM 信息提取
+
+使用 LLM 从 PDF 文本中提取结构化数据，以及从支付截图中提取刷卡信息。
+支持 JSON 格式输出，字段与 CSV_COLUMNS 对齐。
+
+对外接口:
+    extract_invoice_from_text(text, file_name) -> dict    从 PDF 文本提取发票信息
+    extract_card_info_from_image(image_path) -> dict     从支付截图提取刷卡信息
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+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
+
+log = get_logger("llm_extractor")
+
+
+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=8192,
+        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()
+
+    # 处理 ```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))
+
+
+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
+
+
+# ------------------------------------------------------------------
+# 支付截图信息提取（多模态）
+# ------------------------------------------------------------------
+
+
+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_multimodal(
+    system_prompt: str,
+    text: str,
+    image_b64: str,
+    max_tokens: int = 4096,
+) -> str:
+    """发送多模态请求（文本 + 图片）到 LLM。"""
+    from llama_index.core.base.llms.types import ImageBlock, 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),
+                ImageBlock(
+                    url=f"data:image/jpeg;base64,{image_b64}",
+                    detail="high",
+                ),
+            ],
+        ),
+    ]
+
+    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 extract_card_info_from_image(image_path: Path) -> dict[str, Any]:
+    """从支付截图中提取刷卡信息。
+
+    Args:
+        image_path: 支付截图图片路径。
+
+    Returns:
+        包含刷卡日期、刷卡金额、公务卡号的字典。
+    """
+    system_prompt = build_card_info_system_prompt()
+    user_text = f"请分析以下支付截图并提取信息：\n\n文件名: {image_path.name}"
+
+    image_b64 = _image_to_base64(image_path)
+
+    try:
+        response = _llm_query_multimodal(system_prompt, user_text, image_b64, max_tokens=4096)
+        result = _parse_json_response(response)
+        log.info("LLM 支付截图提取成功: %s", image_path.name)
+        return result
+    except Exception as e:
+        log.error("LLM 支付截图提取失败: %s (%s)", image_path.name, e)
+        raise