397 lines
12 KiB
Python
397 lines
12 KiB
Python
import asyncio
|
||
import os
|
||
from pathlib import Path
|
||
|
||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||
from langchain_community.document_loaders import (
|
||
CSVLoader,
|
||
JSONLoader,
|
||
PyPDFLoader,
|
||
TextLoader,
|
||
UnstructuredHTMLLoader,
|
||
UnstructuredMarkdownLoader,
|
||
UnstructuredWordDocumentLoader,
|
||
)
|
||
|
||
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}"
|