chat-ai-api/app/utils/document_processor.py

"""
文档处理工具
支持 Excel, Word, PDF 的文本提取和智能切片
"""
import re
from typing import List, Dict, Optional
import docx
import openpyxl
import PyPDF2
from io import BytesIO


class DocumentProcessor:
    def __init__(self):
        self.chunk_size = 800  # 每个切片的字符数
        self.chunk_overlap = 100  # 切片之间的重叠字符数

    def extract_text(self, file_content: bytes, file_type: str) -> Optional[str]:
        """
        从文件中提取文本
        :param file_content: 文件内容（字节）
        :param file_type: 文件类型（.xlsx, .docx, .pdf等）
        :return: 提取的文本内容
        """
        try:
            if file_type in ['.xlsx', '.xls']:
                return self._extract_excel(file_content)
            elif file_type in ['.docx', '.doc']:
                return self._extract_word(file_content)
            elif file_type == '.pdf':
                return self._extract_pdf(file_content)
            else:
                return None
        except Exception as e:
            print(f"文本提取失败: {str(e)}")
            return None

    def _extract_excel(self, file_content: bytes) -> str:
        """提取 Excel 文本"""
        wb = openpyxl.load_workbook(BytesIO(file_content), data_only=True)
        text_parts = []

        for sheet_name in wb.sheetnames:
            sheet = wb[sheet_name]
            text_parts.append(f"\n=== 工作表: {sheet_name} ===\n")

            for row in sheet.iter_rows(values_only=True):
                row_text = '\t'.join([str(cell) if cell is not None else '' for cell in row])
                if row_text.strip():
                    text_parts.append(row_text)

        return '\n'.join(text_parts)

    def _extract_word(self, file_content: bytes) -> str:
        """提取 Word 文本"""
        doc = docx.Document(BytesIO(file_content))
        text_parts = []

        for para in doc.paragraphs:
            if para.text.strip():
                text_parts.append(para.text)

        # 提取表格
        for table in doc.tables:
            for row in table.rows:
                row_text = '\t'.join([cell.text.strip() for cell in row.cells])
                if row_text.strip():
                    text_parts.append(row_text)

        return '\n'.join(text_parts)

    def _extract_pdf(self, file_content: bytes) -> str:
        """提取 PDF 文本"""
        pdf_reader = PyPDF2.PdfReader(BytesIO(file_content))
        text_parts = []

        for page_num in range(len(pdf_reader.pages)):
            page = pdf_reader.pages[page_num]
            text = page.extract_text()
            if text.strip():
                text_parts.append(f"\n=== 第 {page_num + 1} 页 ===\n")
                text_parts.append(text)

        return '\n'.join(text_parts)

    def split_text_into_chunks(self, text: str) -> List[Dict]:
        """
        将文本智能切片
        :param text: 原始文本
        :return: 切片列表 [{"content": "...", "metadata": {...}}, ...]
        """
        # 按段落分割
        paragraphs = [p.strip() for p in text.split('\n') if p.strip()]

        chunks = []
        current_chunk = ""
        current_length = 0

        for para in paragraphs:
            para_length = len(para)

            # 如果当前段落加入后超过chunk_size，则保存当前chunk
            if current_length + para_length > self.chunk_size and current_chunk:
                chunks.append({
                    "content": current_chunk.strip(),
                    "metadata": {"char_count": current_length}
                })

                # 保留overlap部分
                if current_length > self.chunk_overlap:
                    overlap_text = current_chunk[-self.chunk_overlap:]
                    current_chunk = overlap_text + "\n" + para
                    current_length = len(overlap_text) + para_length
                else:
                    current_chunk = para
                    current_length = para_length
            else:
                # 继续累积
                if current_chunk:
                    current_chunk += "\n" + para
                else:
                    current_chunk = para
                current_length += para_length

        # 保存最后一个chunk
        if current_chunk.strip():
            chunks.append({
                "content": current_chunk.strip(),
                "metadata": {"char_count": current_length}
            })

        return chunks

    def clean_text(self, text: str) -> str:
        """
        清理文本：去除多余空格、特殊字符等
        """
        # 去除多余空白
        text = re.sub(r'\s+', ' ', text)
        # 去除特殊控制字符
        text = re.sub(r'[\x00-\x08\x0b-\x0c\x0e-\x1f\x7f]', '', text)
        return text.strip()


# 全局单例
_document_processor = None


def get_document_processor() -> DocumentProcessor:
    """获取文档处理器单例"""
    global _document_processor
    if _document_processor is None:
        _document_processor = DocumentProcessor()
    return _document_processor