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