deepdata: core product code (backend, frontend, runtime, deepflow module)
This commit is contained in:
figmar
@@ -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
|
||||
)
|
||||
Reference in New Issue
Block a user