feat: add subgraphs support in langgraph integrations

Author: ranst91

typescript-sdk/integrations/langgraph/python/ag_ui_langgraph/agent.py

@@ -155,6 +155,11 @@ async def _handle_stream_events(self, input: RunAgentInput) -> AsyncGenerator[st
         current_graph_state = state
 
         async for event in stream:
+            subgraphs_stream_enabled = input.forwarded_props.get('stream_subgraphs') if input.forwarded_props else False
+            is_subgraph_stream = (subgraphs_stream_enabled and (
+                event.get("event", "").startswith("events") or 
+                event.get("event", "").startswith("values")
+            ))
             if event["event"] == "error":
                 yield self._dispatch_event(
                     RunErrorEvent(type=EventType.RUN_ERROR, message=event["data"]["message"], raw_event=event)
@@ -323,8 +328,18 @@ async def prepare_stream(self, input: RunAgentInput, agent_state: State, config:
             )
             stream_input = {**forwarded_props, **payload_input} if payload_input else None
 
+
+        subgraphs_stream_enabled = input.forwarded_props.get('stream_subgraphs') if input.forwarded_props else False
+
+        stream = self.graph.astream_events(
+            stream_input,
+            config=config,
+            subgraps=bool(subgraphs_stream_enabled),
+            version="v2"
+        )
+
         return {
-            "stream": self.graph.astream_events(stream_input, config, version="v2"),
+            "stream": stream,
             "state": state,
             "config": config
         }
@@ -349,7 +364,13 @@ async def prepare_regenerate_stream( # pylint: disable=too-many-arguments
         )
 
         stream_input = self.langgraph_default_merge_state(time_travel_checkpoint.values, [message_checkpoint], tools)
-        stream = self.graph.astream_events(stream_input, fork, version="v2")
+        subgraphs_stream_enabled = input.forwarded_props.get('stream_subgraphs') if input.forwarded_props else False
+        stream = self.graph.astream_events(
+            stream_input,
+            fork,
+            subgraps=bool(subgraphs_stream_enabled),
+            version="v2"
+        )
 
         return {
             "stream": stream,

typescript-sdk/integrations/langgraph/src/agent.ts

@@ -194,13 +194,15 @@ export class LangGraphAgent extends AbstractAgent {
       }
     );
 
+      const payload = {
+          ...(input.forwardedProps ?? {}),
+          input: this.langGraphDefaultMergeState(timeTravelCheckpoint.values, [messageCheckpoint], tools),
+          // @ts-ignore
+          checkpointId: fork.checkpoint.checkpoint_id!,
+          streamMode,
+      };
     return {
-      streamResponse: this.client.runs.stream(threadId, this.assistant.assistant_id, {
-        input: this.langGraphDefaultMergeState(timeTravelCheckpoint.values, [messageCheckpoint], tools),
-        // @ts-ignore
-        checkpointId: fork.checkpoint.checkpoint_id!,
-        streamMode,
-      }),
+      streamResponse: this.client.runs.stream(threadId, this.assistant.assistant_id, payload),
       state: timeTravelCheckpoint as ThreadState<State>,
       streamMode,
     };
@@ -360,8 +362,14 @@ export class LangGraphAgent extends AbstractAgent {
       }
 
       for await (let streamResponseChunk of streamResponse) {
+          const subgraphsStreamEnabled = input.forwardedProps?.streamSubgraphs
+          const isSubgraphStream = (subgraphsStreamEnabled && (
+              streamResponseChunk.event.startsWith("events") ||
+              streamResponseChunk.event.startsWith("values")
+          ))
+
         // @ts-ignore
-        if (!streamMode.includes(streamResponseChunk.event as StreamMode)) {
+        if (!streamMode.includes(streamResponseChunk.event as StreamMode) && !isSubgraphStream) {
           continue;
         }
 
@@ -391,6 +399,12 @@ export class LangGraphAgent extends AbstractAgent {
         if (streamResponseChunk.event === "values") {
           latestStateValues = chunk.data;
           continue;
+        } else if (subgraphsStreamEnabled && chunk.event.startsWith("values|")) {
+            latestStateValues = {
+                ...latestStateValues,
+                ...chunk.data,
+            };
+            continue;
         }
 
         const chunkData = chunk.data;