ailine/backend/app/subgraphs/dictionary/graph.py

"""
词典子图构建器 - 完善版
Dictionary Subgraph Builder - Complete
"""

from langgraph.graph import StateGraph, START, END

from .state import DictionaryState
from .nodes import (
    parse_intent,
    query_word,
    translate_text,
    extract_terms,
    get_daily_word,
    lookup_word_book,
    add_to_word_book,
    format_result,
    should_continue
)


def build_dictionary_subgraph() -> StateGraph:
    """
    构建词典子图
    
    Returns:
        配置好的 StateGraph
    """
    # 创建图
    graph = StateGraph(DictionaryState)
    
    # 添加节点
    graph.add_node("parse_intent", parse_intent)
    graph.add_node("query_word", query_word)
    graph.add_node("translate_text", translate_text)
    graph.add_node("extract_terms", extract_terms)
    graph.add_node("get_daily_word", get_daily_word)
    graph.add_node("lookup_word_book", lookup_word_book)
    graph.add_node("add_to_word_book", add_to_word_book)
    graph.add_node("format_result", format_result)
    
    # 添加边
    # 从START开始
    graph.add_edge(START, "parse_intent")
    
    # 从parse_intent根据条件路由
    graph.add_conditional_edges(
        "parse_intent",
        should_continue,
        {
            "query_word": "query_word",
            "translate_text": "translate_text",
            "extract_terms": "extract_terms",
            "get_daily_word": "get_daily_word",
            "lookup_word_book": "lookup_word_book",
            "add_to_word_book": "add_to_word_book",
        }
    )
    
    # 从各个操作节点到format_result
    graph.add_edge("query_word", "format_result")
    graph.add_edge("translate_text", "format_result")
    graph.add_edge("extract_terms", "format_result")
    graph.add_edge("get_daily_word", "format_result")
    graph.add_edge("lookup_word_book", "format_result")
    graph.add_edge("add_to_word_book", "format_result")
    
    # 最终到END
    graph.add_edge("format_result", END)
    
    return graph