🎯 What You'll Learn Today

import asyncio
from typing import Annotated, Any, TypedDict
!pip install langchain-core
!pip install langgraph
from langchain_core.messages import BaseMessage, SystemMessage
from langchain_core.runnables import RunnableLambda
from langchain_core.tools import tool
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages

class State(TypedDict):
    """State for testing.

    Attributes:
        messages: Conversation history
        pending_tools: Tools awaiting execution
        results: Tool execution results
        errors: Error messages
        validation_results: Validation check results
    """

    messages: Annotated[list[BaseMessage], add_messages]
    pending_tools: list[dict[str, Any]]
    results: dict[str, Any]
    errors: dict[str, str]
    validation_results: dict[str, bool]

@tool
async def mock_tool(query: str) -> str:
    """Mock tool for testing.

    Args:
        query: Test query string

    Returns:
        Mock result string

    Raises:
        ValueError: If query contains "error"
    """
    await asyncio.sleep(0.1)  # Simulate latency
    if "error" in query:
        raise ValueError("Simulated error")
    return f"Mock result: {query}"

def validate_state(state: State) -> dict[str, bool]:
    """Validate state structure and content.

    Args:
        state: State to validate

    Returns:
        Dictionary of validation results
    """
    validations = {
        "has_messages": len(state.get("messages", [])) > 0,
        "has_valid_tools": all(
            {"id", "tool_name", "args"} <= set(t.keys())
            for t in state.get("pending_tools", [])
        ),
        "valid_results": all(
            isinstance(v, str) for v in state.get("results", {}).values()
        ),
    }
    return validations

async def mock_executor(state: State) -> State:
    """Execute mock tools with validation.

    Args:
        state: Current test state

    Returns:
        Updated state with results
    """
    if not state.get("pending_tools"):
        return state

    validations = validate_state(state)
    if not all(validations.values()):
        return {
            **state,
            "errors": {"validation": "State validation failed"},
            "validation_results": validations,
        }

    results = {}
    errors = {}

    for tool_call in state["pending_tools"]:
        try:
            result = await mock_tool.ainvoke(tool_call["args"]["query"])
            results[tool_call["id"]] = result
        except Exception as e:
            errors[tool_call["id"]] = str(e)

    return {
        **state,
        "results": results,
        "errors": errors,
        "validation_results": validations,
    }

def get_test_state(scenario: str = "basic") -> State:
    """Create test states for different scenarios.

    Args:
        scenario: Test scenario name

    Returns:
        State configured for scenario
    """
    states = {
        "basic": {
            "messages": [SystemMessage(content="Test execution")],
            "pending_tools": [
                {
                    "id": "test_1",
                    "tool_name": "mock_tool",
                    "args": {"query": "test query"},
                }
            ],
            "results": {},
            "errors": {},
            "validation_results": {},
        },
        "error": {
            "messages": [SystemMessage(content="Error test")],
            "pending_tools": [
                {
                    "id": "error_1",
                    "tool_name": "mock_tool",
                    "args": {"query": "error test"},
                }
            ],
            "results": {},
            "errors": {},
            "validation_results": {},
        },
        "invalid": {
            "messages": [],  # Invalid: no messages
            "pending_tools": [{"id": "invalid"}],  # Invalid structure
            "results": {},
            "errors": {},
            "validation_results": {},
        },
    }
    return states.get(scenario, states["basic"])

def create_test_graph() -> StateGraph:
    """Create test graph with validation.

    Returns:
        Configured StateGraph for testing
    """
    graph = StateGraph(State)

    # Add nodes with validation
    graph.add_node("executor", RunnableLambda(mock_executor))
    graph.add_node("validator", RunnableLambda(validate_state))

    # Configure edges
    graph.add_edge(START, "validator")
    graph.add_edge("validator", "executor")
    graph.add_edge("executor", END)

    return graph

async def run_test_scenario(scenario: str):
    """Run test with specific scenario.

    Args:
        scenario: Name of test scenario to run
    """
    graph = create_test_graph()
    chain = graph.compile()

    test_state = get_test_state(scenario)
    print(f"\nRunning {scenario} scenario:")
    print("Initial state:", test_state["pending_tools"])

    try:
        result = await chain.ainvoke(test_state)
        print("\nValidations:", result["validation_results"])
        print("Results:", result["results"])
        print("Errors:", result["errors"])
    except Exception as e:
        print(f"Test failed: {e!s}")

async def demonstrate_testing():
    """Run test demonstrations."""
    print("Test Configuration Demo")
    print("=" * 50)

    scenarios = ["basic", "error", "invalid"]
    for scenario in scenarios:
        await run_test_scenario(scenario)
        print("-" * 50)

Initial state: [{'id': 'test_1', 'tool_name': 'mock_tool', 'args': {'query': 'test query'}}]
Validations: {'has_messages': True, 'has_valid_tools': True, 'valid_results': True}
Results: {'test_1': 'Mock result: test query'}

## Errors: {}

if __name__ == "__main__":
    import nest_asyncio
    nest_asyncio.apply()
    asyncio.run(demonstrate_testing())