128 lines
5.2 KiB
Python
128 lines
5.2 KiB
Python
"""Splitter 协调器，编排文件解析、AI 分块和输出的完整流程"""
|
||
|
||
import os
|
||
from typing import Optional
|
||
|
||
from api_client import ApiClient
|
||
from chunker import AIChunker
|
||
from exceptions import ApiError, ParseError, UnsupportedFormatError
|
||
from parsers.base import ParserRegistry
|
||
from parsers.csv_parser import CsvParser
|
||
from parsers.doc_parser import DocParser
|
||
from parsers.html_parser import HtmlParser
|
||
from parsers.image_parser import ImageParser
|
||
from parsers.legacy_doc_parser import LegacyDocParser
|
||
from parsers.pdf_parser import PdfParser
|
||
from parsers.text_parser import TextParser
|
||
from parsers.xls_parser import XlsParser
|
||
from parsers.xlsx_parser import XlsxParser
|
||
from prompts import detect_content_type, CONTENT_TYPE_IMAGE, CONTENT_TYPE_TABLE
|
||
from writer import JsonWriter, MarkdownWriter
|
||
|
||
# 表格类扩展名 — Coze 有专门的表格知识库，不需要 AI 分块
|
||
TABLE_EXTENSIONS = {".xlsx", ".xls", ".csv"}
|
||
|
||
|
||
class Splitter:
|
||
"""主协调器，编排文件解析、AI 分块和输出的完整流程"""
|
||
|
||
def __init__(
|
||
self,
|
||
api_key: str,
|
||
delimiter: str = "---",
|
||
pre_split_size: Optional[int] = None,
|
||
vision_prompt: Optional[str] = None,
|
||
output_format: str = "markdown",
|
||
):
|
||
self._api_client = ApiClient(api_key=api_key)
|
||
self._delimiter = delimiter
|
||
|
||
# 创建并注册所有解析器
|
||
self._registry = ParserRegistry()
|
||
self._registry.register(TextParser())
|
||
self._registry.register(CsvParser())
|
||
self._registry.register(HtmlParser())
|
||
self._registry.register(PdfParser())
|
||
self._registry.register(DocParser())
|
||
self._registry.register(LegacyDocParser())
|
||
self._registry.register(XlsxParser())
|
||
self._registry.register(XlsParser())
|
||
self._registry.register(ImageParser(self._api_client, vision_prompt=vision_prompt))
|
||
|
||
self._chunker = AIChunker(self._api_client, delimiter, pre_split_size=pre_split_size)
|
||
|
||
# 根据输出格式选择写入器
|
||
if output_format == "json":
|
||
self._writer = JsonWriter()
|
||
else:
|
||
self._writer = MarkdownWriter()
|
||
|
||
def process(self, input_path: str, output_path: str) -> None:
|
||
"""
|
||
执行完整的处理流程：
|
||
1. 验证输入文件存在
|
||
2. 根据文件类型选择解析器
|
||
3. 解析文件获取文本
|
||
4. AI 语义分块
|
||
5. 写入输出
|
||
6. 输出处理摘要
|
||
"""
|
||
# 1. 验证文件存在
|
||
if not os.path.exists(input_path):
|
||
raise FileNotFoundError(f"输入文件不存在: {input_path}")
|
||
|
||
file_name = os.path.basename(input_path)
|
||
file_size = os.path.getsize(input_path)
|
||
size_str = f"{file_size / 1024:.1f}KB" if file_size < 1024 * 1024 else f"{file_size / 1024 / 1024:.1f}MB"
|
||
|
||
# 表格类文件：Coze 有专门的表格知识库，不需要 AI 分块
|
||
ext = os.path.splitext(input_path)[1].lower()
|
||
if ext in TABLE_EXTENSIONS:
|
||
print(f" ⚠️ 表格文件 {file_name} 请直接上传到 Coze「表格知识库」，无需分块处理")
|
||
print(f" 批量模式（-b）会自动将表格文件复制到 tables/ 文件夹")
|
||
return
|
||
|
||
# 2. 选择解析器
|
||
parser = self._registry.get_parser(input_path)
|
||
|
||
# 3. 解析文件
|
||
print(f" [1/4] 解析文件: {file_name} ({size_str})")
|
||
text = parser.parse(input_path)
|
||
print(f" 提取文本: {len(text)} 字符")
|
||
|
||
# 3.5 检测内容类型
|
||
ext = os.path.splitext(input_path)[1].lower()
|
||
content_type = detect_content_type(ext, text)
|
||
type_labels = {
|
||
"document": "文档", "table": "表格", "qa": "问答", "image": "图片",
|
||
}
|
||
print(f" 内容类型: {type_labels.get(content_type, content_type)}")
|
||
|
||
# 4. AI 分块（图片类跳过 API 调用）
|
||
if content_type == CONTENT_TYPE_IMAGE:
|
||
print(f" [2/4] 图片内容，跳过 AI 分块")
|
||
else:
|
||
print(f" [2/4] AI 语义分块中（{type_labels.get(content_type, '通用')}策略）...")
|
||
|
||
def progress_callback(current: int, total: int) -> None:
|
||
print(f" 分块进度: {current}/{total}")
|
||
|
||
chunks = self._chunker.chunk(text, content_type=content_type, source_file=file_name, on_progress=progress_callback)
|
||
|
||
# 5. 写入输出
|
||
print(f" [3/4] 写入输出: {output_path}")
|
||
source_file = os.path.basename(input_path)
|
||
self._writer.write(chunks, output_path, source_file, self._delimiter)
|
||
|
||
# 6. 输出摘要
|
||
chunk_sizes = [len(c.content) for c in chunks]
|
||
avg_size = sum(chunk_sizes) // max(len(chunk_sizes), 1)
|
||
max_size = max(chunk_sizes) if chunk_sizes else 0
|
||
min_size = min(chunk_sizes) if chunk_sizes else 0
|
||
print(f" [4/4] 完成! 共 {len(chunks)} 个分块 (平均 {avg_size} 字, 最大 {max_size} 字, 最小 {min_size} 字)")
|
||
|
||
# 检查是否有超大分块（可能不适合知识库平台）
|
||
oversized = [i + 1 for i, s in enumerate(chunk_sizes) if s > 800]
|
||
if oversized:
|
||
print(f" ⚠️ 分块 {oversized} 超过 800 字，上传知识库时建议检查")
|