run.py

#!/usr/bin/env python3
"""
============================================================
PDF OCR 智能问答系统 — 端到端运行脚本
============================================================

用法:
    # 交互模式: 处理文档后进入问答 REPL
    python run.py -f document.pdf

    # 单次问答
    python run.py -f document.pdf -q "文档主要内容是什么?"

    # 批量处理多个文档
    python run.py -f doc1.pdf doc2.png scan3.jpg

    # 指定分块参数
    python run.py -f document.pdf --chunk-size 1000 --chunk-overlap 200

    # 从已有向量库加载 (跳过 OCR, 直接问答)
    python run.py --load

    # 清空旧数据重新处理
    python run.py -f document.pdf --clear

    # 显示检索到的原文
    python run.py -f document.pdf -q "问题" --show-sources

环境变量 (或 .env 文件):
    EMBEDDING_API_BASE   Embedding API 地址
    EMBEDDING_MODEL_NAME Embedding 模型名
    LLM_API_BASE         LLM API 地址
    LLM_API_KEY          LLM API Key
    LLM_MODEL_NAME       LLM 模型名
"""

import argparse
import json
import os
import sys
import time
from pathlib import Path
from typing import List, Optional

# ---- 环境补丁 (必须在其他导入之前) ----
def _patch():
    import types as _types
    if "langchain_text_splitters" not in sys.modules:
        m = _types.ModuleType("langchain_text_splitters")
        m.__path__ = []
        sys.modules["langchain_text_splitters"] = m
    try:
        import torch  # noqa: F401
    except ImportError:
        pass


_patch()

# 项目导入
sys.path.insert(0, str(Path(__file__).resolve().parent))

import config
from ocr_loader import PaddleOCRLoader
from text_processor import TextProcessingPipeline, RecursiveCharacterTextSplitter
from embeddings import get_embedding_model
from vector_store import VectorStoreManager, build_vector_store
from rag_chain import RAGChain, create_llm, PDFRAGPipeline

# 将内置分割器注入到 mock 模块
import sys as _sys
_lts = _sys.modules.get("langchain_text_splitters")
if _lts is not None:
    _lts.RecursiveCharacterTextSplitter = RecursiveCharacterTextSplitter

from loguru import logger

# ============================================================
# Banner
# ============================================================

BANNER = r"""
  ┌──────────────────────────────────────────────────────┐
  │         📄 PDF OCR 智能问答系统                        │
  │                                                      │
  │  OCR:   PaddleOCR-VL-1.5 (本地)                      │
  │  嵌入:  {emb_model}                                  │
  │  LLM:   {llm_model}                                  │
  │  向量库: {vec_store}                                  │
  └──────────────────────────────────────────────────────┘
"""


def print_banner():
    emb_name = config.EMBEDDING_MODEL_NAME
    llm_name = config.LLM_MODEL_NAME
    vs = config.VECTOR_STORE_TYPE
    # 截断过长的模型名
    if len(emb_name) > 35:
        emb_name = emb_name[:32] + "..."
    if len(llm_name) > 35:
        llm_name = llm_name[:32] + "..."
    print(BANNER.format(emb_model=emb_name, llm_model=llm_name, vec_store=vs))


# ============================================================
# 步骤函数
# ============================================================

def _save_documents(docs: list, path: Path, label: str = "文档"):
    """将 LangChain Document 列表保存为 JSON"""
    path.parent.mkdir(parents=True, exist_ok=True)
    data = []
    for doc in docs:
        data.append({
            "page_content": doc.page_content,
            "metadata": {k: v for k, v in doc.metadata.items()
                         if isinstance(v, (str, int, float, bool, type(None)))}
        })
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    print(f"  💾 {label}已保存: {path} ({len(data)} 条)")


def step_ocr(file_paths: List[str], output_dir: Optional[Path] = None) -> list:
    """Step 1: OCR 识别所有文件, 全部结果合并保存到一个文件"""
    all_docs = []
    for fp in file_paths:
        fp = Path(fp)
        if not fp.exists():
            logger.error(f"文件不存在: {fp}")
            continue
        suffix = fp.suffix.lower()
        if suffix not in config.SUPPORTED_FORMATS:
            logger.warning(f"跳过不支持格式: {fp} (支持: {config.SUPPORTED_FORMATS})")
            continue

        icon = "📄" if suffix == ".pdf" else "🖼️"
        print(f"  {icon} 正在识别: {fp.name} ...", end=" ", flush=True)
        t0 = time.time()
        loader = PaddleOCRLoader(str(fp), verbose=True)
        docs = loader.load()
        elapsed = time.time() - t0
        print(f"{len(docs)} 页/文档 ({elapsed:.1f}s)")
        all_docs.extend(docs)

    # 所有文件识别完后统一保存
    if output_dir and all_docs:
        save_path = output_dir / "ocr_results.json"
        _save_documents(all_docs, save_path, "OCR结果 ")

    return all_docs


