74 lines
2.4 KiB
Python
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")
|