tx_flow_analysis/examples/state_machine.py

#!/usr/bin/env python3
"""
LangGraph状态机示例 - 学习工作流设计
==================================

这个文件展示了如何使用LangGraph创建状态机工作流，包含：
1. 状态定义
2. 节点函数
3. 条件路由
4. 工作流执行

运行方法：
python examples/state_machine.py
"""

import sys
from typing import TypedDict, Annotated, List

try:
    from langgraph.graph import StateGraph, START, END
except ImportError as e:
    print(f"❌ 缺少依赖包: {e}")
    print("请运行: pip install langgraph")
    sys.exit(1)


class AnalysisState(TypedDict):
    """分析工作流的状态"""
    question: str
    data: str
    current_step: str
    analysis_result: str
    report: str
    steps_completed: Annotated[List[str], "add"]  # 使用注解支持列表追加


def planning_node(state: AnalysisState) -> AnalysisState:
    """
    规划节点 - 分析用户需求并制定计划

    这是一个决策节点，决定下一步的分析方向
    """
    print("📋 规划阶段: 分析用户需求")

    question = state["question"].lower()

    # 简单的决策逻辑
    if "趋势" in question or "变化" in question:
        next_step = "trend_analysis"
        plan = "执行趋势分析"
    elif "对比" in question or "比较" in question:
        next_step = "comparison_analysis"
        plan = "执行对比分析"
    elif "汇总" in question or "总结" in question:
        next_step = "summary_analysis"
        plan = "执行汇总分析"
    else:
        next_step = "general_analysis"
        plan = "执行一般性分析"

    print(f"   决策: {plan}")

    return {
        **state,
        "current_step": next_step,
        "steps_completed": ["planning"]
    }


def data_processing_node(state: AnalysisState) -> AnalysisState:
    """数据处理节点 - 准备数据用于分析"""
    print("🔧 数据处理阶段: 准备分析数据")

    # 模拟数据处理
    processed_data = f"已处理的数据: {state['data'][:50]}..."

    print("   数据处理完成")

    return {
        **state,
        "data": processed_data,
        "steps_completed": ["data_processing"]
    }


def trend_analysis_node(state: AnalysisState) -> AnalysisState:
    """趋势分析节点 - 专门处理趋势分析"""
    print("📈 趋势分析阶段: 识别数据变化趋势")

    # 模拟趋势分析
    result = f"趋势分析结果: 数据显示上升趋势，增长率为15%"

    print(f"   分析结果: {result}")

    return {
        **state,
        "analysis_result": result,
        "steps_completed": ["trend_analysis"]
    }


def comparison_analysis_node(state: AnalysisState) -> AnalysisState:
    """对比分析节点 - 专门处理对比分析"""
    print("⚖️ 对比分析阶段: 比较不同数据集")

    # 模拟对比分析
    result = f"对比分析结果: A组数据优于B组数据，差异显著"

    print(f"   分析结果: {result}")

    return {
        **state,
        "analysis_result": result,
        "steps_completed": ["comparison_analysis"]
    }


def general_analysis_node(state: AnalysisState) -> AnalysisState:
    """一般分析节点 - 处理通用分析需求"""
    print("🔍 一般分析阶段: 执行标准数据分析")

    # 模拟一般分析
    result = f"一般分析结果: 数据整体表现良好，符合预期"

    print(f"   分析结果: {result}")

    return {
        **state,
        "analysis_result": result,
        "steps_completed": ["general_analysis"]
    }


def report_generation_node(state: AnalysisState) -> AnalysisState:
    """报告生成节点 - 生成最终分析报告"""
    print("📄 报告生成阶段: 整理分析结果")

    # 生成报告
    report = f"""
📊 数据分析报告
================

问题: {state['question']}
数据: {state['data']}
分析结果: {state['analysis_result']}

执行步骤: {', '.join(state['steps_completed'])}

结论: 分析完成，建议根据具体情况采取相应措施。
"""

    print("   报告生成完成")

    return {
        **state,
        "report": report.strip(),
        "steps_completed": ["report_generation"]
    }


