ForcePilot/src/knowledge/indexing.py

import asyncio
import os
from pathlib import Path

from langchain_community.document_loaders import (
    CSVLoader,
    JSONLoader,
    PyPDFLoader,
    TextLoader,
    UnstructuredHTMLLoader,
    UnstructuredMarkdownLoader,
    UnstructuredWordDocumentLoader,
)
from langchain_text_splitters import RecursiveCharacterTextSplitter

from src.utils import logger

SUPPORTED_FILE_EXTENSIONS: tuple[str, ...] = (
    ".txt",
    ".md",
    ".doc",
    ".docx",
    ".html",
    ".htm",
    ".json",
    ".csv",
    ".xls",
    ".xlsx",
    ".pdf",
    ".jpg",
    ".jpeg",
    ".png",
    ".bmp",
    ".tiff",
    ".tif",
)


def is_supported_file_extension(file_name: str | os.PathLike[str]) -> bool:
    """Check whether the given file path has a supported extension."""
    return Path(file_name).suffix.lower() in SUPPORTED_FILE_EXTENSIONS


def _extract_word_text(file_path: Path) -> str:
    """
    Parse Word documents (.doc/.docx) into plain text.

    Try python-docx first for docx files and fall back to the unstructured
    loader so legacy .doc files are still parsed when possible.
    """
    try:
        from docx import Document  # type: ignore

        doc = Document(file_path)
        text = "\n".join(paragraph.text for paragraph in doc.paragraphs).strip()
        if text:
            return text
    except Exception as docx_error:  # noqa: BLE001
        logger.warning(f"python-docx failed to parse {file_path.name}: {docx_error}")

    try:
        loader = UnstructuredWordDocumentLoader(str(file_path))
        docs = loader.load()
        return "\n".join(doc.page_content for doc in docs).strip()
    except Exception as unstructured_error:  # noqa: BLE001
        logger.error(f"Unstructured failed to parse {file_path.name}: {unstructured_error}")
        raise ValueError(f"无法解析 Word 文档: {file_path.name}") from unstructured_error


def chunk_with_parser(file_path, params=None):
    """
    使用文件解析器将文件切分成固定大小的块

    Args:
        file_path: 文件路径
        params: 参数
    """
    params = params or {}
    chunk_size = int(params.get("chunk_size", 500))
    chunk_overlap = int(params.get("chunk_overlap", 100))

    file_type = Path(file_path).suffix.lower()

    # 选择合适的加载器
    if file_type in [".txt"]:
        loader = TextLoader(file_path)

    elif file_type in [".md"]:
        loader = UnstructuredMarkdownLoader(file_path)

    elif file_type in [".docx", ".doc"]:
        loader = UnstructuredWordDocumentLoader(file_path)

    elif file_type in [".html", ".htm"]:
        loader = UnstructuredHTMLLoader(file_path)

    elif file_type in [".json"]:
        loader = JSONLoader(file_path, jq_schema=".")

    elif file_type in [".csv"]:
        loader = CSVLoader(file_path)

    else:
        raise ValueError(f"不支持的文件类型: {file_type}")

    # 加载文档
    docs = loader.load()

    # 创建文本分割器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", ".", " ", ""],
    )

    # 分割文档
    nodes = text_splitter.split_documents(docs)

    # 添加序号信息到metadata
    for i, node in enumerate(nodes):
        if node.metadata is None:
            node.metadata = {}
        node.metadata["chunk_idx"] = i

    return nodes


def chunk_text(text, params=None):
    """
    将文本切分成固定大小的块
    """
    params = params or {}
    chunk_size = int(params.get("chunk_size", 500))
    chunk_overlap = int(params.get("chunk_overlap", 100))

    # 创建文本分割器
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size, chunk_overlap=chunk_overlap, separators=["\n\n", "\n", ".", " ", ""]
    )

    # 分割文档
    nodes = text_splitter.split_text(text)

    # 添加序号信息到metadata
    nodes = [{"text": node, "metadata": {"chunk_idx": i}} for i, node in enumerate(nodes)]
    return nodes