def step_process(
    documents: list, chunk_size: int, chunk_overlap: int,
    output_dir: Optional[Path] = None
) -> list:
    """Step 2: 文本清洗 + 分割, 全部结果合并保存到一个文件"""
    print(f"  ✂️  正在分割: {len(documents)} 个文档 ...", end=" ", flush=True)
    t0 = time.time()
    pipeline = TextProcessingPipeline(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
    )
    chunks = pipeline.process(documents)
    elapsed = time.time() - t0
    print(f"→ {len(chunks)} 个文本块 ({elapsed:.1f}s)")

    if output_dir and chunks:
        save_path = output_dir / "chunks.json"
        _save_documents(chunks, save_path, "分块结果 ")

    return chunks


def step_embed(chunks: list) -> VectorStoreManager:
    """Step 3: 向量嵌入 + 入库"""
    print(f"  🧠 正在向量化: {len(chunks)} 个文本块 ...", end=" ", flush=True)
    t0 = time.time()
    manager = build_vector_store(chunks, clear_existing=True)
    elapsed = time.time() - t0
    print(f"完成 ({elapsed:.1f}s)")
    return manager


def step_rag(manager: VectorStoreManager):
    """Step 4: 初始化 RAG 链"""
    llm = create_llm()
    chain = RAGChain(vector_store_manager=manager, llm=llm)
    return chain


# ============================================================
# 核心流程
# ============================================================

def run_ingest(
    file_paths: List[str],
    chunk_size: int = config.CHUNK_SIZE,
    chunk_overlap: int = config.CHUNK_OVERLAP,
    clear: bool = True,
    output_dir: Optional[Path] = None,
) -> VectorStoreManager:
    """完整入库流程: OCR → 处理 → 嵌入 → 入库"""
    print("\n" + "─" * 55)
    print("  📥 阶段 1: 文档入库")
    print("─" * 55)

    # Step 1: OCR
    t_start = time.time()
    documents = step_ocr(file_paths, output_dir=output_dir)
    if not documents:
        logger.error("未识别到任何文本内容, 请检查文件是否包含可读文字")
        sys.exit(1)
    print(f"      总计: {len(documents)} 个原始文档页")

    # Step 2: 处理
    chunks = step_process(documents, chunk_size, chunk_overlap,
                          output_dir=output_dir)

    # Step 3: 嵌入入库
    manager = step_embed(chunks)

    total_time = time.time() - t_start
    print(f"\n  ✅ 入库完成 (总耗时 {total_time:.1f}s)")
    print(f"     文档: {len(documents)} 页 → {len(chunks)} 个文本块")
    print(f"     向量维度: {config.EMBEDDING_MODEL_NAME}")
    print(f"     存储: {config.VECTOR_STORE_TYPE} @ {config.VECTOR_DB_DIR}")

    return manager


def run_qa(chain: RAGChain, question: str, show_sources: bool = False):
    """执行单次问答"""
    print("\n" + "─" * 55)
    print(f"  ❓ 问题: {question}")
    print("─" * 55)

    t0 = time.time()
    result = chain.query(question)
    elapsed = time.time() - t0

    print(f"\n  🤖 回答 ({elapsed:.1f}s):")
    print("─" * 55)
    print(result["answer"])

    if show_sources:
        print(f"\n  📚 参考来源 ({len(result['sources'])} 条):")
        print("─" * 55)
        for src in result["sources"]:
            print(f"  [{src['rank']}] {src['document']} | 第{src['page']}页 "
                  f"| {src['content_type']}")
            print(f"      {src['content'][:120]}...")

    return result


def run_repl(chain: RAGChain):
    """交互式问答 REPL"""
    print("\n" + "─" * 55)
    print("  💬 交互问答模式")
    print("─" * 55)
    print("  输入问题后回车, 输入 :s 切换来源显示")
    print("  输入 :q 退出, :c 清屏, :h 帮助")
    print("─" * 55)

    chat_history = []
    show_sources = False

    while True:
        try:
            user_input = input("\n  🔍 > ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n  再见! 👋")
            break

        if not user_input:
            continue

        # 命令处理
        if user_input.startswith(":"):
            cmd = user_input[1:].strip().lower()
            if cmd in ("q", "quit", "exit"):
                print("  再见! 👋")
                break
            elif cmd == "s":
                show_sources = not show_sources
                print(f"  来源显示: {'开启' if show_sources else '关闭'}")
            elif cmd == "c":
                os.system("clear" if os.name != "nt" else "cls")
            elif cmd == "h":
                print("  命令: :q 退出 | :s 切换来源 | :c 清屏 | :h 帮助")
            else:
                print(f"  未知命令: {user_input}")
            continue

        # 问答
        t0 = time.time()
        result = chain.query_with_history(user_input, chat_history)
        elapsed = time.time() - t0

        print(f"\n  🤖 ({elapsed:.1f}s):")
        print(f"  {result['answer']}")

        if show_sources:
            print(f"\n  📚 来源 ({len(result['sources'])} 条):")
            for src in result["sources"]:
                print(f"  [{src['rank']}] {src['document']} "
                      f"第{src['page']}页 | {src['content_type']}")

        chat_history.append({"role": "user", "content": user_input})
        chat_history.append({"role": "assistant", "content": result["answer"]})


