🎯 What You'll Learn Today

import asyncio
from typing import Annotated, Any, Literal, TypedDict
#!pip install langchain-core
#!pip install langgraph
from langchain_core.messages import AIMessage, BaseMessage, SystemMessage, ToolMessage
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.graph import END, START, StateGraph
from langgraph.graph.message import add_messages

class State(TypedDict):
    """State container for graph execution.

    Attributes:
        messages: Conversation history
        pending_tools: Tools awaiting execution
        results: Aggregated tool results
        errors: Aggregated error messages
        tool_configs: Tool-specific configurations
    """

    messages: Annotated[list[BaseMessage], add_messages]
    pending_tools: list[dict[str, Any]]
    results: Annotated[
        dict[str, Any],
        BinaryOperatorAggregate(dict[str, Any], lambda a, b: {**(b or {}), **a}),
    ]
    errors: Annotated[
        dict[str, str],
        BinaryOperatorAggregate(dict[str, str], lambda a, b: {**(b or {}), **a}),
    ]
    tool_configs: dict[str, dict[str, Any]]

async def parallel_executor(state: State) -> State:
    """Execute pending tools in parallel.

    Args:
        state: Current execution state

    Returns:
        Updated state with results/errors
    """
    if not state["pending_tools"]:
        return state

    tasks = []
    for tool_call in state["pending_tools"]:
        tasks.append(
            asyncio.sleep(0.1)  # Simulate tool execution
        )

    await asyncio.gather(*tasks)

    # Simulate results/errors
    new_results = {}
    new_errors = {}
    for tool in state["pending_tools"]:
        if "error" in tool["args"].get("query", ""):
            new_errors[tool["id"]] = "Simulated error"
        else:
            new_results[tool["id"]] = f"Result for {tool['args']['query']}"

    return {**state, "results": new_results, "errors": new_errors}

def result_aggregator(state: State) -> State:
    """Aggregate successful results.

    Args:
        state: Current state with results

    Returns:
        Updated state with aggregated messages
    """
    messages = list(state["messages"])

    for tool_id, result in state["results"].items():
        messages.append(
            ToolMessage(
                content=str(result), tool_call_id=tool_id, name=tool_id.split("_")[0]
            )
        )

    messages.append(
        AIMessage(content=f"Processed {len(state['results'])} results successfully.")
    )

    return {**state, "messages": messages}

def error_handler(state: State) -> State:
    """Process execution errors.

    Args:
        state: Current state with errors

    Returns:
        Updated state with error messages
    """
    messages = list(state["messages"])

    for tool_id, error in state["errors"].items():
        messages.append(AIMessage(content=f"Error in {tool_id}: {error}"))

    return {**state, "messages": messages}

def get_initial_state(state: State = None) -> State:
    """Initialize state with test tools.

    Args:
        state: Optional existing state

    Returns:
        Initialized state with test configuration
    """
    return {
        "messages": [SystemMessage(content="Starting execution")],
        "pending_tools": [
            {
                "id": f"tool_{i}",
                "tool_name": "test_tool",
                "args": {"query": f"test_query_{i}"},
            }
            for i in range(3)
        ],
        "results": {},
        "errors": {},
        "tool_configs": {},
    }

def route_results(state: State) -> Literal["result_aggregator", "error_handler"]:
    """Route based on execution results.

    Args:
        state: Current execution state

    Returns:
        Next node identifier
    """
    return "error_handler" if state.get("errors") else "result_aggregator"

def create_parallel_graph() -> StateGraph:
    """Create execution graph with parallel processing.

    Graph Structure:
    START → init → parallel_executor → [result_aggregator|error_handler] → END

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

    # Add nodes
    graph.add_node("init", get_initial_state)
    graph.add_node("parallel_executor", parallel_executor)
    graph.add_node("result_aggregator", result_aggregator)
    graph.add_node("error_handler", error_handler)

    # Configure edges
    graph.add_edge(START, "init")
    graph.add_edge("init", "parallel_executor")

    # Add conditional routing
    graph.add_conditional_edges(
        "parallel_executor",
        route_results,
        {"result_aggregator": "result_aggregator", "error_handler": "error_handler"},
    )

    # Connect to END
    graph.add_edge("result_aggregator", END)
    graph.add_edge("error_handler", END)

    return graph

async def demonstrate_graph():
    """Demonstrate graph execution."""
    print("Graph Execution Demo")
    print("=" * 50)

    graph = create_parallel_graph()
    chain = graph.compile()

    # Initialize state properly
    initial_state = {
        "messages": [SystemMessage(content="Starting execution")],
        "pending_tools": [],
        "results": {},
        "errors": {},
        "tool_configs": {},
    }

    result = await chain.ainvoke(initial_state)

    print("\nExecution Messages:")
    for msg in result["messages"]:
        prefix = type(msg).__name__.replace("Message", "")
        print(f"\n{prefix}: {msg.content}")

    print("\nResults:", result["results"])
    print("Errors:", result["errors"])

## Execution Messages

System: Starting execution
Tool: Result for test_query_0
Tool: Result for test_query_1
Tool: Result for test_query_2
AI: Processed 3 results successfully
Results: {'tool_0': 'Result for test_query_0', 'tool_1': 'Result for test_query_1', 'tool_2': 'Result for test_query_2'}
Errors: {}

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