def chunk(text_or_path, params=None):
    raise NotImplementedError("chunk is deprecated, use chunk_with_parser or chunk_text instead")


def pdfreader(file_path, params=None):
    """读取PDF文件并返回text文本"""
    if isinstance(file_path, str):
        file_path = Path(file_path)

    assert file_path.exists(), "File not found"
    assert file_path.suffix.lower() == ".pdf", "File format not supported"

    # 使用LangChain的PDF加载器
    loader = PyPDFLoader(str(file_path))
    docs = loader.load()

    # 简单的拼接起来之后返回纯文本
    text = "\n\n".join([d.page_content for d in docs])
    return text


def plainreader(file_path):
    """读取普通文本文件并返回text文本"""
    assert os.path.exists(file_path), "File not found"

    # 使用LangChain的文本加载器
    loader = TextLoader(str(file_path))
    docs = loader.load()
    text = "\n\n".join([d.page_content for d in docs])
    return text


def parse_pdf(file, params=None):
    """
    解析PDF文件，支持多种OCR方式

    Args:
        file: PDF文件路径
        params: 参数字典，包含enable_ocr设置

    Returns:
        str: 解析得到的文本

    Raises:
        OCRServiceException: OCR服务不可用时抛出
    """
    from src.plugins._ocr import OCRServiceException

    params = params or {}
    opt_ocr = params.get("enable_ocr", "disable")

    if opt_ocr == "disable":
        return pdfreader(file, params=params)

    try:
        if opt_ocr == "onnx_rapid_ocr":
            from src.plugins import ocr

            return ocr.process_pdf(file, params=params)

        elif opt_ocr == "mineru_ocr":
            from src.plugins import ocr

            return ocr.process_file_mineru(file, params=params)

        elif opt_ocr == "paddlex_ocr":
            from src.plugins import ocr

            return ocr.process_file_paddlex(file, params=params)

        else:
            raise ValueError(f"不支持的OCR方式: {opt_ocr}")

    except OCRServiceException as e:
        logger.error(f"OCR service failed: {e.service_name} - {str(e)}")
        raise
    except Exception as e:
        logger.error(f"PDF parsing failed: {str(e)}")
        raise OCRServiceException(f"PDF解析失败: {str(e)}", opt_ocr, "parsing_failed")


def parse_image(file, params=None):
    """
    解析图像文件，支持多种OCR方式
    """
    from src.plugins._ocr import OCRServiceException

    params = params or {}
    opt_ocr = params.get("enable_ocr", "disable")

    if opt_ocr == "disable":
        logger.warning(f"OCR is disabled for image file: {file}, Using `onnx_rapid_ocr` instead")
        opt_ocr = "onnx_rapid_ocr"

    try:
        if opt_ocr == "onnx_rapid_ocr":
            from src.plugins import ocr

            return ocr.process_image(file, params=params)

        elif opt_ocr == "mineru_ocr":
            from src.plugins import ocr

            return ocr.process_file_mineru(file, params=params)

        elif opt_ocr == "paddlex_ocr":
            from src.plugins import ocr

            return ocr.process_file_paddlex(file, params=params)

        else:
            raise ValueError(f"不支持的OCR方式: {opt_ocr}")

    except OCRServiceException as e:
        logger.error(f"OCR service failed: {e.service_name} - {str(e)}")
        raise
    except Exception as e:
        logger.error(f"Image parsing failed: {str(e)}")
        raise OCRServiceException(f"Image解析失败: {str(e)}", opt_ocr, "parsing_failed")


async def parse_pdf_async(file, params=None):
    return await asyncio.to_thread(parse_pdf, file, params=params)


async def parse_image_async(file, params=None):
    return await asyncio.to_thread(parse_image, file, params=params)