# ============================================================
# API 连通性检查
# ============================================================

def check_apis() -> bool:
    """检查 Embedding API 和 LLM API 是否可达"""
    import urllib.request

    all_ok = True

    # 检查 Embedding API
    emb_url = config.EMBEDDING_API_BASE.rstrip("/")
    try:
        req = urllib.request.Request(f"{emb_url}/models", method="HEAD")
        urllib.request.urlopen(req, timeout=5)
        print(f"  ✅ Embedding API: {emb_url}")
    except Exception as e:
        print(f"  ⚠️  Embedding API: {emb_url} — {e}")
        all_ok = False

    # 检查 LLM API
    llm_url = config.LLM_API_BASE.rstrip("/")
    try:
        req = urllib.request.Request(f"{llm_url}/models", method="HEAD")
        urllib.request.urlopen(req, timeout=5)
        print(f"  ✅ LLM API: {llm_url}")
    except Exception as e:
        print(f"  ⚠️  LLM API: {llm_url} — {e}")
        all_ok = False

    return all_ok


# ============================================================
# 主入口
# ============================================================

def main():
    parser = argparse.ArgumentParser(
        description="PDF OCR 智能问答系统 — 端到端运行脚本",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
示例:
  python run.py -f document.pdf                    # 交互问答
  python run.py -f doc.pdf -q "主要内容?"           # 单次问答
  python run.py -f a.pdf b.png --clear             # 批量处理
  python run.py --load                             # 加载已有向量库
        """,
    )
    parser.add_argument(
        "-f", "--files", nargs="+",
        default=["/data/huangjie/Project/dProject/pdfocr/过滤网modify.pdf",
                 "/data/huangjie/Project/dProject/pdfocr/videoagent.png",
                 "/data/huangjie/Project/dProject/pdfocr/biaozhun.jpg"],
        help="要处理的文档路径 (PDF/PNG/JPG/BMP/TIF)",
    )
    parser.add_argument(
        "-q", "--question",
        help="单次问答 (不进入交互模式)",
    )
    parser.add_argument(
        "--load", action="store_true",
        help="加载已有向量库, 跳过 OCR 处理",
    )
    parser.add_argument(
        "--clear", action="store_true",
        help="清空旧向量库数据后重新处理",
    )
    parser.add_argument(
        "--chunk-size", type=int, default=config.CHUNK_SIZE,
        help=f"文本块大小 (默认: {config.CHUNK_SIZE})",
    )
    parser.add_argument(
        "--chunk-overlap", type=int, default=config.CHUNK_OVERLAP,
        help=f"块间重叠字符数 (默认: {config.CHUNK_OVERLAP})",
    )
    parser.add_argument(
        "--show-sources", action="store_true",
        help="在回答中显示参考来源",
    )
    parser.add_argument(
        "--top-k", type=int, default=config.RETRIEVAL_TOP_K,
        help=f"检索返回文档数 (默认: {config.RETRIEVAL_TOP_K})",
    )
    parser.add_argument(
        "--skip-api-check", action="store_true",
        help="跳过 API 连通性检查",
    )
    parser.add_argument(
        "--output-dir", type=str, default=None,
        help=f"中间结果保存目录 (默认: {config.OCR_OUTPUT_DIR})",
    )

    args = parser.parse_args()

    # Banner
    print_banner()

    # API 检查
    if not args.skip_api_check:
        print("  🔌 API 连通性检查:")
        check_apis()
        print()

    # 模式判断
    if args.load:
        # 加载已有向量库
        print("  📂 加载已有向量库...")
        manager = VectorStoreManager(store_type=config.VECTOR_STORE_TYPE)
        count = manager.get_document_count()
        if count == 0:
            logger.error("向量库为空! 请先用 -f 指定文件进行入库")
            sys.exit(1)
        print(f"  ✅ 已加载: {count} 个文档块")
    elif args.files:
        # 处理文件
        output_dir = Path(args.output_dir) if args.output_dir else config.OCR_OUTPUT_DIR
        manager = run_ingest(
            args.files,
            chunk_size=args.chunk_size,
            chunk_overlap=args.chunk_overlap,
            clear=args.clear,
            output_dir=output_dir,
        )
    else:
        parser.print_help()
        print("\n  ❌ 请指定 -f/--files 或 --load")
        sys.exit(1)

    # 初始化 RAG 链
    print("\n" + "─" * 55)
    print("  🔗 阶段 2: 初始化 RAG 问答引擎")
    print("─" * 55)
    llm = create_llm()
    chain = RAGChain(
        vector_store_manager=manager,
        llm=llm,
        top_k=args.top_k,
    )
    print(f"  ✅ RAG 引擎就绪 (LLM={config.LLM_MODEL_NAME})")

    # 问答
    if args.question:
        run_qa(chain, args.question, show_sources=args.show_sources)
    else:
        run_repl(chain)


if __name__ == "__main__":
    main()