def route_from_planning(state: AnalysisState) -> str:
    """
    从规划节点路由到具体分析节点

    这是一个条件路由函数，根据规划结果决定下一个节点
    """
    return state["current_step"]


def create_analysis_workflow() -> StateGraph:
    """创建分析工作流"""

    # 创建状态图
    workflow = StateGraph(AnalysisState)

    # 添加节点
    workflow.add_node("planning", planning_node)
    workflow.add_node("data_processing", data_processing_node)
    workflow.add_node("trend_analysis", trend_analysis_node)
    workflow.add_node("comparison_analysis", comparison_analysis_node)
    workflow.add_node("general_analysis", general_analysis_node)
    workflow.add_node("report_generation", report_generation_node)

    # 设置入口点
    workflow.set_entry_point("planning")

    # 添加固定边
    workflow.add_edge("data_processing", "report_generation")
    workflow.add_edge("trend_analysis", "report_generation")
    workflow.add_edge("comparison_analysis", "report_generation")
    workflow.add_edge("general_analysis", "report_generation")
    workflow.add_edge("report_generation", END)

    # 添加条件边
    workflow.add_conditional_edges(
        "planning",
        route_from_planning,
        {
            "trend_analysis": "data_processing",  # 趋势分析需要先处理数据
            "comparison_analysis": "data_processing",  # 对比分析也需要先处理数据
            "general_analysis": "data_processing",  # 一般分析也需要先处理数据
        }
    )

    return workflow


def main():
    """主函数 - 演示状态机工作流"""
    print("🚀 LangGraph状态机示例")
    print("=" * 50)

    try:
        # 创建工作流
        workflow = create_analysis_workflow()
        app = workflow.compile()

        # 测试用例
        test_cases = [
            {
                "question": "这个季度的数据趋势如何？",
                "data": "Q1: 100, Q2: 120, Q3: 140, Q4: 160"
            },
            {
                "question": "比较A产品和B产品的销量",
                "data": "A产品: 500件, B产品: 450件"
            },
            {
                "question": "请分析这份数据",
                "data": "各类数据指标: 正常范围"
            }
        ]

        print("\n🧪 测试不同类型的分析:")
        print("-" * 40)

        for i, test_case in enumerate(test_cases, 1):
            print(f"\n📋 测试用例 {i}:")
            print(f"问题: {test_case['question']}")
            print(f"数据: {test_case['data']}")

            # 执行工作流
            result = app.invoke({
                "question": test_case["question"],
                "data": test_case["data"],
                "current_step": "",
                "analysis_result": "",
                "report": "",
                "steps_completed": []
            })

            print("\n✅ 执行完成:")
            print(f"执行步骤: {result['steps_completed']}")
            print(f"分析结果: {result['analysis_result']}")

            # 显示报告预览
            report_preview = result['report'][:200] + "..." if len(result['report']) > 200 else result['report']
            print(f"报告预览: {report_preview}")

        print("\n🎉 状态机示例完成！")
        print("\n💡 LangGraph学习要点:")
        print("1. 状态定义: 使用TypedDict定义工作流状态")
        print("2. 节点函数: 每个节点处理特定逻辑")
        print("3. 条件路由: 根据状态动态选择下一节点")
        print("4. 状态传递: 通过return更新状态")
        print("5. 工作流编译: 调用compile()创建可执行应用")

        print("\n📚 下一步学习:")
        print("- 查看 examples/advanced_agent.py - 学习高级Agent功能")
        print("- 阅读 llmops/ 目录下的实际Agent代码")
        print("- 学习 PRACTICE_GUIDE.md 中的Phase 3内容")

    except Exception as e:
        print(f"❌ 运行出错: {e}")
        print("\n🔧 故障排除:")
        print("1. 确认已安装langgraph: pip install langgraph")
        print("2. 检查Python版本是否支持（推荐3.10+）")


if __name__ == "__main__":
    main()