Build lightweight AI agent admin
This commit is contained in:
Codex
@@ -0,0 +1,238 @@
|
||||
"""
|
||||
知识库检索节点
|
||||
|
||||
在 AI 工作流中检索知识库，返回与查询最相关的文档分段
|
||||
支持向量检索、全文检索和混合检索
|
||||
"""
|
||||
import logging
|
||||
import time
|
||||
from typing import Any, Dict, List
|
||||
|
||||
from ..base import BaseNode, NodeContext, NodeResult
|
||||
from ..registry import NodeRegistry
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@NodeRegistry.register
|
||||
class KnowledgeRetrievalNode(BaseNode):
|
||||
"""
|
||||
知识库检索节点
|
||||
|
||||
从指定知识库中检索与查询文本最相关的文档分段，
|
||||
输出可直接作为 LLM 节点的上下文使用
|
||||
"""
|
||||
|
||||
node_type = 'knowledge_retrieval'
|
||||
node_name = '知识库检索'
|
||||
node_category = 'knowledge'
|
||||
node_icon = 'BookOpen'
|
||||
node_description = '从知识库中检索相关文档，为 LLM 提供上下文'
|
||||
|
||||
inputs = [
|
||||
{
|
||||
'name': 'query',
|
||||
'type': 'string',
|
||||
'description': '检索查询文本',
|
||||
},
|
||||
]
|
||||
|
||||
outputs = [
|
||||
{
|
||||
'name': 'results',
|
||||
'type': 'array',
|
||||
'description': '检索结果列表',
|
||||
},
|
||||
{
|
||||
'name': 'context',
|
||||
'type': 'string',
|
||||
'description': '拼接后的上下文文本',
|
||||
},
|
||||
]
|
||||
|
||||
def execute(self, context: NodeContext) -> NodeResult:
|
||||
"""同步执行（不支持，需要异步）"""
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='知识库检索节点必须异步执行',
|
||||
)
|
||||
|
||||
async def execute_async(self, context: NodeContext) -> NodeResult:
|
||||
"""异步执行知识库检索"""
|
||||
start_time = time.time()
|
||||
|
||||
try:
|
||||
# 读取配置
|
||||
knowledge_base_ids = self.config.get('knowledge_base_ids', [])
|
||||
query_template = self.config.get('query', '')
|
||||
top_k = self.config.get('top_k', 5)
|
||||
score_threshold = self.config.get('score_threshold', 0.5)
|
||||
retrieval_mode = self.config.get('retrieval_mode', None)
|
||||
rerank_enabled = self.config.get('rerank_enabled', None)
|
||||
rerank_model_id = self.config.get('rerank_model_id', None)
|
||||
output_variable = self.config.get('output_variable', 'knowledge_results')
|
||||
context_variable = self.config.get('context_variable', 'knowledge_context')
|
||||
context_template = self.config.get('context_template', '')
|
||||
|
||||
if not knowledge_base_ids:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='未配置知识库',
|
||||
elapsed_time=int((time.time() - start_time) * 1000),
|
||||
)
|
||||
|
||||
# 解析查询模板中的变量
|
||||
query = context.resolve_template(query_template) if query_template else context.user_input
|
||||
if not query or not query.strip():
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='检索查询文本为空',
|
||||
elapsed_time=int((time.time() - start_time) * 1000),
|
||||
)
|
||||
|
||||
# 调用检索服务
|
||||
db = context.db_session
|
||||
if not db:
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error='数据库会话不可用',
|
||||
elapsed_time=int((time.time() - start_time) * 1000),
|
||||
)
|
||||
|
||||
from ai_platform.knowledge.services.retrieval_service import RetrievalService
|
||||
|
||||
service = RetrievalService(db)
|
||||
results = await service.retrieve(
|
||||
query=query,
|
||||
knowledge_base_ids=knowledge_base_ids,
|
||||
top_k=top_k,
|
||||
score_threshold=score_threshold,
|
||||
retrieval_mode=retrieval_mode,
|
||||
rerank_enabled=rerank_enabled,
|
||||
rerank_model_id=rerank_model_id,
|
||||
)
|
||||
|
||||
# 构建结果列表
|
||||
result_list = []
|
||||
for r in results:
|
||||
result_list.append({
|
||||
'segment_id': r.segment_id,
|
||||
'document_id': r.document_id,
|
||||
'document_name': r.document_name,
|
||||
'knowledge_base_id': r.knowledge_base_id,
|
||||
'knowledge_base_name': r.knowledge_base_name,
|
||||
'content': r.content,
|
||||
'score': r.score,
|
||||
'token_count': r.token_count,
|
||||
})
|
||||
|
||||
# 构建上下文文本
|
||||
if context_template:
|
||||
# 自定义模板
|
||||
context_text = context.resolve_template(context_template)
|
||||
else:
|
||||
# 默认：拼接所有检索结果内容
|
||||
context_parts = []
|
||||
for i, r in enumerate(result_list, 1):
|
||||
context_parts.append(
|
||||
f"[{i}] (来源: {r['document_name']}, 相似度: {r['score']:.2f})\n{r['content']}"
|
||||
)
|
||||
context_text = '\n\n'.join(context_parts)
|
||||
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
|
||||
return NodeResult(
|
||||
success=True,
|
||||
output={
|
||||
'results': result_list,
|
||||
'context': context_text,
|
||||
'total': len(result_list),
|
||||
'query': query,
|
||||
},
|
||||
output_variables={
|
||||
output_variable: result_list,
|
||||
context_variable: context_text,
|
||||
f'{output_variable}_total': len(result_list),
|
||||
},
|
||||
elapsed_time=elapsed_time,
|
||||
metadata={
|
||||
'query': query,
|
||||
'knowledge_base_ids': knowledge_base_ids,
|
||||
'top_k': top_k,
|
||||
'score_threshold': score_threshold,
|
||||
'retrieval_mode': retrieval_mode,
|
||||
'rerank_enabled': rerank_enabled,
|
||||
'result_count': len(result_list),
|
||||
},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(f'知识库检索节点执行失败: {e}')
|
||||
elapsed_time = int((time.time() - start_time) * 1000)
|
||||
return NodeResult(
|
||||
success=False,
|
||||
error=f'检索失败: {str(e)}',
|
||||
elapsed_time=elapsed_time,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_config_schema(cls) -> Dict[str, Any]:
|
||||
"""获取配置 Schema"""
|
||||
return {
|
||||
'type': 'object',
|
||||
'properties': {
|
||||
'knowledge_base_ids': {
|
||||
'type': 'array',
|
||||
'title': '知识库',
|
||||
'description': '选择要检索的知识库',
|
||||
'items': {'type': 'string'},
|
||||
},
|
||||
'query': {
|
||||
'type': 'string',
|
||||
'title': '查询文本',
|
||||
'description': '支持变量引用，如 {{user_input}}，留空则使用用户输入',
|
||||
},
|
||||
'top_k': {
|
||||
'type': 'integer',
|
||||
'title': '返回数量',
|
||||
'default': 5,
|
||||
'minimum': 1,
|
||||
'maximum': 20,
|
||||
},
|
||||
'score_threshold': {
|
||||
'type': 'number',
|
||||
'title': '相似度阈值',
|
||||
'default': 0.5,
|
||||
'minimum': 0,
|
||||
'maximum': 1,
|
||||
},
|
||||
'retrieval_mode': {
|
||||
'type': 'string',
|
||||
'title': '检索模式',
|
||||
'description': '留空则使用知识库默认配置',
|
||||
'enum': ['vector', 'fulltext', 'hybrid'],
|
||||
},
|
||||
'rerank_enabled': {
|
||||
'type': 'boolean',
|
||||
'title': '启用重排序',
|
||||
'description': '不设置则使用知识库默认配置',
|
||||
'default': None,
|
||||
},
|
||||
'rerank_model_id': {
|
||||
'type': 'string',
|
||||
'title': '重排序模型',
|
||||
'description': '不设置则使用知识库默认配置',
|
||||
},
|
||||
'output_variable': {
|
||||
'type': 'string',
|
||||
'title': '结果变量名',
|
||||
'default': 'knowledge_results',
|
||||
},
|
||||
'context_variable': {
|
||||
'type': 'string',
|
||||
'title': '上下文变量名',
|
||||
'default': 'knowledge_context',
|
||||
},
|
||||
},
|
||||
'required': ['knowledge_base_ids'],
|
||||
}
|
||||
Reference in New Issue
Block a user