proposa-ai/bid_proposal/bid_reader.py

"""
bid_reader.py - 招标文件读取模块

负责读取招标 PDF，提取文本和表格内容，返回结构化的 BidDocument 对象。
"""

import logging
import os
from typing import Optional

from bid_proposal.models import BidDocument, BidPage

logger = logging.getLogger(__name__)


class BidReadError(Exception):
    """招标文件读取失败异常"""

    pass


def read_bid_pdf(pdf_path: str) -> BidDocument:
    """读取招标 PDF，提取文本和表格。

    实现步骤：
    1. 使用 pymupdf (import fitz) 打开 PDF，逐页提取文本
    2. 使用 pdf_table_to_docx 提取表格
    3. 组合成 BidDocument 返回

    Args:
        pdf_path: 招标文件 PDF 路径

    Returns:
        BidDocument: 结构化的招标文件内容

    Raises:
        FileNotFoundError: PDF 文件不存在
        BidReadError: PDF 读取/解析失败（如文件损坏、加密等）
    """
    if not os.path.exists(pdf_path):
        raise FileNotFoundError(f"招标文件不存在: {pdf_path}")

    logger.info("开始读取招标文件: %s", pdf_path)

    # 1. 使用 pymupdf 提取文本
    pages = _extract_text(pdf_path)

    # 2. 使用 pdf_table_to_docx 提取表格
    tables = _extract_tables(pdf_path)

    # 3. 构建 BidDocument
    doc = BidDocument(
        pdf_path=os.path.abspath(pdf_path),
        pages=pages,
        tables=tables,
    )

    logger.info("读取完成: %d 页, %d 个表格, %d 字符", doc.total_pages, doc.table_count, len(doc.full_text))
    return doc


def _extract_text(pdf_path: str) -> list[BidPage]:
    """使用 pymupdf 从 PDF 中逐页提取文本。

    Args:
        pdf_path: PDF 文件路径

    Returns:
        list[BidPage]: 每页的文本内容列表

    Raises:
        BidReadError: PDF 打开或读取失败
    """
    try:
        import fitz  # pymupdf
    except ImportError:
        raise ImportError("请安装 pymupdf: pip install pymupdf")

    try:
        doc = fitz.open(pdf_path)
    except Exception as e:
        raise BidReadError(f"无法打开 PDF 文件: {e}")

    pages: list[BidPage] = []
    try:
        for page_num in range(len(doc)):
            page = doc[page_num]
            text = page.get_text().strip()
            bid_page = BidPage(page_num=page_num + 1, text=text)
            pages.append(bid_page)
            logger.debug("第 %d 页: %d 字符", page_num + 1, len(text))
    finally:
        doc.close()

    if not pages:
        raise BidReadError(f"PDF 文件为空或无法提取任何文本: {pdf_path}")

    return pages


def _extract_tables(pdf_path: str) -> list:
    """使用 pdf_table_to_docx 从 PDF 中提取表格。

    Args:
        pdf_path: PDF 文件路径

    Returns:
        list[TableInfo]: 表格列表

    Note:
        如果表格提取失败，记录警告并返回空列表，不影响文本提取结果。
    """
    try:
        from pdf_table_to_docx.extractor import PDFTableExtractor

        extractor = PDFTableExtractor(pdf_path)
        extracted_tables = extractor.extract()
        if extracted_tables:
            tables = list(extracted_tables)
            logger.info("提取到 %d 个表格", len(tables))
            return tables
        else:
            logger.info("PDF 中未发现表格")
            return []
    except Exception as e:
        logger.warning("表格提取失败，仅使用文本内容继续: %s", e)
        return []