ailine/backend/app/graph/subgraph_builder.py

"""
子图整合主图构建器
Subgraph Integration Main Graph Builder
"""

from langgraph.graph import StateGraph, START, END
from typing import Dict, Any

from .state import MainGraphState, CurrentAction
from ..agent_subgraphs.contact import build_contact_subgraph
from ..agent_subgraphs.dictionary import build_dictionary_subgraph
from ..agent_subgraphs.news_analysis import build_news_analysis_subgraph


def parse_user_intent(state: MainGraphState) -> MainGraphState:
    """
    解析用户意图节点
    
    确定该路由到哪个子图
    """
    state.current_phase = "intent_parsing"
    
    # 从messages中提取用户查询（如果user_query为空）
    if not state.user_query and state.messages:
        # 获取最后一条消息的内容
        last_msg = state.messages[-1]
        state.user_query = last_msg.content
    
    query_lower = state.user_query.lower()
    
    # 简单的关键词匹配
    if any(keyword in query_lower for keyword in ["通讯录", "联系人", "contact", "email"]):
        state.current_action = CurrentAction.CONTACT
        state.intent_confidence = 0.9
    
    elif any(keyword in query_lower for keyword in ["词典", "单词", "翻译", "dictionary", "translate"]):
        state.current_action = CurrentAction.DICTIONARY
        state.intent_confidence = 0.9
    
    elif any(keyword in query_lower for keyword in ["资讯", "新闻", "分析", "news", "report"]):
        state.current_action = CurrentAction.NEWS_ANALYSIS
        state.intent_confidence = 0.9
    
    else:
        # 默认是普通聊天
        state.current_action = CurrentAction.GENERAL_CHAT
        state.intent_confidence = 0.8
    
    return state


def route_to_subgraph(state: MainGraphState) -> str:
    """
    条件路由：决定路由到哪个子图
    """
    if state.current_action == CurrentAction.NONE:
        return "general_chat"
    elif state.current_action == CurrentAction.GENERAL_CHAT:
        return "general_chat"
    elif state.current_action == CurrentAction.CONTACT:
        return "contact_subgraph"
    elif state.current_action == CurrentAction.DICTIONARY:
        return "dictionary_subgraph"
    elif state.current_action == CurrentAction.NEWS_ANALYSIS:
        return "news_analysis_subgraph"
    else:
        return "general_chat"


def general_chat_node(state: MainGraphState) -> MainGraphState:
    """
    普通聊天节点
    （目前是占位符，后续整合旧的LLM调用逻辑）
    """
    state.current_phase = "general_chat"
    state.final_result = f"普通聊天模式：{state.user_query}"
    state.success = True
    return state


def integrate_results(state: MainGraphState) -> MainGraphState:
    """
    整合子图结果节点
    """
    state.current_phase = "integrating"
    
    # 整合通讯录子图结果
    if state.contact_result:
        state.final_result = state.contact_result.get("final_result", "")
    
    # 整合词典子图结果
    elif state.dictionary_result:
        state.final_result = state.dictionary_result.get("final_result", "")
    
    # 整合资讯子图结果
    elif state.news_result:
        state.final_result = state.news_result.get("final_result", "")
    
    else:
        # 没有子图结果
        if not state.final_result:
            state.final_result = "处理完成"
    
    state.current_phase = "done"
    return state


def build_main_graph() -> StateGraph:
    """
    构建整合了子图的主图
    
    Returns:
        配置好的 StateGraph
    """
    # 创建图
    graph = StateGraph(MainGraphState)
    
    # 添加节点
    graph.add_node("parse_intent", parse_user_intent)
    graph.add_node("general_chat", general_chat_node)
    graph.add_node("integrate_results", integrate_results)
    
    # 添加子图节点
    contact_graph = build_contact_subgraph()
    dictionary_graph = build_dictionary_subgraph()
    news_analysis_graph = build_news_analysis_subgraph()
    
    graph.add_node("contact_subgraph", contact_graph.compile())
    graph.add_node("dictionary_subgraph", dictionary_graph.compile())
    graph.add_node("news_analysis_subgraph", news_analysis_graph.compile())
    
    # 添加边
    # 从START开始
    graph.add_edge(START, "parse_intent")
    
    # 从parse_intent根据条件路由
    graph.add_conditional_edges(
        "parse_intent",
        route_to_subgraph,
        {
            "general_chat": "general_chat",
            "contact_subgraph": "contact_subgraph",
            "dictionary_subgraph": "dictionary_subgraph",
            "news_analysis_subgraph": "news_analysis_subgraph",
        }
    )
    
    # 从普通聊天和子图到结果整合
    graph.add_edge("general_chat", "integrate_results")
    graph.add_edge("contact_subgraph", "integrate_results")
    graph.add_edge("dictionary_subgraph", "integrate_results")
    graph.add_edge("news_analysis_subgraph", "integrate_results")
    
    # 最终到END
    graph.add_edge("integrate_results", END)
    
    return graph