from typing import Annotated, TypedDict
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.tools import tool
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode


class State(TypedDict):
    messages: Annotated[list, add_messages]


# Create two simple subgraphs
def create_subgraph(name: str):
    def node(state: State):
        return {"messages": [AIMessage(content=f"Response from {name}")]}
    
    builder = StateGraph(State)
    builder.add_node("respond", node)
    builder.add_edge(START, "respond")
    builder.add_edge("respond", END)
    return builder.compile()


subgraph_a = create_subgraph("subgraph_a")
subgraph_b = create_subgraph("subgraph_b")


@tool
def call_subgraph_a(query: str) -> str:
    """Call subgraph A"""
    result = subgraph_a.invoke({"messages": [HumanMessage(content=query)]})
    return result["messages"][-1].content


@tool
def call_subgraph_b(query: str) -> str:
    """Call subgraph B"""
    result = subgraph_b.invoke({"messages": [HumanMessage(content=query)]})
    return result["messages"][-1].content


def model_node(state: State):
    # Model returns two parallel tool calls
    return {
        "messages": [
            AIMessage(
                content="",
                tool_calls=[
                    {"id": "tool_call_A", "name": "call_subgraph_a", "args": {"query": "hello"}},
                    {"id": "tool_call_B", "name": "call_subgraph_b", "args": {"query": "world"}},
                ],
            )
        ]
    }


builder = StateGraph(State)
builder.add_node("model", model_node)
builder.add_node("tools", ToolNode([call_subgraph_a, call_subgraph_b]))
builder.add_edge(START, "model")
builder.add_edge("model", "tools")
builder.add_edge("tools", END)

graph = builder.compile()


async def main():
    async for chunk in graph.astream(
        {"messages": [HumanMessage(content="Call both subgraphs")]},
        stream_mode=["updates", "messages"],
        subgraphs=True,
    ):
        namespace, event_type, payload = chunk

        # Try to extract tool_call_id from payload
        tool_call_id = "?"
        if event_type == "messages" and isinstance(payload, tuple):
            msg = payload[0]
            if hasattr(msg, 'tool_call_id'):
                tool_call_id = msg.tool_call_id

        print(f"namespace: {namespace}, tool_call_id: {tool_call_id}")


import asyncio
asyncio.run(main())

----------
Output:

namespace: ()
event_type: messages
tool_call_id: ?
---
namespace: ()
event_type: updates
tool_call_id: ?
---
namespace: ('tools:1032c8d0-64e7-7979-524c-31671560f76a',)
event_type: messages
tool_call_id: ?
---
namespace: ('tools:1032c8d0-64e7-7979-524c-31671560f76a', '1')
event_type: messages
tool_call_id: ?
---
namespace: ('tools:1032c8d0-64e7-7979-524c-31671560f76a',)
event_type: updates
tool_call_id: ?
---
namespace: ('tools:1032c8d0-64e7-7979-524c-31671560f76a', '1')
event_type: updates
tool_call_id: ?
---
namespace: ()
event_type: messages
tool_call_id: tool_call_A
---
namespace: ()
event_type: messages
tool_call_id: tool_call_B
---
namespace: ()
event_type: updates
tool_call_id: ?
---