489 lines
18 KiB
Python
489 lines
18 KiB
Python
"""
|
|
知识库业务服务
|
|
|
|
实现知识库的 CRUD 操作和文件管理。
|
|
对应 Java: com.deepdata.rag.indexer.application.KnowledgeBaseService
|
|
"""
|
|
import logging
|
|
import uuid
|
|
from typing import List, Tuple, Optional
|
|
|
|
from fastapi import BackgroundTasks
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.exception import BusinessError, ErrorCodes
|
|
from app.db.models.dataset_management import DatasetFiles
|
|
from app.db.models.knowledge_gen import KnowledgeBase, RagFile, FileStatus, RagType
|
|
from app.db.models.models import Models
|
|
from app.module.rag.infra.embeddings import EmbeddingFactory
|
|
from app.module.rag.infra.vectorstore import (
|
|
drop_collection,
|
|
rename_collection,
|
|
delete_chunks_by_rag_file_ids,
|
|
update_chunk_by_id,
|
|
delete_chunk_by_id,
|
|
)
|
|
from app.module.rag.repository import KnowledgeBaseRepository, RagFileRepository
|
|
from app.module.rag.schema.request import (
|
|
KnowledgeBaseCreateReq,
|
|
KnowledgeBaseUpdateReq,
|
|
KnowledgeBaseQueryReq,
|
|
AddFilesReq,
|
|
DeleteFilesReq,
|
|
RagFileReq,
|
|
)
|
|
from app.module.rag.schema.response import KnowledgeBaseResp, PagedResponse, RagFileResp, ModelConfig
|
|
from app.module.rag.service.file_processor import FileProcessor
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class KnowledgeBaseService:
|
|
"""知识库业务服务类
|
|
|
|
功能：
|
|
1. 知识库 CRUD 操作
|
|
2. 文件管理（添加、删除、查询）
|
|
"""
|
|
|
|
def __init__(self, db: AsyncSession):
|
|
"""初始化服务
|
|
|
|
Args:
|
|
db: 数据库异步 session
|
|
"""
|
|
self.db = db
|
|
self.kb_repo = KnowledgeBaseRepository(db)
|
|
self.file_repo = RagFileRepository(db)
|
|
self.file_processor = FileProcessor()
|
|
|
|
# ==================== 知识库 CRUD ====================
|
|
|
|
async def create(self, request: KnowledgeBaseCreateReq) -> str:
|
|
"""创建知识库
|
|
|
|
Args:
|
|
request: 创建请求
|
|
|
|
Returns:
|
|
知识库 ID
|
|
"""
|
|
knowledge_base = KnowledgeBase(
|
|
id=str(uuid.uuid4()),
|
|
name=request.name,
|
|
description=request.description,
|
|
type=request.type,
|
|
embedding_model=request.embedding_model,
|
|
chat_model=request.chat_model,
|
|
)
|
|
|
|
knowledge_base = await self.kb_repo.create(knowledge_base)
|
|
await self.db.commit()
|
|
|
|
logger.info("成功创建知识库: %s", request.name)
|
|
return knowledge_base.id
|
|
|
|
async def update(self, knowledge_base_id: str, request: KnowledgeBaseUpdateReq) -> None:
|
|
"""更新知识库
|
|
|
|
Args:
|
|
knowledge_base_id: 知识库 ID
|
|
request: 更新请求
|
|
"""
|
|
knowledge_base = await self.kb_repo.get_by_id(knowledge_base_id)
|
|
if not knowledge_base:
|
|
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
|
|
|
old_name = str(knowledge_base.name)
|
|
new_name = request.name
|
|
kb_type = knowledge_base.type
|
|
|
|
if new_name and new_name != old_name:
|
|
if await self.kb_repo.exists_by_name(new_name, exclude_id=knowledge_base_id):
|
|
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_ALREADY_EXISTS, data={"name": new_name})
|
|
|
|
knowledge_base.name = request.name
|
|
knowledge_base.description = request.description
|
|
|
|
await self.kb_repo.update(knowledge_base)
|
|
|
|
if old_name != new_name:
|
|
try:
|
|
if kb_type == RagType.DOCUMENT.value:
|
|
rename_collection(old_name, new_name)
|
|
elif kb_type == RagType.GRAPH.value:
|
|
from app.module.rag.service.strategy.graph_strategy import GraphKnowledgeBaseStrategy
|
|
GraphKnowledgeBaseStrategy.rename_workspace(old_name, new_name)
|
|
GraphKnowledgeBaseStrategy.clear_cache(old_name)
|
|
except BusinessError:
|
|
await self.db.rollback()
|
|
raise
|
|
|
|
await self.db.commit()
|
|
|
|
async def delete(self, knowledge_base_id: str) -> None:
|
|
"""删除知识库
|
|
|
|
Args:
|
|
knowledge_base_id: 知识库 ID
|
|
"""
|
|
knowledge_base = await self.kb_repo.get_by_id(knowledge_base_id)
|
|
if not knowledge_base:
|
|
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
|
|
|
kb_name = str(knowledge_base.name)
|
|
kb_type = knowledge_base.type
|
|
|
|
await self.file_repo.delete_by_knowledge_base(knowledge_base_id)
|
|
await self.kb_repo.delete(knowledge_base_id)
|
|
|
|
if kb_type == RagType.DOCUMENT.value:
|
|
try:
|
|
drop_collection(kb_name)
|
|
except Exception as e:
|
|
logger.error("删除 Milvus 集合失败: %s", e)
|
|
elif kb_type == RagType.GRAPH.value:
|
|
try:
|
|
from app.module.rag.service.strategy.graph_strategy import GraphKnowledgeBaseStrategy
|
|
import shutil
|
|
from pathlib import Path
|
|
from app.core.config import settings
|
|
workspace_path = Path(settings.rag_storage_dir) / kb_name
|
|
if workspace_path.exists():
|
|
shutil.rmtree(workspace_path)
|
|
logger.info("已删除知识图谱 workspace: %s", kb_name)
|
|
GraphKnowledgeBaseStrategy.clear_cache(kb_name)
|
|
except Exception as e:
|
|
logger.error("删除知识图谱 workspace 失败: %s", e)
|
|
|
|
await self.db.commit()
|
|
|
|
async def get_by_id(self, knowledge_base_id: str) -> KnowledgeBaseResp:
|
|
"""获取知识库详情"""
|
|
knowledge_base = await self.kb_repo.get_by_id(knowledge_base_id)
|
|
if not knowledge_base:
|
|
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
|
|
|
file_count = await self.file_repo.count_by_knowledge_base(knowledge_base_id)
|
|
chunk_count = await self.file_repo.count_chunks_by_knowledge_base(knowledge_base_id)
|
|
|
|
data = self._kb_to_dict(knowledge_base)
|
|
data.update({
|
|
"file_count": file_count,
|
|
"chunk_count": chunk_count,
|
|
"embedding": await self._get_model_config(knowledge_base.embedding_model),
|
|
"chat": await self._get_model_config(knowledge_base.chat_model),
|
|
})
|
|
return KnowledgeBaseResp(**data)
|
|
|
|
@staticmethod
|
|
def _kb_to_dict(kb: KnowledgeBase) -> dict:
|
|
"""知识库实体转字典"""
|
|
return {
|
|
"id": kb.id,
|
|
"name": kb.name,
|
|
"description": kb.description,
|
|
"type": kb.type,
|
|
"embedding_model": kb.embedding_model,
|
|
"chat_model": kb.chat_model,
|
|
"created_at": kb.created_at,
|
|
"updated_at": kb.updated_at,
|
|
"created_by": kb.created_by,
|
|
"updated_by": kb.updated_by,
|
|
}
|
|
|
|
async def _get_model_config(self, model_id: Optional[str]) -> Optional[ModelConfig]:
|
|
"""获取模型配置"""
|
|
if not model_id:
|
|
return None
|
|
|
|
result = await self.db.execute(
|
|
select(Models).where(
|
|
Models.id == model_id,
|
|
(Models.is_deleted == False) | (Models.is_deleted.is_(None))
|
|
)
|
|
)
|
|
model = result.scalar_one_or_none()
|
|
return ModelConfig.model_validate(model) if model else None
|
|
|
|
async def list(self, request: KnowledgeBaseQueryReq) -> PagedResponse:
|
|
"""分页查询知识库列表"""
|
|
items, total = await self.kb_repo.list(
|
|
keyword=request.keyword,
|
|
rag_type=request.type,
|
|
page=request.page,
|
|
page_size=request.page_size,
|
|
)
|
|
|
|
responses = []
|
|
for item in items:
|
|
file_count = await self.file_repo.count_by_knowledge_base(item.id)
|
|
chunk_count = await self.file_repo.count_chunks_by_knowledge_base(item.id)
|
|
data = self._kb_to_dict(item)
|
|
data.update({
|
|
"file_count": file_count,
|
|
"chunk_count": chunk_count,
|
|
})
|
|
responses.append(KnowledgeBaseResp(**data))
|
|
|
|
return PagedResponse.create(
|
|
content=responses,
|
|
total_elements=total,
|
|
page=request.page,
|
|
size=request.page_size,
|
|
)
|
|
|
|
# ==================== 文件管理 ====================
|
|
|
|
async def add_files(
|
|
self,
|
|
request: AddFilesReq,
|
|
background_tasks: BackgroundTasks = None,
|
|
) -> dict:
|
|
"""添加文件到知识库
|
|
|
|
验证知识库、创建文件记录、启动后台处理。
|
|
数据库提交后立即返回，不等待文件处理完成。
|
|
|
|
Args:
|
|
request: 添加文件请求
|
|
background_tasks: FastAPI 后台任务
|
|
|
|
Returns:
|
|
包含成功和跳过文件数量的字典
|
|
"""
|
|
knowledge_base = await self.kb_repo.get_by_id(request.knowledge_base_id)
|
|
if not knowledge_base:
|
|
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
|
|
|
rag_files, skipped_file_ids = await self._create_rag_files(request)
|
|
|
|
await self.db.commit()
|
|
|
|
if rag_files and background_tasks:
|
|
kb_type = knowledge_base.type if knowledge_base.type else "DOCUMENT"
|
|
self.file_processor.start_background_processing(
|
|
background_tasks=background_tasks,
|
|
knowledge_base_id=str(knowledge_base.id),
|
|
knowledge_base_name=str(knowledge_base.name),
|
|
knowledge_base_type=str(kb_type),
|
|
request_data=request.model_dump(),
|
|
)
|
|
|
|
return {
|
|
"success_count": len(rag_files),
|
|
"skipped_count": len(skipped_file_ids),
|
|
"skipped_file_ids": skipped_file_ids,
|
|
}
|
|
|
|
async def _create_rag_files(self, request: AddFilesReq) -> Tuple[List[RagFile], List[str]]:
|
|
"""创建 RAG 文件记录"""
|
|
if not request.files:
|
|
raise BusinessError(ErrorCodes.BAD_REQUEST, "文件列表不能为空")
|
|
|
|
rag_files = []
|
|
skipped_file_ids = []
|
|
|
|
for file_info in request.files:
|
|
try:
|
|
result = await self.db.execute(
|
|
select(DatasetFiles).where(DatasetFiles.id == file_info.id)
|
|
)
|
|
dataset_file = result.scalar_one_or_none()
|
|
|
|
if not dataset_file:
|
|
logger.warning("文件不存在，跳过: file_id=%s", file_info.id)
|
|
skipped_file_ids.append(file_info.id)
|
|
continue
|
|
|
|
rag_file = RagFile(
|
|
id=str(uuid.uuid4()),
|
|
knowledge_base_id=request.knowledge_base_id,
|
|
file_name=dataset_file.file_name,
|
|
file_id=file_info.id,
|
|
file_metadata={
|
|
"process_type": request.process_type.value,
|
|
"dataset_id": dataset_file.dataset_id,
|
|
"file_path": dataset_file.file_path,
|
|
},
|
|
status=FileStatus.UNPROCESSED,
|
|
)
|
|
rag_files.append(rag_file)
|
|
|
|
except Exception as e:
|
|
logger.error("处理文件信息失败: file_id=%s, error=%s", file_info.id, e)
|
|
skipped_file_ids.append(file_info.id)
|
|
|
|
if rag_files:
|
|
await self.file_repo.batch_create(rag_files)
|
|
logger.info("成功添加 %d 个文件到知识库", len(rag_files))
|
|
|
|
return rag_files, skipped_file_ids
|
|
|
|
async def list_files(self, knowledge_base_id: str, request: RagFileReq) -> PagedResponse:
|
|
"""获取知识库文件列表"""
|
|
if not await self.kb_repo.get_by_id(knowledge_base_id):
|
|
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
|
|
|
items, total = await self.file_repo.list_by_knowledge_base(
|
|
knowledge_base_id=knowledge_base_id,
|
|
keyword=request.keyword,
|
|
status=request.status,
|
|
page=request.page,
|
|
page_size=request.page_size,
|
|
)
|
|
|
|
responses = [RagFileResp.model_validate(item) for item in items]
|
|
|
|
return PagedResponse.create(
|
|
content=responses,
|
|
total_elements=total,
|
|
page=request.page,
|
|
size=request.page_size,
|
|
)
|
|
|
|
async def delete_files(self, knowledge_base_id: str, request: DeleteFilesReq) -> None:
|
|
"""删除知识库文件
|
|
|
|
Args:
|
|
knowledge_base_id: 知识库 ID
|
|
request: 删除文件请求
|
|
"""
|
|
knowledge_base = await self.kb_repo.get_by_id(knowledge_base_id)
|
|
if not knowledge_base:
|
|
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
|
|
|
if not request.file_ids:
|
|
raise BusinessError(ErrorCodes.BAD_REQUEST, "文件ID列表不能为空")
|
|
|
|
kb_type = knowledge_base.type
|
|
kb_name = str(knowledge_base.name)
|
|
|
|
rag_files = []
|
|
for file_id in request.file_ids:
|
|
rag_file = await self.file_repo.get_by_id(file_id)
|
|
if rag_file:
|
|
rag_files.append(rag_file)
|
|
|
|
if rag_files:
|
|
if kb_type == RagType.DOCUMENT.value:
|
|
try:
|
|
delete_chunks_by_rag_file_ids(
|
|
kb_name,
|
|
[r.id for r in rag_files],
|
|
)
|
|
except Exception as e:
|
|
logger.error("删除 Milvus 数据失败: %s", e)
|
|
elif kb_type == RagType.GRAPH.value:
|
|
try:
|
|
from app.module.rag.service.strategy.graph_strategy import GraphKnowledgeBaseStrategy
|
|
strategy = GraphKnowledgeBaseStrategy(self.db)
|
|
rag_instance = await strategy._get_or_create_graph_rag(knowledge_base)
|
|
for rag_file in rag_files:
|
|
doc_id = str(rag_file.id)
|
|
await rag_instance.adelete_by_doc_id(doc_id)
|
|
logger.info("已从知识图谱删除文件: %s, doc_id=%s", rag_file.file_name, doc_id)
|
|
except Exception as e:
|
|
logger.error("删除知识图谱数据失败: %s", e)
|
|
|
|
for file_id in request.file_ids:
|
|
try:
|
|
await self.file_repo.delete(file_id)
|
|
except Exception as e:
|
|
logger.error("删除数据库记录失败: %s, error=%s", file_id, e)
|
|
|
|
await self.db.commit()
|
|
logger.info("成功删除 %d 个文件", len(rag_files))
|
|
|
|
async def update_chunk(
|
|
self,
|
|
knowledge_base_id: str,
|
|
chunk_id: str,
|
|
text: str,
|
|
metadata: dict = None,
|
|
) -> None:
|
|
"""更新指定分块的文本和元数据
|
|
|
|
Args:
|
|
knowledge_base_id: 知识库 ID
|
|
chunk_id: 分块 ID
|
|
text: 新的文本内容
|
|
metadata: 新的元数据（可选）
|
|
"""
|
|
knowledge_base = await self.kb_repo.get_by_id(knowledge_base_id)
|
|
if not knowledge_base:
|
|
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
|
|
|
if knowledge_base.type != RagType.DOCUMENT.value:
|
|
raise BusinessError(
|
|
ErrorCodes.RAG_INVALID_REQUEST,
|
|
f"知识库类型 {knowledge_base.type} 不支持分块更新"
|
|
)
|
|
|
|
from app.module.system.service.common_service import get_model_by_id
|
|
import asyncio
|
|
|
|
embedding_entity = await get_model_by_id(self.db, knowledge_base.embedding_model)
|
|
if not embedding_entity:
|
|
raise BusinessError(ErrorCodes.RAG_MODEL_NOT_FOUND)
|
|
|
|
embedding = 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),
|
|
)
|
|
|
|
await asyncio.to_thread(
|
|
update_chunk_by_id,
|
|
collection_name=str(knowledge_base.name),
|
|
chunk_id=chunk_id,
|
|
text=text,
|
|
metadata=metadata,
|
|
embedding_instance=embedding,
|
|
)
|
|
|
|
logger.info(
|
|
"成功更新分块: kb=%s chunk_id=%s",
|
|
knowledge_base_id, chunk_id
|
|
)
|
|
|
|
async def delete_chunk(
|
|
self,
|
|
knowledge_base_id: str,
|
|
chunk_id: str,
|
|
) -> None:
|
|
"""删除指定分块
|
|
|
|
Args:
|
|
knowledge_base_id: 知识库 ID
|
|
chunk_id: 分块 ID
|
|
"""
|
|
knowledge_base = await self.kb_repo.get_by_id(knowledge_base_id)
|
|
if not knowledge_base:
|
|
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
|
|
|
if knowledge_base.type != RagType.DOCUMENT.value:
|
|
raise BusinessError(
|
|
ErrorCodes.RAG_INVALID_REQUEST,
|
|
f"知识库类型 {knowledge_base.type} 不支持分块删除"
|
|
)
|
|
|
|
import asyncio
|
|
rag_file_id = await asyncio.to_thread(
|
|
delete_chunk_by_id,
|
|
collection_name=str(knowledge_base.name),
|
|
chunk_id=chunk_id,
|
|
)
|
|
|
|
if rag_file_id:
|
|
rag_file = await self.file_repo.get_by_id(rag_file_id)
|
|
if rag_file and rag_file.chunk_count and rag_file.chunk_count > 0:
|
|
rag_file.chunk_count = rag_file.chunk_count - 1
|
|
await self.db.commit()
|
|
|
|
logger.info(
|
|
"成功删除分块: kb=%s chunk_id=%s",
|
|
knowledge_base_id, chunk_id
|
|
)
|