async def process_file_to_markdown(file_path: str, params: dict | None = None) -> str:
    """
    将不同类型的文件转换为markdown格式

    Args:
        file_path: 文件路径
        params: 处理参数

    Returns:
        markdown格式内容
    """
    file_path_obj = Path(file_path)
    file_ext = file_path_obj.suffix.lower()

    if file_ext == ".pdf":
        # 使用 OCR 处理 PDF
        text = await parse_pdf_async(str(file_path_obj), params=params)
        return f"# {file_path_obj.name}\n\n{text}"

    elif file_ext in [".txt", ".md"]:
        # 直接读取文本文件
        with open(file_path_obj, encoding="utf-8") as f:
            content = f.read()
        return f"# {file_path_obj.name}\n\n{content}"

    elif file_ext in [".doc", ".docx"]:
        # 处理 Word 文档
        text = _extract_word_text(file_path_obj)
        return f"# {file_path_obj.name}\n\n{text}"

    elif file_ext in [".jpg", ".jpeg", ".png", ".bmp", ".tiff", ".tif"]:
        # 使用 OCR 处理图片
        text = await parse_image_async(str(file_path_obj), params=params)
        return f"# {file_path_obj.name}\n\n{text}"

    elif file_ext in [".html", ".htm"]:
        # 使用 BeautifulSoup 处理 HTML 文件
        from markdownify import markdownify as md

        with open(file_path_obj, encoding="utf-8") as f:
            content = f.read()
        text = md(content, heading_style="ATX")
        return f"# {file_path_obj.name}\n\n{text}"

    elif file_ext == ".csv":
        # 处理 CSV 文件
        import pandas as pd

        df = pd.read_csv(file_path_obj)
        # 将每一行数据与表头组合成独立的表格
        markdown_content = f"# {file_path_obj.name}\n\n"

        for index, row in df.iterrows():
            # 创建包含表头和当前行的小表格
            row_df = pd.DataFrame([row], columns=df.columns)
            markdown_table = row_df.to_markdown(index=False)
            markdown_content += f"{markdown_table}\n\n"

        return markdown_content.strip()

    elif file_ext in [".xls", ".xlsx"]:
        # 处理 Excel 文件
        import pandas as pd

        # 读取所有工作表
        excel_file = pd.ExcelFile(file_path_obj)
        markdown_content = f"# {file_path_obj.name}\n\n"

        for sheet_name in excel_file.sheet_names:
            df = pd.read_excel(file_path_obj, sheet_name=sheet_name)
            markdown_content += f"## {sheet_name}\n\n"

            # 将每一行数据与表头组合成独立的表格
            for index, row in df.iterrows():
                # 创建包含表头和当前行的小表格
                row_df = pd.DataFrame([row], columns=df.columns)
                markdown_table = row_df.to_markdown(index=False)
                markdown_content += f"{markdown_table}\n\n"

        return markdown_content.strip()

    elif file_ext == ".json":
        # 处理 JSON 文件
        import json

        with open(file_path_obj, encoding="utf-8") as f:
            data = json.load(f)
        # 将 JSON 数据格式化为 markdown 代码块
        json_str = json.dumps(data, ensure_ascii=False, indent=2)
        return f"# {file_path_obj.name}\n\n```json\n{json_str}\n```"

    else:
        # 尝试作为文本文件读取
        raise ValueError(f"Unsupported file type: {file_ext}")


async def process_url_to_markdown(url: str, params: dict | None = None) -> str:
    """
    将URL转换为markdown格式

    Args:
        url: URL地址
        params: 处理参数

    Returns:
        markdown格式内容
    """
    import requests
    from bs4 import BeautifulSoup

    try:
        response = requests.get(url, timeout=30)
        soup = BeautifulSoup(response.content, "html.parser")
        text_content = soup.get_text()
        return f"# {url}\n\n{text_content}"
    except Exception as e:
        logger.error(f"Failed to process URL {url}: {e}")
        return f"# {url}\n\nFailed to process URL: {e}"