PlantAI-dev/test_langgraph_state.py

"""
测试 LangGraph TypedDict state 传递机制

验证 execution_results 类型定义是否影响 state 传递
"""

import asyncio
from typing import Any, TypedDict
from langgraph.graph import StateGraph, END

# 正确的 TypedDict 定义
class CorrectGraphState(TypedDict):
    execution_results: list[dict[str, Any]]
    validation_result: dict[str, Any]
    current_agent: str | None

# 错误的 TypedDict 定义（之前的 bug）
class WrongGraphState(TypedDict):
    execution_results: dict[str, Any]  # 错误：应该是 list
    validation_result: dict[str, Any]
    current_agent: str | None


def execution_node(state: dict) -> dict:
    """模拟 Execution Agent 设置 execution_results"""
    print(f"[Execution] 输入 state keys: {list(state.keys())}")
    print(f"[Execution] 输入 execution_results: {state.get('execution_results', 'NOT FOUND')}")
    
    # 设置 list 类型
    execution_results = [
        {
            "stage": "identify",
            "result": {
                "status": "success",
                "manifest": {"stage": "identify", "status": "success"},
                "message": "✅ Identify completed"
            }
        }
    ]
    
    state["execution_results"] = execution_results
    state["current_agent"] = "execution"
    
    print(f"[Execution] 设置 execution_results: {len(execution_results)} items")
    print(f"[Execution] 输出 state['execution_results']: {state.get('execution_results')}")
    
    return state


def validator_node(state: dict) -> dict:
    """模拟 Validator Agent 读取 execution_results"""
    print(f"[Validator] 输入 state keys: {list(state.keys())}")
    print(f"[Validator] 输入 execution_results type: {type(state.get('execution_results'))}")
    print(f"[Validator] 输入 execution_results: {state.get('execution_results', 'NOT FOUND')}")
    
    execution_results = state.get("execution_results", [])
    
    if isinstance(execution_results, list) and len(execution_results) > 0:
        last_execution = execution_results[-1]
        execution_result = last_execution.get("result", {})
        status = execution_result.get("status", "unknown")
        
        overall_passed = status == "success"
        print(f"[Validator] ✅ execution_results 有数据，status={status}, overall_passed={overall_passed}")
    else:
        print(f"[Validator] ⚠️ execution_results 为空或类型错误")
        overall_passed = False
    
    state["validation_result"] = {
        "overall_passed": overall_passed,
        "stage": "identify"
    }
    state["current_agent"] = "validator"
    
    return state


def validate_decision(state: dict) -> str:
    """模拟 _validate_decision 决策"""
    validation = state.get("validation_result", {})
    overall_passed = validation.get("overall_passed", False)
    
    print(f"[Decision] validation_result: {validation}")
    print(f"[Decision] overall_passed: {overall_passed}")
    
    if overall_passed:
        return "success"
    else:
        return "failure"


async def test_with_correct_state():
    """测试正确的 TypedDict 定义"""
    print("\n" + "="*60)
    print("测试 1: 正确的 TypedDict 定义 (list[dict[str, Any]])")
    print("="*60)
    
    workflow = StateGraph(CorrectGraphState)
    workflow.add_node("execution", execution_node)
    workflow.add_node("validator", validator_node)
    workflow.add_conditional_edges(
        "validator",
        validate_decision,
        {
            "success": END,
            "failure": END
        }
    )
    workflow.add_edge("execution", "validator")
    workflow.set_entry_point("execution")
    
    graph = workflow.compile()
    
    initial_state = {
        "execution_results": [],
        "validation_result": {},
        "current_agent": None
    }
    
    result = await graph.ainvoke(initial_state)
    
    print(f"\n[结果] 最终 state: {result}")
    print(f"[结果] execution_results type: {type(result.get('execution_results'))}")
    print(f"[结果] validation_result: {result.get('validation_result')}")
    
    return result


async def test_with_wrong_state():
    """测试错误的 TypedDict 定义"""
    print("\n" + "="*60)
    print("测试 2: 错误的 TypedDict 定义 (dict[str, Any])")
    print("="*60)
    
    workflow = StateGraph(WrongGraphState)
    workflow.add_node("execution", execution_node)
    workflow.add_node("validator", validator_node)
    workflow.add_conditional_edges(
        "validator",
        validate_decision,
        {
            "success": END,
            "failure": END
        }
    )
    workflow.add_edge("execution", "validator")
    workflow.set_entry_point("execution")
    
    graph = workflow.compile()
    
    initial_state = {
        "execution_results": {},  # 初始化为 dict（匹配错误定义）
        "validation_result": {},
        "current_agent": None
    }
    
    result = await graph.ainvoke(initial_state)
    
    print(f"\n[结果] 最终 state: {result}")
    print(f"[结果] execution_results type: {type(result.get('execution_results'))}")
    print(f"[结果] validation_result: {result.get('validation_result')}")
    
    return result


async def main():
    """运行所有测试"""
    print("\n" + "="*80)
    print("LangGraph TypedDict State 传递测试")
    print("="*80)
    
    result_correct = await test_with_correct_state()
    result_wrong = await test_with_wrong_state()
    
    print("\n" + "="*80)
    print("测试结果对比")
    print("="*80)
    
    correct_passed = result_correct.get("validation_result", {}).get("overall_passed", False)
    wrong_passed = result_wrong.get("validation_result", {}).get("overall_passed", False)
    
    print(f"正确定义: overall_passed = {correct_passed}")
    print(f"错误定义: overall_passed = {wrong_passed}")
    
    if correct_passed and not wrong_passed:
        print("\n✅ 验证成功：TypedDict 类型定义影响 LangGraph state 传递！")
        print("   正确定义 (list) -> state 正确传递 -> validation_passed=True")
        print("   错误定义 (dict) -> state 传递异常 -> validation_passed=False")
    elif correct_passed and wrong_passed:
        print("\n⚠️ 两种定义都能正常工作，问题可能在其他地方")
    else:
        print("\n❌ 测试失败：两种定义都没有正常工作")


if __name__ == "__main__":
    asyncio.run(main())