🎯 What You'll Learn Today

from typing import Annotated, Any, TypedDict
!pip install langchain-core
!pip install langgraph
from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
from langgraph.channels.binop import BinaryOperatorAggregate
from langgraph.graph.message import add_messages

def dict_reducer(a: dict, b: dict | None) -> dict:
    """Combine two dictionaries for result aggregation.

    Args:
        a: First dictionary
        b: Second dictionary (may be None)

    Returns:
        Combined dictionary with values from both inputs

    Examples:
        >>> dict_reducer({'a': 1}, {'b': 2})
        {'a': 1, 'b': 2}
        >>> dict_reducer({'a': 1}, None)
        {'a': 1}
    """
    if b is None:
        return a
    return {**a, **b}

class State(TypedDict):
    """State with aggregation support.

    Attributes:
        messages: Conversation history
        pending_tools: List of tools to execute
        results: Tool execution results with aggregation
        errors: Error messages with aggregation
    """

    messages: Annotated[list[BaseMessage], add_messages]
    pending_tools: list[dict]
    results: Annotated[
        dict[str, Any], BinaryOperatorAggregate(dict[str, Any], dict_reducer)
    ]
    errors: Annotated[
        dict[str, str], BinaryOperatorAggregate(dict[str, str], dict_reducer)
    ]

def result_aggregator(state: State) -> State:
    """Aggregate results from parallel tool execution.

    Creates appropriate message types for tool results and assistant responses.

    Args:
        state: Current state with results to aggregate

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

    # Process successful results
    if state["results"]:
        # First add tool messages for each result
        for tool_id, result in state["results"].items():
            tool_message = ToolMessage(
                content=str(result),
                tool_call_id=tool_id,
                name=tool_id.split("_")[0],  # Extract tool name from ID
            )
            messages.append(tool_message)

        # Add assistant summary of results
        summary = "Here are the tool results:\n" + "\n".join(
            f"- {tool_id}: {result}" for tool_id, result in state["results"].items()
        )
        messages.append(AIMessage(content=summary))

    # Process errors if any
    if state["errors"]:
        error_summary = "Some tools encountered errors:\n" + "\n".join(
            f"- {tool_id}: {error}" for tool_id, error in state["errors"].items()
        )
        messages.append(AIMessage(content=error_summary))

    return {
        "messages": messages,
        "pending_tools": [],
        "results": state["results"],
        "errors": state["errors"],
    }

async def demonstrate_aggregation():
    """Demonstrate result aggregation with proper message types."""
    # Initial state with some results and errors
    state = {
        "messages": [],
        "pending_tools": [],
        "results": {"calculator_1": "2 + 2 = 4", "weather_1": "22°C, Sunny"},
        "errors": {"search_1": "API rate limit exceeded"},
    }

    print("Result Aggregation Demo")
    print("=" * 50)

    result = result_aggregator(state)

    print("\nFinal Messages:")
    for msg in result["messages"]:
        if isinstance(msg, ToolMessage):
            print(f"\nTool ({msg.name}): {msg.content}")
        elif isinstance(msg, AIMessage):
            print(f"\nAssistant: {msg.content}")

Tool (calculator): 2 + 2 = 4
Tool (weather): 22°C, Sunny

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