ailine/backend/app/graph/subgraph_tools.py

"""
子图工具模块 - 将三个子图包装成 LangChain 工具
Subgraph Tools Module - Wrap three subgraphs as LangChain tools
"""

from langchain_core.tools import tool
from typing import Optional


# ============== 词典子图工具 ==============
@tool
def dictionary_tool(query: str, action: Optional[str] = None) -> str:
    """
    词典/翻译工具 - 查询单词、翻译文本、提取术语、获取每日一词
    
    Args:
        query: 用户查询内容（单词、句子、文本等）
        action: 可选，指定操作类型（"query" 查询单词，"translate" 翻译，
                "extract" 提取术语，"daily" 每日一词，不指定则自动识别）
    
    Returns:
        格式化的结果文本
    """
    try:
        from backend.app.agent_subgraphs.dictionary import (
            DictionaryState,
            DictionaryAction,
            parse_intent,
            format_result
        )
        from backend.app.agent_subgraphs.dictionary.nodes import (
            query_word, translate_text, extract_terms, get_daily_word
        )
        
        # 创建初始状态
        state = DictionaryState(user_query=query, user_id="default")
        
        # 处理 action
        if action == "query":
            state.action = DictionaryAction.QUERY
            state.action_params = {"word": query}
            state = query_word(state)
        elif action == "translate":
            state.action = DictionaryAction.TRANSLATE
            state.source_text = query
            state = translate_text(state)
        elif action == "daily":
            state.action = DictionaryAction.DAILY_WORD
            state = get_daily_word(state)
        elif action == "extract":
            state.action = DictionaryAction.EXTRACT
            state.action_params = {"text": query}
            state = extract_terms(state)
        else:
            # 自动解析意图
            state = parse_intent(state)
            if state.action == DictionaryAction.QUERY:
                state = query_word(state)
            elif state.action == DictionaryAction.TRANSLATE:
                state = translate_text(state)
            elif state.action == DictionaryAction.DAILY_WORD:
                state = get_daily_word(state)
            elif state.action == DictionaryAction.EXTRACT:
                state = extract_terms(state)
        
        # 格式化结果
        state = format_result(state)
        
        return state.final_result or "操作完成"
        
    except Exception as e:
        return f"词典工具执行出错：{str(e)}"


# ============== 资讯分析子图工具 ==============
@tool
def news_analysis_tool(query: str, action: Optional[str] = None) -> str:
    """
    资讯分析工具 - 查询新闻、分析URL、提取关键词、生成报告
    
    Args:
        query: 用户查询内容（关键词、URL、文本等）
        action: 可选，指定操作类型（"query" 查询新闻，"analyze" 分析URL，
                "keywords" 提取关键词，"report" 生成报告，不指定则自动识别）
    
    Returns:
        格式化的结果文本
    """
    try:
        from backend.app.agent_subgraphs.news_analysis import (
            NewsAnalysisState,
            NewsAction,
            parse_intent,
            format_result
        )
        from backend.app.agent_subgraphs.news_analysis.nodes import (
            query_news, analyze_url, extract_keywords, generate_report
        )
        
        # 创建初始状态
        state = NewsAnalysisState(user_query=query, user_id="default")
        
        # 处理 action
        if action == "query":
            state.action = NewsAction.QUERY_NEWS
            state = query_news(state)
        elif action == "analyze":
            state.action = NewsAction.ANALYZE_URL
            state.custom_urls = [query]
            state = analyze_url(state)
        elif action == "keywords":
            state.action = NewsAction.EXTRACT_KEYWORDS
            state = extract_keywords(state)
        elif action == "report":
            state.action = NewsAction.GENERATE_REPORT
            state = generate_report(state)
        else:
            # 自动解析意图
            state = parse_intent(state)
            if state.action == NewsAction.QUERY_NEWS:
                state = query_news(state)
            elif state.action == NewsAction.ANALYZE_URL:
                state.custom_urls = [query]
                state = analyze_url(state)
            elif state.action == NewsAction.EXTRACT_KEYWORDS:
                state = extract_keywords(state)
            elif state.action == NewsAction.GENERATE_REPORT:
                state = generate_report(state)
        
        # 格式化结果
        state = format_result(state)
        
        return state.final_result or "操作完成"
        
    except Exception as e:
        return f"资讯分析工具执行出错：{str(e)}"


# ============== 通讯录子图工具 ==============
@tool
def contact_tool(query: str, action: Optional[str] = None) -> str:
    """
    通讯录工具 - 查询联系人、添加联系人、管理通讯录
    
    Args:
        query: 用户查询内容（姓名、电话、信息等）
        action: 可选，指定操作类型（不指定则自动识别）
    
    Returns:
        格式化的结果文本
    """
    try:
        from backend.app.agent_subgraphs.contact import (
            ContactState,
            ContactAction,
            parse_intent,
            format_result
        )
        from backend.app.agent_subgraphs.contact.nodes import (
            query_contact, add_contact, list_contacts
        )
        
        # 创建初始状态
        state = ContactState(user_query=query, user_id="default")
        
        # 自动解析意图
        state = parse_intent(state)
        
        # 根据解析结果执行操作
        if state.action == ContactAction.QUERY:
            state = query_contact(state)
        elif state.action == ContactAction.ADD:
            state = add_contact(state)
        elif state.action == ContactAction.LIST:
            state = list_contacts(state)
        
        # 格式化结果
        state = format_result(state)
        
        return state.final_result or "操作完成"
        
    except Exception as e:
        return f"通讯录工具执行出错：{str(e)}"


# ============== 工具列表 ==============
SUBGRAPH_TOOLS = [
    dictionary_tool,
    news_analysis_tool,
    contact_tool
]

SUBGRAPH_TOOLS_BY_NAME = {tool.name: tool for tool in SUBGRAPH_TOOLS}