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

"""
资讯子图构建器
News Analysis Subgraph Builder
"""

from langgraph.graph import StateGraph, START, END

from .state import NewsAnalysisState
from .nodes import (
    parse_intent,
    query_news,
    analyze_url,
    extract_keywords,
    generate_report,
    format_result,
    should_continue
)


def build_news_analysis_subgraph() -> StateGraph:
    """
    构建资讯子图
    
    Returns:
        配置好的 StateGraph
    """
    # 创建图
    graph = StateGraph(NewsAnalysisState)
    
    # 添加节点
    graph.add_node("parse_intent", parse_intent)
    graph.add_node("query_news", query_news)
    graph.add_node("analyze_url", analyze_url)
    graph.add_node("extract_keywords", extract_keywords)
    graph.add_node("generate_report", generate_report)
    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_news": "query_news",
            "analyze_url": "analyze_url",
            "extract_keywords": "extract_keywords",
            "generate_report": "generate_report",
        }
    )
    
    # 从各个操作节点到format_result
    graph.add_edge("query_news", "format_result")
    graph.add_edge("analyze_url", "format_result")
    graph.add_edge("extract_keywords", "format_result")
    graph.add_edge("generate_report", "format_result")
    
    # 最终到END
    graph.add_edge("format_result", END)
    
    return graph