deepdata: core product code (backend, frontend, runtime, deepflow module)
This commit is contained in:
figmar
@@ -0,0 +1,16 @@
|
||||
"""
|
||||
公共工具模块
|
||||
|
||||
提供文本清理、元数据构建、批量处理等工具类。
|
||||
"""
|
||||
from .text_cleaner import TextCleaner
|
||||
from .metadata_builder import MetadataBuilder
|
||||
from .batch_processor import BatchProcessor
|
||||
from .file_utils import get_file_path
|
||||
|
||||
__all__ = [
|
||||
"TextCleaner",
|
||||
"MetadataBuilder",
|
||||
"BatchProcessor",
|
||||
"get_file_path",
|
||||
]
|
||||
@@ -0,0 +1,94 @@
|
||||
"""
|
||||
批量处理工具
|
||||
|
||||
提供分批处理文档分块的工具方法。
|
||||
"""
|
||||
import uuid
|
||||
import logging
|
||||
from typing import List
|
||||
from app.module.rag.infra.vectorstore import chunks_to_documents
|
||||
from app.module.rag.infra.document.types import DocumentChunk
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BatchProcessor:
|
||||
"""批量处理工具类"""
|
||||
|
||||
DEFAULT_BATCH_SIZE: int = 20
|
||||
|
||||
@staticmethod
|
||||
async def store_in_batches(
|
||||
vectorstore,
|
||||
chunks: List[DocumentChunk],
|
||||
batch_size: int = DEFAULT_BATCH_SIZE
|
||||
) -> int:
|
||||
"""分批存储分块到向量数据库
|
||||
|
||||
Args:
|
||||
vectorstore: 向量存储实例
|
||||
chunks: 分块列表
|
||||
batch_size: 批次大小
|
||||
|
||||
Returns:
|
||||
成功存储的分块数量
|
||||
"""
|
||||
if not chunks:
|
||||
return 0
|
||||
|
||||
batch_size = batch_size or BatchProcessor.DEFAULT_BATCH_SIZE
|
||||
total_chunks = len(chunks)
|
||||
stored_count = 0
|
||||
|
||||
for batch_start in range(0, total_chunks, batch_size):
|
||||
batch_end = min(batch_start + batch_size, total_chunks)
|
||||
batch_chunks = chunks[batch_start:batch_end]
|
||||
|
||||
await BatchProcessor._store_single_batch(
|
||||
vectorstore,
|
||||
batch_chunks,
|
||||
batch_start,
|
||||
batch_end,
|
||||
total_chunks
|
||||
)
|
||||
stored_count += len(batch_chunks)
|
||||
|
||||
logger.info("批量存储完成，总数量: %d", stored_count)
|
||||
return stored_count
|
||||
|
||||
@staticmethod
|
||||
async def _store_single_batch(
|
||||
vectorstore,
|
||||
batch_chunks: List[DocumentChunk],
|
||||
batch_start: int,
|
||||
batch_end: int,
|
||||
total_chunks: int,
|
||||
) -> None:
|
||||
"""存储单个批次
|
||||
|
||||
Args:
|
||||
vectorstore: 向量存储实例
|
||||
batch_chunks: 批次分块列表
|
||||
batch_start: 批次起始索引
|
||||
batch_end: 批次结束索引
|
||||
total_chunks: 总分块数
|
||||
"""
|
||||
logger.info("处理分块批次 %d-%d / %d", batch_start + 1, batch_end, total_chunks)
|
||||
|
||||
ids = [str(uuid.uuid4()) for _ in batch_chunks]
|
||||
documents, doc_ids = chunks_to_documents(batch_chunks, ids=ids)
|
||||
|
||||
try:
|
||||
# 使用异步方法避免阻塞事件循环
|
||||
await vectorstore.aadd_documents(documents=documents, ids=doc_ids)
|
||||
logger.info("批次 %d-%d 存储成功", batch_start + 1, batch_end)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"批次 %d-%d 存储失败: %s\n第一个文档内容: %.200s",
|
||||
batch_start + 1,
|
||||
batch_end,
|
||||
str(e),
|
||||
documents[0].page_content if documents else "N/A"
|
||||
)
|
||||
raise
|
||||
@@ -0,0 +1,25 @@
|
||||
"""
|
||||
文件工具类
|
||||
|
||||
提供文件路径提取等通用工具方法。
|
||||
"""
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
|
||||
def get_file_path(rag_file) -> Optional[str]:
|
||||
"""获取 RAG 文件的绝对路径
|
||||
|
||||
Args:
|
||||
rag_file: RAG 文件实体（需要有 file_metadata 属性）
|
||||
|
||||
Returns:
|
||||
文件绝对路径，不存在返回 None
|
||||
"""
|
||||
if not rag_file.file_metadata:
|
||||
return None
|
||||
|
||||
file_path = rag_file.file_metadata.get("file_path")
|
||||
if file_path:
|
||||
return str(Path(file_path).absolute())
|
||||
return None
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
元数据构建工具
|
||||
|
||||
提供构建和管理分块元数据的工具方法。
|
||||
"""
|
||||
from typing import Dict, Any
|
||||
from app.db.models.knowledge_gen import KnowledgeBase, RagFile
|
||||
|
||||
|
||||
class MetadataBuilder:
|
||||
"""元数据构建工具类"""
|
||||
|
||||
@staticmethod
|
||||
def build_chunk_metadata(rag_file: RagFile, knowledge_base: KnowledgeBase) -> Dict[str, Any]:
|
||||
"""构建分块基础元数据
|
||||
|
||||
Args:
|
||||
rag_file: RAG 文件实体
|
||||
knowledge_base: 知识库实体
|
||||
|
||||
Returns:
|
||||
元数据字典
|
||||
"""
|
||||
return {
|
||||
"rag_file_id": str(rag_file.id),
|
||||
"original_file_id": str(rag_file.file_id),
|
||||
"knowledge_base_id": str(knowledge_base.id),
|
||||
"file_name": rag_file.file_name,
|
||||
"knowledge_base_name": knowledge_base.name,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def add_to_chunks(chunks: list, base_metadata: Dict[str, Any]) -> None:
|
||||
"""为基础元数据添加到分块列表
|
||||
|
||||
Args:
|
||||
chunks: 分块列表
|
||||
base_metadata: 基础元数据
|
||||
"""
|
||||
for chunk in chunks:
|
||||
if hasattr(chunk, 'metadata') and isinstance(chunk.metadata, dict):
|
||||
chunk.metadata.update(base_metadata)
|
||||
@@ -0,0 +1,67 @@
|
||||
"""
|
||||
文本清理工具
|
||||
|
||||
提供文本清理和标准化的工具方法。
|
||||
"""
|
||||
import re
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class TextCleaner:
|
||||
"""文本清理工具类"""
|
||||
|
||||
@staticmethod
|
||||
def clean(text: Optional[str]) -> str:
|
||||
"""清理文本
|
||||
|
||||
Args:
|
||||
text: 原始文本
|
||||
|
||||
Returns:
|
||||
清理后的文本
|
||||
"""
|
||||
if not text:
|
||||
return ""
|
||||
|
||||
text = TextCleaner._remove_control_characters(text)
|
||||
text = TextCleaner._normalize_whitespace(text)
|
||||
text = TextCleaner._remove_empty_lines(text)
|
||||
|
||||
return text.strip()
|
||||
|
||||
@staticmethod
|
||||
def _remove_control_characters(text: str) -> str:
|
||||
"""移除控制字符"""
|
||||
control_chars = re.compile(r'[\x00-\x08\x0b-\x0c\x0e-\x1f\x7f-\x9f]')
|
||||
return control_chars.sub('', text)
|
||||
|
||||
@staticmethod
|
||||
def _normalize_whitespace(text: str) -> str:
|
||||
"""标准化空白字符"""
|
||||
text = re.sub(r'[ \t]+', ' ', text)
|
||||
text = re.sub(r'\n\s*\n\s*\n+', '\n\n', text)
|
||||
return text
|
||||
|
||||
@staticmethod
|
||||
def _remove_empty_lines(text: str) -> str:
|
||||
"""移除空行"""
|
||||
lines = text.split('\n')
|
||||
return '\n'.join(line for line in lines if line.strip())
|
||||
|
||||
@staticmethod
|
||||
def has_printable_content(text: str) -> bool:
|
||||
"""检查是否包含可打印内容
|
||||
|
||||
Args:
|
||||
text: 文本
|
||||
|
||||
Returns:
|
||||
是否包含可打印内容
|
||||
"""
|
||||
if not text:
|
||||
return False
|
||||
return any(c.isprintable() and not c.isspace() for c in text)
|
||||
@@ -0,0 +1,325 @@
|
||||
"""
|
||||
文件处理器
|
||||
|
||||
负责文件的后台 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()
|
||||
@@ -0,0 +1,488 @@
|
||||
"""
|
||||
知识库业务服务
|
||||
|
||||
实现知识库的 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
|
||||
)
|
||||
@@ -0,0 +1,27 @@
|
||||
"""
|
||||
策略模式实现
|
||||
|
||||
提供知识库检索和处理的策略模式实现。
|
||||
"""
|
||||
from app.db.models.knowledge_gen import RagType
|
||||
from .base import KnowledgeBaseStrategy
|
||||
from .factory import KnowledgeBaseStrategyFactory
|
||||
from .vector_strategy import VectorKnowledgeBaseStrategy
|
||||
from .graph_strategy import GraphKnowledgeBaseStrategy
|
||||
|
||||
KnowledgeBaseStrategyFactory.register_strategy(
|
||||
RagType.DOCUMENT.value,
|
||||
VectorKnowledgeBaseStrategy
|
||||
)
|
||||
|
||||
KnowledgeBaseStrategyFactory.register_strategy(
|
||||
RagType.GRAPH.value,
|
||||
GraphKnowledgeBaseStrategy
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"KnowledgeBaseStrategy",
|
||||
"KnowledgeBaseStrategyFactory",
|
||||
"VectorKnowledgeBaseStrategy",
|
||||
"GraphKnowledgeBaseStrategy",
|
||||
]
|
||||
@@ -0,0 +1,105 @@
|
||||
"""
|
||||
策略模式 - 知识库策略基类
|
||||
|
||||
定义知识库检索和处理的抽象接口，支持不同类型的知识库（向量、知识图谱等）。
|
||||
"""
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
|
||||
class KnowledgeBaseStrategy(ABC):
|
||||
"""知识库策略抽象基类
|
||||
|
||||
所有知识库策略（向量、知识图谱等）必须实现此接口。
|
||||
|
||||
接口设计原则：
|
||||
- query(): 直接获取数据，不涉及相似度计算，返回值格式不统一
|
||||
- search(): 基于输入文本的相似度检索，返回值必须统一为 UnifiedSearchResult
|
||||
- process_file(): 处理单个文件，不同策略有不同的处理逻辑
|
||||
"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
"""初始化策略
|
||||
|
||||
Args:
|
||||
db: 数据库异步 session
|
||||
"""
|
||||
self.db = db
|
||||
|
||||
@abstractmethod
|
||||
async def query(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
**kwargs
|
||||
) -> Any:
|
||||
"""直接查询知识库数据（不涉及相似度计算）
|
||||
|
||||
不同知识库类型的返回格式可以不同：
|
||||
- 向量知识库：返回 PagedResponse[RagChunkResp]（分页分块列表）
|
||||
- 知识图谱：返回图谱数据（格式由 LightRAG 定义）
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
**kwargs: 额外参数（如 rag_file_id, paging_query 等）
|
||||
|
||||
Returns:
|
||||
查询结果（格式由具体策略决定）
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def search(
|
||||
self,
|
||||
query_text: str,
|
||||
knowledge_base_ids: List[str],
|
||||
top_k: int = 10,
|
||||
threshold: Optional[float] = None,
|
||||
**kwargs
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""基于输入文本的相似度检索
|
||||
|
||||
所有知识库类型必须返回统一格式：UnifiedSearchResult
|
||||
|
||||
Args:
|
||||
query_text: 查询文本
|
||||
knowledge_base_ids: 知识库 ID 列表
|
||||
top_k: 返回结果数量
|
||||
threshold: 相似度阈值（可选）
|
||||
**kwargs: 额外参数
|
||||
|
||||
Returns:
|
||||
统一格式的检索结果列表，每个元素包含：
|
||||
- id: 结果 ID
|
||||
- text: 文本内容
|
||||
- score: 相似度分数
|
||||
- metadata: 元数据字典
|
||||
- resultType: 结果类型（"vector" 或 "graph"）
|
||||
- knowledgeBaseId: 知识库 ID
|
||||
- knowledgeBaseName: 知识库名称
|
||||
"""
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def process_file(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
rag_file_id: str,
|
||||
**kwargs
|
||||
) -> None:
|
||||
"""处理单个文件
|
||||
|
||||
不同知识库类型有不同的处理逻辑：
|
||||
- 向量知识库：文档分块 -> 向量化 -> 存储到 Milvus
|
||||
- 知识图谱：文档分块 -> 构建知识图谱 -> 存储到 LightRAG
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
rag_file_id: RAG 文件 ID
|
||||
**kwargs: 额外参数（如 process_type, chunk_size 等）
|
||||
|
||||
Raises:
|
||||
BusinessError: 处理失败时抛出
|
||||
"""
|
||||
pass
|
||||
@@ -0,0 +1,69 @@
|
||||
"""
|
||||
策略工厂
|
||||
|
||||
根据知识库类型（DOCUMENT/GRAPH）创建对应的知识库策略实例。
|
||||
"""
|
||||
from typing import Dict
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.db.models.knowledge_gen import RagType
|
||||
from app.core.exception import BusinessError, ErrorCodes
|
||||
from .base import KnowledgeBaseStrategy
|
||||
|
||||
|
||||
class KnowledgeBaseStrategyFactory:
|
||||
"""知识库策略工厂
|
||||
|
||||
使用工厂模式创建知识库策略实例，支持策略缓存以提高性能。
|
||||
"""
|
||||
|
||||
_strategy_cache: Dict[str, type] = {}
|
||||
|
||||
@classmethod
|
||||
def register_strategy(cls, rag_type: str, strategy_class: type):
|
||||
"""注册策略类
|
||||
|
||||
Args:
|
||||
rag_type: 知识库类型（DOCUMENT/GRAPH）
|
||||
strategy_class: 策略类
|
||||
"""
|
||||
cls._strategy_cache[rag_type.upper()] = strategy_class
|
||||
|
||||
@classmethod
|
||||
def create_strategy(
|
||||
cls,
|
||||
rag_type: str,
|
||||
db: AsyncSession
|
||||
) -> KnowledgeBaseStrategy:
|
||||
"""创建知识库策略实例
|
||||
|
||||
Args:
|
||||
rag_type: 知识库类型（DOCUMENT/GRAPH）
|
||||
db: 数据库异步 session
|
||||
|
||||
Returns:
|
||||
知识库策略实例
|
||||
|
||||
Raises:
|
||||
BusinessError: 不支持的知识库类型
|
||||
"""
|
||||
rag_type_upper = rag_type.upper() if isinstance(rag_type, str) else str(rag_type).upper()
|
||||
|
||||
if rag_type_upper not in cls._strategy_cache:
|
||||
raise BusinessError(
|
||||
ErrorCodes.RAG_UNSUPPORTED_TYPE,
|
||||
f"不支持的 RAG 类型: {rag_type}"
|
||||
)
|
||||
|
||||
strategy_class = cls._strategy_cache[rag_type_upper]
|
||||
return strategy_class(db)
|
||||
|
||||
@classmethod
|
||||
def get_supported_types(cls) -> list:
|
||||
"""获取支持的知识库类型列表
|
||||
|
||||
Returns:
|
||||
支持的类型列表
|
||||
"""
|
||||
return list(cls._strategy_cache.keys())
|
||||
@@ -0,0 +1,366 @@
|
||||
"""
|
||||
知识图谱策略
|
||||
|
||||
实现 GRAPH 类型知识库的检索和处理逻辑。
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any, Awaitable, Callable, Dict, List, Optional
|
||||
|
||||
import numpy as np
|
||||
from lightrag import LightRAG
|
||||
from lightrag.base import QueryParam
|
||||
from lightrag.constants import DEFAULT_ENTITY_TYPES
|
||||
from lightrag.llm.openai import openai_complete_if_cache, openai_embed
|
||||
from lightrag.utils import EmbeddingFunc, get_env_value, setup_logger
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.config import settings
|
||||
from app.core.exception import BusinessError, ErrorCodes
|
||||
from app.db.models.knowledge_gen import KnowledgeBase
|
||||
from app.module.rag.repository import KnowledgeBaseRepository
|
||||
from app.module.rag.service.common import get_file_path
|
||||
from app.module.system.service.common_service import get_model_by_id
|
||||
from .base import KnowledgeBaseStrategy
|
||||
|
||||
|
||||
setup_logger("lightrag", level="INFO")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
DEFAULT_WORKING_DIR = str(settings.rag_storage_dir)
|
||||
|
||||
|
||||
def _create_llm_func(model_name: str, base_url: str, api_key: str) -> Callable[..., Awaitable[str]]:
|
||||
async def _llm(prompt: str, system_prompt: str = None, history_messages: list = None, **kwargs) -> str:
|
||||
return await openai_complete_if_cache(
|
||||
model_name, prompt,
|
||||
system_prompt=system_prompt,
|
||||
history_messages=history_messages or [],
|
||||
api_key=api_key, base_url=base_url, **kwargs,
|
||||
)
|
||||
return _llm
|
||||
|
||||
|
||||
def _create_embedding_func(model_name: str, base_url: str, api_key: str, embedding_dim: int) -> EmbeddingFunc:
|
||||
async def _embed(texts: list[str]) -> np.ndarray:
|
||||
return await openai_embed.func(texts, model=model_name, api_key=api_key, base_url=base_url)
|
||||
return EmbeddingFunc(embedding_dim=embedding_dim, func=_embed, max_token_size=8192)
|
||||
|
||||
|
||||
async def _create_rag(
|
||||
llm_func: Callable[..., Awaitable[str]],
|
||||
embedding_func: EmbeddingFunc,
|
||||
working_dir: str,
|
||||
workspace: str = "",
|
||||
) -> LightRAG:
|
||||
os.makedirs(working_dir, exist_ok=True)
|
||||
rag = LightRAG(
|
||||
working_dir=working_dir,
|
||||
workspace=workspace,
|
||||
llm_model_func=llm_func,
|
||||
embedding_func=embedding_func,
|
||||
addon_params={
|
||||
"language": "Chinese",
|
||||
"entity_types": get_env_value("ENTITY_TYPES", DEFAULT_ENTITY_TYPES, list),
|
||||
},
|
||||
)
|
||||
await rag.initialize_storages()
|
||||
return rag
|
||||
|
||||
|
||||
class GraphKnowledgeBaseStrategy(KnowledgeBaseStrategy):
|
||||
# 类级别的缓存，允许跨实例共享
|
||||
_rag_cache: Dict[str, Any] = {}
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
super().__init__(db)
|
||||
self.kb_repo = KnowledgeBaseRepository(db)
|
||||
|
||||
async def query(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
**kwargs
|
||||
) -> Any:
|
||||
node_label = kwargs.get("node_label")
|
||||
if not node_label:
|
||||
raise BusinessError(
|
||||
ErrorCodes.RAG_INVALID_REQUEST,
|
||||
"Missing 'node_label' parameter for graph query"
|
||||
)
|
||||
|
||||
kb = await self._get_knowledge_base(knowledge_base_id)
|
||||
if not kb:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
|
||||
rag_instance = await self._get_or_create_graph_rag(kb)
|
||||
return await rag_instance.get_knowledge_graph(node_label=node_label)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query_text: str,
|
||||
knowledge_base_ids: List[str],
|
||||
top_k: int = 10,
|
||||
threshold: Optional[float] = None,
|
||||
**kwargs
|
||||
) -> List[Dict[str, Any]]:
|
||||
if len(knowledge_base_ids) != 1:
|
||||
raise BusinessError(
|
||||
ErrorCodes.RAG_INVALID_REQUEST,
|
||||
"At least one knowledge base required for graph search"
|
||||
)
|
||||
|
||||
all_results = []
|
||||
for kb_id in knowledge_base_ids:
|
||||
kb = await self._get_knowledge_base(kb_id)
|
||||
if not kb:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
|
||||
rag_instance = await self._get_or_create_graph_rag(kb)
|
||||
# Use aquery_data for content retrieval (not get_knowledge_graph)
|
||||
query_param = QueryParam(mode="mix", top_k=top_k, only_need_context=True)
|
||||
retrieval_results = await rag_instance.aquery_data(query_text, query_param)
|
||||
|
||||
unified_results = self._convert_retrieval_results_into_unified(
|
||||
retrieval_results, str(kb.id), str(kb.name)
|
||||
)
|
||||
all_results.extend(unified_results)
|
||||
|
||||
all_results.sort(key=lambda x: x.get("score", 1.0), reverse=True)
|
||||
logger.info("知识图谱检索完成: 结果数=%d", len(all_results))
|
||||
return all_results
|
||||
|
||||
@staticmethod
|
||||
def _convert_retrieval_results_into_unified(
|
||||
retrieval_results: Dict[str, Any],
|
||||
knowledge_base_id: str,
|
||||
knowledge_base_name: str
|
||||
) -> List[Dict[str, Any]]:
|
||||
if not isinstance(retrieval_results, dict):
|
||||
return []
|
||||
|
||||
data = retrieval_results.get("data", retrieval_results)
|
||||
entities = data.get("entities", [])
|
||||
relationships = data.get("relationships", [])
|
||||
chunks = data.get("chunks", [])
|
||||
|
||||
if not entities and not relationships and not chunks:
|
||||
return []
|
||||
|
||||
reference_files: Dict[str, str] = {}
|
||||
ref_counter = 1
|
||||
|
||||
for chunk in chunks:
|
||||
file_path = chunk.get("file_path", "")
|
||||
if file_path and file_path not in reference_files:
|
||||
reference_files[file_path] = str(ref_counter)
|
||||
ref_counter += 1
|
||||
|
||||
for entity in entities:
|
||||
file_path = entity.get("file_path", "")
|
||||
if file_path and file_path not in reference_files:
|
||||
reference_files[file_path] = str(ref_counter)
|
||||
ref_counter += 1
|
||||
|
||||
for rel in relationships:
|
||||
file_path = rel.get("file_path", "")
|
||||
if file_path and file_path not in reference_files:
|
||||
reference_files[file_path] = str(ref_counter)
|
||||
ref_counter += 1
|
||||
|
||||
text_parts = []
|
||||
|
||||
if entities:
|
||||
text_parts.append("\nKnowledge Graph Data (Entity):\n")
|
||||
text_parts.append("\n```json")
|
||||
for entity in entities:
|
||||
entity_data = {
|
||||
"entity": entity.get("entity_name", ""),
|
||||
"type": entity.get("entity_type", ""),
|
||||
"description": entity.get("description", ""),
|
||||
}
|
||||
text_parts.append(f"\n{json.dumps(entity_data, ensure_ascii=False)}")
|
||||
text_parts.append("\n```")
|
||||
|
||||
if relationships:
|
||||
text_parts.append("\n\nKnowledge Graph Data (Relationship):\n")
|
||||
text_parts.append("\n```json")
|
||||
for rel in relationships:
|
||||
rel_data = {
|
||||
"entity1": rel.get("src_id", ""),
|
||||
"entity2": rel.get("tgt_id", ""),
|
||||
"description": rel.get("description", ""),
|
||||
}
|
||||
text_parts.append(f"\n{json.dumps(rel_data, ensure_ascii=False)}")
|
||||
text_parts.append("\n```")
|
||||
|
||||
if chunks:
|
||||
text_parts.append(
|
||||
"\n\nDocument Chunks (Each entry has a reference_id refer to the `Reference Document List`):\n"
|
||||
)
|
||||
text_parts.append("\n```json")
|
||||
for chunk in chunks:
|
||||
file_path = chunk.get("file_path", "")
|
||||
ref_id = reference_files.get(file_path, "")
|
||||
chunk_data = {
|
||||
"reference_id": ref_id,
|
||||
"content": chunk.get("content", ""),
|
||||
}
|
||||
text_parts.append(f"\n{json.dumps(chunk_data, ensure_ascii=False)}")
|
||||
text_parts.append("\n```")
|
||||
|
||||
if reference_files:
|
||||
text_parts.append(
|
||||
"\n\nReference Document List (Each entry starts with a [reference_id] "
|
||||
"that corresponds to entries in the Document Chunks):\n"
|
||||
)
|
||||
text_parts.append("\n```")
|
||||
for file_path, ref_id in reference_files.items():
|
||||
file_name = Path(file_path).name if file_path else ""
|
||||
text_parts.append(f"\n[{ref_id}] {file_name}")
|
||||
text_parts.append("\n```")
|
||||
|
||||
formatted_text = "".join(text_parts)
|
||||
|
||||
return [{
|
||||
"id": knowledge_base_id,
|
||||
"text": formatted_text,
|
||||
"score": 1.0,
|
||||
"metadata": {
|
||||
"entity_count": len(entities),
|
||||
"relationship_count": len(relationships),
|
||||
"chunk_count": len(chunks),
|
||||
"reference_files": {v: k for k, v in reference_files.items()},
|
||||
},
|
||||
"resultType": "graph_retrieval",
|
||||
"knowledgeBaseId": knowledge_base_id,
|
||||
"knowledgeBaseName": knowledge_base_name,
|
||||
}]
|
||||
|
||||
async def process_file(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
rag_file_id: str,
|
||||
**kwargs
|
||||
) -> None:
|
||||
from app.module.rag.infra.document.processor import ingest_file_to_chunks
|
||||
from app.db.models.knowledge_gen import FileStatus
|
||||
from app.module.rag.repository import RagFileRepository
|
||||
|
||||
kb_repo = KnowledgeBaseRepository(self.db)
|
||||
file_repo = RagFileRepository(self.db)
|
||||
|
||||
kb = await kb_repo.get_by_id(knowledge_base_id)
|
||||
if not kb:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
|
||||
rag_file = await file_repo.get_by_id(rag_file_id)
|
||||
if not rag_file:
|
||||
raise BusinessError(ErrorCodes.RAG_FILE_NOT_FOUND)
|
||||
|
||||
try:
|
||||
await file_repo.update_status(rag_file_id, FileStatus.PROCESSING)
|
||||
await self.db.commit()
|
||||
|
||||
file_path = get_file_path(rag_file)
|
||||
if not file_path or not Path(file_path).exists():
|
||||
await file_repo.update_status(rag_file_id, FileStatus.PROCESS_FAILED)
|
||||
await self.db.commit()
|
||||
logger.error("文件不存在: %s", file_path)
|
||||
return
|
||||
|
||||
chunks = await ingest_file_to_chunks(
|
||||
file_path,
|
||||
process_type=kwargs.get("process_type", "DEFAULT_CHUNK"),
|
||||
chunk_size=kwargs.get("chunk_size", 500),
|
||||
overlap_size=kwargs.get("overlap_size", 50),
|
||||
delimiter=kwargs.get("delimiter"),
|
||||
)
|
||||
|
||||
if not chunks:
|
||||
await file_repo.update_status(rag_file_id, FileStatus.PROCESS_FAILED)
|
||||
await self.db.commit()
|
||||
logger.error("文件解析失败，未生成任何文档")
|
||||
return
|
||||
|
||||
rag_instance = await self._get_or_create_graph_rag(kb)
|
||||
|
||||
for idx, chunk in enumerate(chunks):
|
||||
logger.info(
|
||||
"插入文档到知识图谱: %s, 进度: %d/%d",
|
||||
rag_file.file_name, idx + 1, len(chunks)
|
||||
)
|
||||
await rag_instance.ainsert(
|
||||
input=chunk.text,
|
||||
file_paths=[file_path]
|
||||
)
|
||||
|
||||
await file_repo.update_status(rag_file_id, FileStatus.PROCESSED)
|
||||
await file_repo.update_chunk_count(rag_file_id, len(chunks))
|
||||
await self.db.commit()
|
||||
|
||||
logger.info("文件 %s 知识图谱构建完成，文档数=%d", rag_file.file_name, len(chunks))
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("文件 %s 知识图谱构建失败: %s", rag_file.file_name, e)
|
||||
await file_repo.update_status(rag_file_id, FileStatus.PROCESS_FAILED)
|
||||
await self.db.commit()
|
||||
raise BusinessError(
|
||||
ErrorCodes.RAG_FILE_PROCESS_FAILED,
|
||||
f"知识图谱构建失败: {str(e)}"
|
||||
) from e
|
||||
|
||||
async def _get_knowledge_base(self, knowledge_base_id: str) -> KnowledgeBase:
|
||||
kb = await self.kb_repo.get_by_id(knowledge_base_id)
|
||||
if not kb:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
return kb
|
||||
|
||||
async def _get_or_create_graph_rag(self, kb: KnowledgeBase) -> Any:
|
||||
kb_name = str(kb.name)
|
||||
if kb_name in self._rag_cache:
|
||||
return self._rag_cache[kb_name]
|
||||
|
||||
chat_model = await get_model_by_id(self.db, str(kb.chat_model))
|
||||
embedding_model = await get_model_by_id(self.db, str(kb.embedding_model))
|
||||
|
||||
if not chat_model or not embedding_model:
|
||||
raise BusinessError(ErrorCodes.RAG_MODEL_NOT_FOUND)
|
||||
|
||||
llm_func = _create_llm_func(
|
||||
str(chat_model.model_name),
|
||||
str(chat_model.base_url),
|
||||
str(chat_model.api_key),
|
||||
)
|
||||
|
||||
from app.module.shared.llm import LLMFactory
|
||||
embedding_func = _create_embedding_func(
|
||||
str(embedding_model.model_name),
|
||||
str(embedding_model.base_url),
|
||||
str(embedding_model.api_key),
|
||||
LLMFactory.get_embedding_dimension(
|
||||
str(embedding_model.model_name),
|
||||
str(embedding_model.base_url),
|
||||
str(embedding_model.api_key),
|
||||
),
|
||||
)
|
||||
|
||||
rag = await _create_rag(llm_func, embedding_func, DEFAULT_WORKING_DIR, workspace=kb_name)
|
||||
self._rag_cache[kb_name] = rag
|
||||
return rag
|
||||
|
||||
@classmethod
|
||||
def rename_workspace(cls, old_name: str, new_name: str) -> None:
|
||||
old_path = Path(DEFAULT_WORKING_DIR) / old_name
|
||||
new_path = Path(DEFAULT_WORKING_DIR) / new_name
|
||||
if old_path.exists() and old_path.is_dir():
|
||||
old_path.rename(new_path)
|
||||
logger.info("知识图谱 workspace 重命名: %s -> %s", old_name, new_name)
|
||||
cls.clear_cache(old_name)
|
||||
|
||||
@classmethod
|
||||
def clear_cache(cls, name: str) -> None:
|
||||
if name in cls._rag_cache:
|
||||
del cls._rag_cache[name]
|
||||
logger.info("已清除知识图谱缓存: %s", name)
|
||||
@@ -0,0 +1,481 @@
|
||||
"""
|
||||
向量知识库策略
|
||||
|
||||
实现 DOCUMENT 类型知识库的检索和处理逻辑。
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from pymilvus import AnnSearchRequest, Function, FunctionType
|
||||
|
||||
from app.core.exception import BusinessError, ErrorCodes
|
||||
from app.module.rag.infra.embeddings import EmbeddingFactory
|
||||
from app.module.rag.infra.vectorstore.milvus_client import get_milvus_client
|
||||
from app.module.rag.repository import KnowledgeBaseRepository, RagFileRepository
|
||||
from app.module.rag.schema.response import PagedResponse, RagChunkResp
|
||||
from app.module.rag.schema.request import ChunkFilterQuery
|
||||
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
|
||||
from .base import KnowledgeBaseStrategy
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class VectorKnowledgeBaseStrategy(KnowledgeBaseStrategy):
|
||||
"""向量知识库策略实现
|
||||
|
||||
提供 DOCUMENT 类型知识库的 query、 search 和 process_file 功能。
|
||||
"""
|
||||
|
||||
async def query(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
**kwargs
|
||||
) -> PagedResponse:
|
||||
"""查询知识库分块数据（不涉及相似度计算）
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
**kwargs: 额外参数，必须包含:
|
||||
- rag_file_id: RAG 文件 ID
|
||||
- chunk_filter_query: 分页和过滤参数
|
||||
|
||||
Returns:
|
||||
分页分块列表
|
||||
"""
|
||||
rag_file_id = kwargs.get("rag_file_id")
|
||||
chunk_filter_query: ChunkFilterQuery = kwargs.get("chunk_filter_query")
|
||||
|
||||
if not rag_file_id or not chunk_filter_query:
|
||||
raise BusinessError(
|
||||
ErrorCodes.RAG_INVALID_REQUEST,
|
||||
"Missing rag_file_id or chunk_filter_query parameters"
|
||||
)
|
||||
|
||||
kb_repo = KnowledgeBaseRepository(self.db)
|
||||
file_repo = RagFileRepository(self.db)
|
||||
|
||||
knowledge_base = await kb_repo.get_by_id(knowledge_base_id)
|
||||
if not knowledge_base:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
|
||||
rag_file = await file_repo.get_by_id(rag_file_id)
|
||||
if not rag_file:
|
||||
raise BusinessError(ErrorCodes.RAG_FILE_NOT_FOUND)
|
||||
|
||||
client = get_milvus_client()
|
||||
|
||||
try:
|
||||
base_filter = f'metadata["rag_file_id"] == "{rag_file_id}"'
|
||||
combined_filter = self._build_combined_filter(base_filter, chunk_filter_query.expr)
|
||||
|
||||
count_res = client.query(
|
||||
collection_name=knowledge_base.name,
|
||||
filter=combined_filter,
|
||||
output_fields=["id"],
|
||||
)
|
||||
total = len(count_res)
|
||||
|
||||
offset = (chunk_filter_query.page - 1) * chunk_filter_query.size
|
||||
results = client.query(
|
||||
collection_name=knowledge_base.name,
|
||||
filter=combined_filter,
|
||||
output_fields=["id", "text", "metadata"],
|
||||
limit=chunk_filter_query.size,
|
||||
offset=offset,
|
||||
)
|
||||
|
||||
chunks = [
|
||||
RagChunkResp(
|
||||
id=item.get("id", ""),
|
||||
text=item.get("text", ""),
|
||||
metadata=item.get("metadata", {}),
|
||||
score=0.0,
|
||||
distance=None,
|
||||
)
|
||||
for item in results
|
||||
]
|
||||
|
||||
logger.info(
|
||||
"查询文件分块成功: kb=%s file=%s total=%d",
|
||||
knowledge_base_id, rag_file_id, total
|
||||
)
|
||||
|
||||
return PagedResponse.create(
|
||||
content=chunks,
|
||||
total_elements=total,
|
||||
page=chunk_filter_query.page,
|
||||
size=chunk_filter_query.size,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"查询文件分块失败: kb=%s file=%s error=%s",
|
||||
knowledge_base_id, rag_file_id, e
|
||||
)
|
||||
raise BusinessError(
|
||||
ErrorCodes.RAG_MILVUS_ERROR,
|
||||
f"查询文件分块失败: {str(e)}"
|
||||
) from e
|
||||
|
||||
@staticmethod
|
||||
def _build_combined_filter(base_filter: str, user_expr: Optional[str]) -> str:
|
||||
if not user_expr:
|
||||
return base_filter
|
||||
return f"{base_filter} && {user_expr}"
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query_text: str,
|
||||
knowledge_base_ids: List[str],
|
||||
top_k: int = 10,
|
||||
threshold: Optional[float] = None,
|
||||
**kwargs
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""基于输入文本的相似度检索（混合检索： 向量 + BM25)
|
||||
|
||||
Args:
|
||||
query_text: 查询文本
|
||||
knowledge_base_ids: 知识库 ID 列表
|
||||
top_k: 返回结果数量
|
||||
threshold: 相似度阈值（可选)
|
||||
**kwargs: 额外参数
|
||||
|
||||
Returns:
|
||||
统一格式的检索结果列表
|
||||
"""
|
||||
kb_repo = KnowledgeBaseRepository(self.db)
|
||||
|
||||
knowledge_bases = []
|
||||
for kb_id in knowledge_base_ids:
|
||||
kb = await kb_repo.get_by_id(kb_id)
|
||||
if not kb:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
knowledge_bases.append(kb)
|
||||
|
||||
embedding_entity = await get_model_by_id(self.db, knowledge_bases[0].embedding_model)
|
||||
if not embedding_entity:
|
||||
raise BusinessError(ErrorCodes.RAG_MODEL_NOT_FOUND)
|
||||
|
||||
embedding = EmbeddingFactory.create_embeddings(
|
||||
model_name=embedding_entity.model_name,
|
||||
base_url=getattr(embedding_entity, "base_url", None),
|
||||
api_key=getattr(embedding_entity, "api_key", None),
|
||||
)
|
||||
|
||||
try:
|
||||
query_vector = await asyncio.to_thread(embedding.embed_query, query_text)
|
||||
except Exception as e:
|
||||
logger.error("查询向量化失败: %s", e)
|
||||
raise BusinessError(
|
||||
ErrorCodes.RAG_EMBEDDING_FAILED,
|
||||
f"查询向量化失败: {str(e)}"
|
||||
) from e
|
||||
|
||||
all_results = await self._execute_hybrid_search(
|
||||
knowledge_bases, query_vector, query_text, top_k
|
||||
)
|
||||
|
||||
all_results.sort(
|
||||
key=lambda x: x.get("score") or x.get("distance", 0),
|
||||
reverse=True
|
||||
)
|
||||
|
||||
if threshold is not None:
|
||||
all_results = [
|
||||
r for r in all_results
|
||||
if (r.get("score") or r.get("distance", 0)) >= 0
|
||||
]
|
||||
|
||||
formatted = self._format_unified_results(all_results)
|
||||
logger.info("向量检索完成: 结果数=%d", len(formatted))
|
||||
return formatted
|
||||
|
||||
@staticmethod
|
||||
async def _execute_hybrid_search(
|
||||
knowledge_bases: list,
|
||||
query_vector: list,
|
||||
query_text: str,
|
||||
top_k: int,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""执行混合检索"""
|
||||
all_results = []
|
||||
client = get_milvus_client()
|
||||
|
||||
for kb in knowledge_bases:
|
||||
try:
|
||||
if not client.has_collection(kb.name):
|
||||
logger.warning("集合 %s 不存在，跳过", kb.name)
|
||||
continue
|
||||
|
||||
dense_search = AnnSearchRequest(
|
||||
data=[query_vector],
|
||||
anns_field="vector",
|
||||
param={"nprobe": 10},
|
||||
limit=top_k,
|
||||
)
|
||||
|
||||
sparse_search = AnnSearchRequest(
|
||||
data=[query_text],
|
||||
anns_field="sparse",
|
||||
param={"drop_ratio_search": 0.2},
|
||||
limit=top_k,
|
||||
)
|
||||
|
||||
ranker = Function(
|
||||
name="weight",
|
||||
input_field_names=[],
|
||||
function_type=FunctionType.RERANK,
|
||||
params={
|
||||
"reranker": "weighted",
|
||||
"weights": [0.1, 0.9],
|
||||
"norm_score": True,
|
||||
}
|
||||
)
|
||||
search_results = client.hybrid_search(
|
||||
collection_name=kb.name,
|
||||
reqs=[dense_search, sparse_search],
|
||||
ranker=ranker,
|
||||
output_fields=["id", "text", "metadata"],
|
||||
limit=top_k,
|
||||
)
|
||||
|
||||
if search_results and len(search_results) > 0:
|
||||
for result in search_results[0]:
|
||||
result["knowledge_base_id"] = kb.id
|
||||
result["knowledge_base_name"] = kb.name
|
||||
all_results.append(result)
|
||||
|
||||
except Exception as e:
|
||||
logger.error("知识库 %s 混合检索失败: %s", kb.name, e)
|
||||
continue
|
||||
|
||||
return all_results
|
||||
|
||||
@staticmethod
|
||||
def _format_unified_results(
|
||||
all_results: List[Dict[str, Any]]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""格式化为统一检索结果"""
|
||||
formatted = []
|
||||
for r in all_results:
|
||||
entity = r.get("entity", {})
|
||||
metadata = entity.get("metadata", {})
|
||||
if isinstance(metadata, dict):
|
||||
metadata_str = json.dumps(metadata, ensure_ascii=False)
|
||||
else:
|
||||
metadata_str = metadata if metadata else "{}"
|
||||
|
||||
formatted.append({
|
||||
"id": entity.get("id", ""),
|
||||
"text": entity.get("text", ""),
|
||||
"score": r.get("score") or r.get("distance", 0),
|
||||
"metadata": json.loads(metadata_str) if isinstance(metadata_str, str) else metadata,
|
||||
"resultType": "vector",
|
||||
"knowledgeBaseId": r.get("knowledge_base_id", ""),
|
||||
"knowledgeBaseName": r.get("knowledge_base_name", ""),
|
||||
})
|
||||
|
||||
return formatted
|
||||
|
||||
async def process_file(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
rag_file_id: str,
|
||||
**kwargs
|
||||
) -> None:
|
||||
"""处理单个文件（文档向量化)
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
rag_file_id: RAG 文件 ID
|
||||
**kwargs: 额外参数
|
||||
- process_type: 分块类型
|
||||
- chunk_size: 分块大小
|
||||
- overlap_size: 重叠大小
|
||||
- delimiter: 分隔符
|
||||
"""
|
||||
from pathlib import Path
|
||||
from app.module.rag.infra.vectorstore import VectorStoreFactory
|
||||
from app.module.rag.infra.document.processor import ingest_file_to_chunks
|
||||
from app.db.models.knowledge_gen import FileStatus
|
||||
from app.module.rag.service.common import TextCleaner, MetadataBuilder, BatchProcessor
|
||||
|
||||
kb_repo = KnowledgeBaseRepository(self.db)
|
||||
file_repo = RagFileRepository(self.db)
|
||||
|
||||
kb = await kb_repo.get_by_id(knowledge_base_id)
|
||||
if not kb:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
|
||||
rag_file = await file_repo.get_by_id(rag_file_id)
|
||||
if not rag_file:
|
||||
raise BusinessError(ErrorCodes.RAG_FILE_NOT_FOUND)
|
||||
|
||||
try:
|
||||
await file_repo.update_status(rag_file_id, FileStatus.PROCESSING)
|
||||
await self.db.commit()
|
||||
|
||||
file_path = self._get_file_path(rag_file)
|
||||
if not file_path or not Path(file_path).exists():
|
||||
await file_repo.update_status(rag_file_id, FileStatus.PROCESS_FAILED)
|
||||
await self.db.commit()
|
||||
logger.error("文件不存在: %s", file_path)
|
||||
return
|
||||
|
||||
chunks = await ingest_file_to_chunks(
|
||||
file_path,
|
||||
process_type=kwargs.get("process_type", "DEFAULT_CHUNK"),
|
||||
chunk_size=kwargs.get("chunk_size", 500),
|
||||
overlap_size=kwargs.get("overlap_size", 50),
|
||||
delimiter=kwargs.get("delimiter"),
|
||||
)
|
||||
|
||||
if not chunks:
|
||||
await file_repo.update_status(rag_file_id, FileStatus.PROCESS_FAILED)
|
||||
await self.db.commit()
|
||||
logger.error("文件解析失败，未生成任何分块")
|
||||
return
|
||||
|
||||
valid_chunks = self._filter_and_clean_chunks(chunks)
|
||||
if not valid_chunks:
|
||||
await file_repo.update_status(rag_file_id, FileStatus.PROCESS_FAILED)
|
||||
await self.db.commit()
|
||||
logger.error("文件没有有效的分块内容")
|
||||
return
|
||||
|
||||
embedding_entity = await get_model_by_id(self.db, kb.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),
|
||||
)
|
||||
|
||||
vectorstore = VectorStoreFactory.create(
|
||||
collection_name=str(kb.name),
|
||||
embedding=embedding,
|
||||
)
|
||||
|
||||
base_metadata = MetadataBuilder.build_chunk_metadata(rag_file, kb)
|
||||
MetadataBuilder.add_to_chunks(valid_chunks, base_metadata)
|
||||
|
||||
await BatchProcessor.store_in_batches(vectorstore, valid_chunks)
|
||||
|
||||
await file_repo.update_status(rag_file_id, FileStatus.PROCESSED)
|
||||
await file_repo.update_chunk_count(rag_file_id, len(valid_chunks))
|
||||
await self.db.commit()
|
||||
|
||||
logger.info("文件 %s 向量化处理完成，分块数=%d", rag_file.file_name, len(valid_chunks))
|
||||
|
||||
except Exception as e:
|
||||
logger.exception("文件 %s 处理失败: %s", rag_file.file_name, e)
|
||||
await file_repo.update_status(rag_file_id, FileStatus.PROCESS_FAILED)
|
||||
await self.db.commit()
|
||||
raise BusinessError(
|
||||
ErrorCodes.RAG_FILE_PROCESS_FAILED,
|
||||
f"文件处理失败: {str(e)}"
|
||||
) from e
|
||||
|
||||
def _get_file_path(self, rag_file) -> Optional[str]:
|
||||
from app.module.rag.service.common import get_file_path
|
||||
return get_file_path(rag_file)
|
||||
|
||||
def _filter_and_clean_chunks(self, chunks: list) -> list:
|
||||
"""过滤和清理分块"""
|
||||
from app.module.rag.service.common import TextCleaner
|
||||
|
||||
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(
|
||||
"跳过无效分块: chunk_index=%s",
|
||||
chunk.metadata.get("chunk_index")
|
||||
)
|
||||
|
||||
logger.info("有效分块数量: %d / %d", len(valid_chunks), len(chunks))
|
||||
return valid_chunks
|
||||
|
||||
async def update_chunk(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
chunk_id: str,
|
||||
text: str,
|
||||
metadata: Optional[dict] = None,
|
||||
) -> None:
|
||||
"""更新指定分块的文本和元数据
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
chunk_id: 分块 ID
|
||||
text: 新的文本内容
|
||||
metadata: 新的元数据（可选）
|
||||
"""
|
||||
from app.module.rag.infra.vectorstore import update_chunk_by_id
|
||||
|
||||
kb_repo = KnowledgeBaseRepository(self.db)
|
||||
kb = await kb_repo.get_by_id(knowledge_base_id)
|
||||
if not kb:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
|
||||
embedding_entity = await get_model_by_id(self.db, kb.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(kb.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
|
||||
"""
|
||||
from app.module.rag.infra.vectorstore import delete_chunk_by_id
|
||||
|
||||
kb_repo = KnowledgeBaseRepository(self.db)
|
||||
kb = await kb_repo.get_by_id(knowledge_base_id)
|
||||
if not kb:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
|
||||
await asyncio.to_thread(
|
||||
delete_chunk_by_id,
|
||||
collection_name=str(kb.name),
|
||||
chunk_id=chunk_id,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"成功删除分块: kb=%s chunk_id=%s",
|
||||
knowledge_base_id, chunk_id
|
||||
)
|
||||
@@ -0,0 +1,122 @@
|
||||
"""
|
||||
统一检索服务
|
||||
|
||||
提供知识库内容的统一检索接口，支持多种知识库类型(向量、知识图谱等)。
|
||||
使用策略模式实现不同知识库类型的检索逻辑。
|
||||
"""
|
||||
import logging
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.exception import BusinessError, ErrorCodes
|
||||
from app.db.models.knowledge_gen import KnowledgeBase
|
||||
from app.module.rag.repository import KnowledgeBaseRepository
|
||||
from app.module.rag.schema.request import PagingQuery, ChunkFilterQuery, RetrieveReq
|
||||
from app.module.rag.schema.response import PagedResponse
|
||||
from .strategy import KnowledgeBaseStrategyFactory
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class UnifiedRetrievalService:
|
||||
"""统一检索服务
|
||||
|
||||
揯供统一的检索接口,根据知识库类型自动选择策略。
|
||||
"""
|
||||
|
||||
def __init__(self, db: AsyncSession):
|
||||
"""初始化服务
|
||||
|
||||
Args:
|
||||
db: 数据库异步 session
|
||||
"""
|
||||
self.db = db
|
||||
self.kb_repo = KnowledgeBaseRepository(db)
|
||||
|
||||
async def query(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
**kwargs
|
||||
) -> Any:
|
||||
"""直接查询知识库数据(不涉及相似度计算)
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
**kwargs: 额外参数
|
||||
- 向量知识库: rag_file_id, paging_query
|
||||
- 知识图谱: node_label
|
||||
|
||||
Returns:
|
||||
查询结果(格式由具体知识库类型决定)
|
||||
"""
|
||||
kb = await self._get_knowledge_base(knowledge_base_id)
|
||||
strategy = KnowledgeBaseStrategyFactory.create_strategy(kb.type, self.db)
|
||||
return await strategy.query(knowledge_base_id, **kwargs)
|
||||
|
||||
async def search(
|
||||
self,
|
||||
request: RetrieveReq
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""基于输入文本的相似度检索
|
||||
|
||||
Args:
|
||||
request: 检索请求
|
||||
|
||||
Returns:
|
||||
统一格式的检索结果列表(UnifiedSearchResult)
|
||||
"""
|
||||
if not request.knowledge_base_ids:
|
||||
raise BusinessError(
|
||||
ErrorCodes.RAG_INVALID_REQUEST,
|
||||
"至少需要一个知识库 ID"
|
||||
)
|
||||
|
||||
kb = await self._get_knowledge_base(request.knowledge_base_ids[0])
|
||||
strategy = KnowledgeBaseStrategyFactory.create_strategy(kb.type, self.db)
|
||||
|
||||
return await strategy.search(
|
||||
query_text=request.query,
|
||||
knowledge_base_ids=request.knowledge_base_ids,
|
||||
top_k=request.top_k,
|
||||
threshold=request.threshold,
|
||||
)
|
||||
|
||||
async def get_chunks(
|
||||
self,
|
||||
knowledge_base_id: str,
|
||||
rag_file_id: str,
|
||||
query: ChunkFilterQuery,
|
||||
) -> PagedResponse:
|
||||
"""获取指定 RAG 文件的分块列表(仅向量知识库)
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
rag_file_id: RAG 文件 ID
|
||||
query: 分页和过滤参数
|
||||
|
||||
Returns:
|
||||
分块列表(分页)
|
||||
"""
|
||||
return await self.query(
|
||||
knowledge_base_id,
|
||||
rag_file_id=rag_file_id,
|
||||
chunk_filter_query=query,
|
||||
)
|
||||
|
||||
async def _get_knowledge_base(self, knowledge_base_id: str) -> KnowledgeBase:
|
||||
"""获取知识库实体
|
||||
|
||||
Args:
|
||||
knowledge_base_id: 知识库 ID
|
||||
|
||||
Returns:
|
||||
知识库实体
|
||||
|
||||
Raises:
|
||||
BusinessError: 知识库不存在
|
||||
"""
|
||||
kb = await self.kb_repo.get_by_id(knowledge_base_id)
|
||||
if not kb:
|
||||
raise BusinessError(ErrorCodes.RAG_KNOWLEDGE_BASE_NOT_FOUND)
|
||||
return kb
|
||||
Reference in New Issue
Block a user