geo/backend/app/agent_framework/tools/knowledge_tools.py

74 lines
2.4 KiB
Python

"""Knowledge 业务工具 - 将知识库服务注册为 FunctionTool"""
import logging
from typing import Any
from agentkit.tools.function_tool import FunctionTool
from agentkit.tools.registry import ToolRegistry
logger = logging.getLogger(__name__)
async def search_knowledge(
query: str,
knowledge_base_ids: list[str],
top_k: int = 5,
) -> dict:
"""从知识库检索相关内容"""
if not knowledge_base_ids or not query:
return {"content": "暂无相关知识库内容", "sources": []}
try:
from app.database import AsyncSessionLocal
from app.services.knowledge.rag_service import RAGService
async with AsyncSessionLocal() as session:
rag = RAGService()
results = await rag.search(
session=session,
query=query,
knowledge_base_ids=knowledge_base_ids,
top_k=top_k,
)
if results:
content_parts = []
sources = []
for r in results:
title = r.get("document_title", "未知")
content_parts.append(f"[来源: {title}]\n{r.get('content', '')}")
sources.append(title)
return {"content": "\n\n---\n\n".join(content_parts), "sources": sources}
except Exception as e:
logger.warning(f"知识库检索失败: {e}")
return {"content": "暂无相关知识库内容", "sources": []}
async def detect_ai_patterns(content: str, platform_id: str) -> dict:
"""检测内容中的AI生成模式"""
from app.services.distribution.rule_service import platform_rule_service
patterns = platform_rule_service.detect_ai_patterns(content, platform_id)
return {"patterns": patterns, "count": len(patterns)}
def register_knowledge_tools(registry: ToolRegistry) -> None:
"""注册所有知识库相关工具"""
registry.register(
FunctionTool(
name="search_knowledge",
description="从知识库检索相关内容",
func=search_knowledge,
tags=["knowledge", "rag"],
)
)
registry.register(
FunctionTool(
name="detect_ai_patterns",
description="检测内容中的AI生成模式",
func=detect_ai_patterns,
tags=["knowledge", "deai"],
)
)
logger.info("Knowledge tools registered")