ailine/backend/app/tools/__init__.py

"""
Agent Tools - 封装所有功能为 @tool 函数
"""

from langchain_core.tools import tool
from typing import Optional
from backend.app.logger import info


# ========== RAG Pipeline（复用现有）
_rag_pipeline = None


def _get_rag_pipeline():
    """获取 RAG Pipeline 实例（复用 rag_nodes.py 的逻辑）"""
    global _rag_pipeline
    if _rag_pipeline is None:
        from backend.app.rag.pipeline import RAGPipeline
        _rag_pipeline = RAGPipeline(
            num_queries=3,
            rerank_top_n=5,
            use_rerank=True,
            return_parent_docs=True,
        )
    return _rag_pipeline


@tool
async def rag_search(query: str) -> str:
    """
    检索知识库获取相关信息。
    
    当用户询问关于系统、业务、文档相关的问题时使用此工具。
    
    Args:
        query: 用户的问题或搜索关键词
    
    Returns:
        检索到的相关文档内容
    """
    info(f"[RAG Tool] 开始检索: {query[:50]}...")
    
    try:
        pipeline = _get_rag_pipeline()
        documents = await pipeline.aretrieve(query)
        rag_context = pipeline.format_context(documents)
        
        info(f"[RAG Tool] 检索完成，得到 {len(documents)} 个文档")
        
        if rag_context:
            return rag_context
        else:
            return "知识库中没有找到相关内容。"
            
    except Exception as e:
        info(f"[RAG Tool] 检索失败: {e}")
        return f"知识库检索失败: {str(e)}"


@tool
def web_search(query: str) -> str:
    """
    联网搜索获取最新信息。
    
    当用户询问实时新闻、热点事件、最新资讯或知识库中没有的内容时使用此工具。
    
    Args:
        query: 搜索关键词
    
    Returns:
        搜索结果摘要
    """
    info(f"[WebSearch Tool] 开始搜索: {query[:50]}...")
    
    try:
        from backend.app.core import web_search as core_web_search
        search_result = core_web_search(query, max_results=5)
        
        info(f"[WebSearch Tool] 搜索完成")
        return search_result
        
    except Exception as e:
        info(f"[WebSearch Tool] 搜索失败: {e}")
        return f"联网搜索失败: {str(e)}"


# ====== 子图工具封装器
async def _invoke_subgraph(subgraph_builder, query: str, state_class) -> str:
    """
    通用子图调用函数
    
    Args:
        subgraph_builder: 子图构建函数
        query: 用户查询
        state_class: 子图状态类
    
    Returns:
        子图执行结果
    """
    try:
        graph = subgraph_builder()
        compiled_graph = graph.compile()
        
        # 构造初始状态
        initial_state = state_class(user_query=query)
        
        # 调用子图
        result = await compiled_graph.ainvoke(initial_state)
        
        # 返回结果
        return result.get("final_result", "子图执行完成")
        
    except Exception as e:
        info(f"[Subgraph Tool] 执行失败: {e}")
        return f"执行失败: {str(e)}"


@tool
async def contact_lookup(query: str) -> str:
    """
    查询通讯录信息。
    
    当用户询问联系人、邮箱、联系方式、发送邮件时使用此工具。
    
    Args:
        query: 用户查询，描述需要的操作
    
    Returns:
        联系人信息或操作结果
    """
    info(f"[Contact Tool] 查询: {query[:50]}...")
    
    from backend.app.subgraphs.contact.graph import build_contact_subgraph
    from backend.app.subgraphs.contact.state import ContactState
    
    return await _invoke_subgraph(build_contact_subgraph, query, ContactState)


@tool
async def dictionary_lookup(word: str) -> str:
    """
    查询词典，获取单词释义、翻译等。
    
    当用户询问单词、翻译、生词时使用此工具。
    
    Args:
        word: 需要查询的单词或短语
    
    Returns:
        单词释义和翻译
    """
    info(f"[Dictionary Tool] 查询: {word}")
    
    from backend.app.subgraphs.dictionary.graph import build_dictionary_subgraph
    from backend.app.subgraphs.dictionary.state import DictionaryState
    
    return await _invoke_subgraph(build_dictionary_subgraph, word, DictionaryState)


@tool
async def news_analysis(topic: str) -> str:
    """
    分析热点新闻和资讯。
    
    当用户询问新闻分析、热点解读时使用此工具。
    
    Args:
        topic: 新闻主题或关键词
    
    Returns:
        新闻分析结果
    """
    info(f"[NewsAnalysis Tool] 分析: {topic}")
    
    from backend.app.subgraphs.news_analysis.graph import build_news_analysis_subgraph
    from backend.app.subgraphs.news_analysis.state import NewsAnalysisState
    
    return await _invoke_subgraph(build_news_analysis_subgraph, topic, NewsAnalysisState)


# ====== 导出所有工具
ALL_TOOLS = [
    rag_search,
    web_search,
    contact_lookup,
    dictionary_lookup,
    news_analysis,
]