326 lines
13 KiB
Python
326 lines
13 KiB
Python
"""
|
||
文件处理器
|
||
|
||
负责文件的后台 ETL 处理：加载、分块、向量化、存储。
|
||
支持两种知识库类型：DOCUMENT（向量检索）和 GRAPH（知识图谱）。
|
||
使用全局 WorkerPool 实现并发控制，最多 10 个文件并行处理。
|
||
"""
|
||
import asyncio
|
||
import logging
|
||
from pathlib import Path
|
||
from typing import List
|
||
|
||
from fastapi import BackgroundTasks
|
||
from sqlalchemy.ext.asyncio import AsyncSession
|
||
|
||
from app.db.models.knowledge_gen import KnowledgeBase, RagFile, FileStatus, RagType
|
||
from app.db.session import AsyncSessionLocal
|
||
from app.module.rag.infra.document import ingest_file_to_chunks
|
||
from app.module.rag.infra.embeddings import EmbeddingFactory
|
||
from app.module.rag.infra.task.worker_pool import get_global_pool
|
||
from app.module.rag.infra.vectorstore import VectorStoreFactory
|
||
from app.module.rag.repository import RagFileRepository, KnowledgeBaseRepository
|
||
from app.module.rag.schema.request import AddFilesReq
|
||
from app.module.rag.service.common import (
|
||
TextCleaner,
|
||
MetadataBuilder,
|
||
BatchProcessor,
|
||
get_file_path,
|
||
)
|
||
from app.module.system.service.common_service import get_model_by_id
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# xuhuanqing
|
||
def remove_utf8_surrogates(text: str) -> str:
|
||
"""移除字符串内的UTF‑16代理字符 \ud800‑\udfff，规避surrogates not allowed"""
|
||
if not text:
|
||
return text
|
||
# 用unicode转码，忽略代理
|
||
# return text.encode("utf‑16", errors="surrogatepass").decode("utf‑16")
|
||
return text.encode("utf‑8", errors="ignore").decode("utf‑8")
|
||
|
||
class FileProcessor:
|
||
|
||
def __init__(self):
|
||
self.worker_pool = get_global_pool(max_workers=10)
|
||
|
||
def start_background_processing(
|
||
self,
|
||
background_tasks: BackgroundTasks,
|
||
knowledge_base_id: str,
|
||
knowledge_base_name: str,
|
||
knowledge_base_type: str,
|
||
request_data: dict,
|
||
) -> None:
|
||
background_tasks.add_task(
|
||
self._process_files_background,
|
||
knowledge_base_id,
|
||
knowledge_base_name,
|
||
knowledge_base_type,
|
||
request_data,
|
||
)
|
||
logger.info("已注册后台任务: 知识库=%s, 类型=%s", knowledge_base_name, knowledge_base_type)
|
||
|
||
async def _process_files_background(
|
||
self,
|
||
knowledge_base_id: str,
|
||
knowledge_base_name: str,
|
||
knowledge_base_type: str,
|
||
request_data: dict,
|
||
) -> None:
|
||
async with AsyncSessionLocal() as db:
|
||
try:
|
||
kb_repo = KnowledgeBaseRepository(db)
|
||
file_repo = RagFileRepository(db)
|
||
|
||
knowledge_base = await kb_repo.get_by_id(knowledge_base_id)
|
||
if not knowledge_base:
|
||
logger.error("知识库不存在: %s", knowledge_base_id)
|
||
return
|
||
|
||
request = AddFilesReq.model_validate(request_data)
|
||
files = await file_repo.get_unprocessed_files(knowledge_base_id)
|
||
|
||
if not files:
|
||
logger.info("知识库 %s 没有待处理的文件", knowledge_base_name)
|
||
return
|
||
|
||
logger.info("开始处理 %d 个文件，知识库: %s, 类型: %s", len(files), knowledge_base_name, knowledge_base_type)
|
||
|
||
if knowledge_base_type == RagType.GRAPH.value:
|
||
await self._process_graph_files(db, files, knowledge_base)
|
||
else:
|
||
await self._process_document_files(files, knowledge_base, request)
|
||
|
||
logger.info("知识库 %s 文件处理完成", knowledge_base_name)
|
||
|
||
except Exception as e:
|
||
logger.exception("后台处理文件失败: %s", e)
|
||
finally:
|
||
await db.close()
|
||
|
||
async def _process_document_files(
|
||
self,
|
||
files: List[RagFile],
|
||
knowledge_base: KnowledgeBase,
|
||
request: AddFilesReq,
|
||
) -> None:
|
||
async def process_with_semaphore(rag_file: RagFile):
|
||
async with self.worker_pool.semaphore:
|
||
async with AsyncSessionLocal() as file_db:
|
||
try:
|
||
await self._process_single_file(file_db, rag_file, knowledge_base, request)
|
||
finally:
|
||
await file_db.close()
|
||
|
||
tasks = [process_with_semaphore(f) for f in files]
|
||
await asyncio.gather(*tasks, return_exceptions=True)
|
||
|
||
async def _process_graph_files(
|
||
self,
|
||
db: AsyncSession,
|
||
files: List[RagFile],
|
||
knowledge_base: KnowledgeBase,
|
||
) -> None:
|
||
from app.module.shared.common.document_loaders import load_documents
|
||
|
||
try:
|
||
rag_instance = await self._initialize_graph_rag(db, knowledge_base)
|
||
|
||
for rag_file in files:
|
||
await self._process_single_graph_file(db, rag_file, rag_instance, load_documents)
|
||
|
||
except Exception as e:
|
||
logger.exception("初始化知识图谱失败: %s", e)
|
||
for rag_file in files:
|
||
file_repo = RagFileRepository(db)
|
||
await self._mark_failed(db, file_repo, str(rag_file.id), f"知识图谱初始化失败: {str(e)}")
|
||
|
||
@staticmethod
|
||
async def _initialize_graph_rag(db: AsyncSession, knowledge_base: KnowledgeBase):
|
||
from app.module.rag.service.strategy.graph_strategy import GraphKnowledgeBaseStrategy
|
||
strategy = GraphKnowledgeBaseStrategy(db)
|
||
return await strategy._get_or_create_graph_rag(knowledge_base)
|
||
|
||
async def _process_single_graph_file(
|
||
self,
|
||
db: AsyncSession,
|
||
rag_file: RagFile,
|
||
rag_instance,
|
||
load_documents,
|
||
) -> None:
|
||
file_repo = RagFileRepository(db)
|
||
|
||
try:
|
||
await self._update_status(db, file_repo, str(rag_file.id), FileStatus.PROCESSING, 10)
|
||
await db.commit()
|
||
|
||
file_path = get_file_path(rag_file)
|
||
if not file_path or not Path(file_path).exists():
|
||
await self._mark_failed(db, file_repo, str(rag_file.id), "文件不存在")
|
||
return
|
||
|
||
documents = await asyncio.to_thread(load_documents, file_path)
|
||
if not documents:
|
||
await self._mark_failed(db, file_repo, str(rag_file.id), "文件解析失败，未生成文档")
|
||
return
|
||
|
||
await self._update_progress(db, file_repo, str(rag_file.id), 30)
|
||
await db.commit()
|
||
|
||
all_content = "\n\n".join(doc.page_content for doc in documents)
|
||
doc_id = str(rag_file.id)
|
||
logger.info("插入文档到知识图谱: %s, doc_id=%s, 文档数=%d", str(rag_file.file_name), doc_id, len(documents))
|
||
# xuhuanqing
|
||
all_content = remove_utf8_surrogates(all_content)
|
||
await rag_instance.ainsert(input=all_content, file_paths=[file_path], ids=doc_id)
|
||
|
||
doc_status_data = await rag_instance.doc_status.get_by_id(doc_id)
|
||
chunk_count = len(doc_status_data.get("chunks_list", [])) if doc_status_data else 0
|
||
await self._mark_success(db, file_repo, str(rag_file.id), chunk_count)
|
||
logger.info("文件 %s 知识图谱处理完成, 实际分块数: %d", str(rag_file.file_name), chunk_count)
|
||
|
||
except Exception as e:
|
||
logger.exception("文件 %s 知识图谱处理失败: %s", str(rag_file.file_name), e)
|
||
await self._mark_failed(db, file_repo, str(rag_file.id), str(e))
|
||
|
||
async def _process_single_file(
|
||
self,
|
||
db: AsyncSession,
|
||
rag_file: RagFile,
|
||
knowledge_base: KnowledgeBase,
|
||
request: AddFilesReq,
|
||
) -> None:
|
||
file_repo = RagFileRepository(db)
|
||
|
||
try:
|
||
await self._update_status(db, file_repo, rag_file.id, FileStatus.PROCESSING, 5)
|
||
await db.commit()
|
||
|
||
file_path = get_file_path(rag_file)
|
||
if not file_path or not Path(file_path).exists():
|
||
await self._mark_failed(db, file_repo, rag_file.id, "文件不存在")
|
||
return
|
||
|
||
base_metadata = MetadataBuilder.build_chunk_metadata(rag_file, knowledge_base)
|
||
chunks = await ingest_file_to_chunks(
|
||
file_path,
|
||
process_type=request.process_type,
|
||
chunk_size=request.chunk_size,
|
||
overlap_size=request.overlap_size,
|
||
delimiter=request.delimiter,
|
||
**base_metadata,
|
||
)
|
||
|
||
if not chunks:
|
||
await self._mark_failed(db, file_repo, rag_file.id, "文档解析后未生成任何分块")
|
||
return
|
||
|
||
logger.info("文件 %s 分块完成，共 %d 个分块", rag_file.file_name, len(chunks))
|
||
|
||
valid_chunks = self._filter_and_clean_chunks(chunks, rag_file)
|
||
if not valid_chunks:
|
||
await self._mark_failed(db, file_repo, rag_file.id, "文件没有有效的分块内容")
|
||
return
|
||
|
||
embedding = await self._get_embeddings(db, knowledge_base)
|
||
vectorstore = VectorStoreFactory.create(
|
||
collection_name=str(knowledge_base.name),
|
||
embedding=embedding,
|
||
)
|
||
|
||
await self._update_progress(db, file_repo, rag_file.id, 60)
|
||
await db.commit()
|
||
|
||
MetadataBuilder.add_to_chunks(valid_chunks, {
|
||
"rag_file_id": str(rag_file.id),
|
||
"original_file_id": str(rag_file.file_id),
|
||
"knowledge_base_id": str(knowledge_base.id),
|
||
})
|
||
|
||
await BatchProcessor.store_in_batches(vectorstore, valid_chunks)
|
||
|
||
await self._mark_success(db, file_repo, rag_file.id, len(valid_chunks))
|
||
logger.info("文件 %s ETL 处理完成", rag_file.file_name)
|
||
|
||
except Exception as e:
|
||
logger.exception("文件 %s 处理失败: %s", rag_file.file_name, e)
|
||
await self._mark_failed(db, file_repo, rag_file.id, str(e))
|
||
|
||
@staticmethod
|
||
def _filter_and_clean_chunks(chunks: list, rag_file: RagFile) -> list:
|
||
valid_chunks = []
|
||
for chunk in chunks:
|
||
cleaned_text = TextCleaner.clean(chunk.text)
|
||
|
||
if cleaned_text and TextCleaner.has_printable_content(cleaned_text):
|
||
chunk.text = cleaned_text
|
||
valid_chunks.append(chunk)
|
||
else:
|
||
logger.warning(
|
||
"跳过无效分块: rag_file_id=%s, chunk_index=%s",
|
||
rag_file.id,
|
||
chunk.metadata.get("chunk_index")
|
||
)
|
||
|
||
logger.info("文件 %s 有效分块数量: %d / %d", rag_file.file_name, len(valid_chunks), len(chunks))
|
||
return valid_chunks
|
||
|
||
@staticmethod
|
||
async def _get_embeddings(db: AsyncSession, knowledge_base: KnowledgeBase):
|
||
embedding_entity = await get_model_by_id(db, str(knowledge_base.embedding_model))
|
||
if not embedding_entity:
|
||
raise ValueError(f"嵌入模型不存在: {knowledge_base.embedding_model}")
|
||
|
||
return EmbeddingFactory.create_embeddings(
|
||
model_name=str(embedding_entity.model_name),
|
||
base_url=getattr(embedding_entity, "base_url", None),
|
||
api_key=getattr(embedding_entity, "api_key", None),
|
||
)
|
||
|
||
@staticmethod
|
||
async def _update_status(
|
||
db: AsyncSession,
|
||
file_repo: RagFileRepository,
|
||
rag_file_id: str,
|
||
status: FileStatus,
|
||
progress: int = 0,
|
||
) -> None:
|
||
await file_repo.update_status(rag_file_id, status)
|
||
await file_repo.update_progress(rag_file_id, progress)
|
||
await db.flush()
|
||
|
||
@staticmethod
|
||
async def _update_progress(
|
||
db: AsyncSession,
|
||
file_repo: RagFileRepository,
|
||
rag_file_id: str,
|
||
progress: int,
|
||
) -> None:
|
||
await file_repo.update_progress(rag_file_id, progress)
|
||
await db.flush()
|
||
|
||
@staticmethod
|
||
async def _mark_success(
|
||
db: AsyncSession,
|
||
file_repo: RagFileRepository,
|
||
rag_file_id: str,
|
||
chunk_count: int,
|
||
) -> None:
|
||
await file_repo.update_chunk_count(rag_file_id, chunk_count)
|
||
await file_repo.update_status(rag_file_id, FileStatus.PROCESSED)
|
||
await file_repo.update_progress(rag_file_id, 100)
|
||
await db.commit()
|
||
|
||
@staticmethod
|
||
async def _mark_failed(
|
||
db: AsyncSession,
|
||
file_repo: RagFileRepository,
|
||
rag_file_id: str,
|
||
err_msg: str,
|
||
) -> None:
|
||
logger.error("文件处理失败: %s", err_msg)
|
||
await file_repo.update_status(rag_file_id, FileStatus.PROCESS_FAILED, err_msg=err_msg)
|
||
await db.commit()
|