Merge pull request #102 from ag-ui-protocol/feat/langgraph-reasoning-support

Reasoning/thinking support

Author: ranst91

python-sdk/ag_ui/core/__init__.py

@@ -51,10 +51,15 @@
     "TextMessageContentEvent",
     "TextMessageEndEvent",
     "TextMessageChunkEvent",
+    "ThinkingTextMessageStartEvent",
+    "ThinkingTextMessageContentEvent",
+    "ThinkingTextMessageEndEvent",
     "ToolCallStartEvent",
     "ToolCallArgsEvent",
     "ToolCallEndEvent",
     "ToolCallChunkEvent",
+    "ThinkingStartEvent",
+    "ThinkingEndEvent",
     "StateSnapshotEvent",
     "StateDeltaEvent",
     "MessagesSnapshotEvent",

python-sdk/ag_ui/core/events.py

@@ -17,10 +17,15 @@ class EventType(str, Enum):
     TEXT_MESSAGE_CONTENT = "TEXT_MESSAGE_CONTENT"
     TEXT_MESSAGE_END = "TEXT_MESSAGE_END"
     TEXT_MESSAGE_CHUNK = "TEXT_MESSAGE_CHUNK"
+    THINKING_TEXT_MESSAGE_START = "THINKING_TEXT_MESSAGE_START",
+    THINKING_TEXT_MESSAGE_CONTENT = "THINKING_TEXT_MESSAGE_CONTENT",
+    THINKING_TEXT_MESSAGE_END = "THINKING_TEXT_MESSAGE_END",
     TOOL_CALL_START = "TOOL_CALL_START"
     TOOL_CALL_ARGS = "TOOL_CALL_ARGS"
     TOOL_CALL_END = "TOOL_CALL_END"
     TOOL_CALL_CHUNK = "TOOL_CALL_CHUNK"
+    THINKING_START = "THINKING_START",
+    THINKING_END = "THINKING_END",
     STATE_SNAPSHOT = "STATE_SNAPSHOT"
     STATE_DELTA = "STATE_DELTA"
     MESSAGES_SNAPSHOT = "MESSAGES_SNAPSHOT"
@@ -80,6 +85,29 @@ class TextMessageChunkEvent(BaseEvent):
     role: Optional[Literal["assistant"]] = None
     delta: Optional[str] = None
 
+class ThinkingTextMessageStartEvent(BaseEvent):
+    """
+    Event indicating the start of a thinking text message.
+    """
+    type: Literal[EventType.THINKING_TEXT_MESSAGE_START]
+
+class ThinkingTextMessageContentEvent(BaseEvent):
+    """
+    Event indicating a piece of a thinking text message.
+    """
+    type: Literal[EventType.THINKING_TEXT_MESSAGE_CONTENT]
+    delta: str  # This should not be an empty string
+
+    def model_post_init(self, __context):
+        if len(self.delta) == 0:
+            raise ValueError("Delta must not be an empty string")
+
+class ThinkingTextMessageEndEvent(BaseEvent):
+    """
+    Event indicating the end of a thinking text message.
+    """
+    type: Literal[EventType.THINKING_TEXT_MESSAGE_END]
+
 class ToolCallStartEvent(BaseEvent):
     """
     Event indicating the start of a tool call.
@@ -116,6 +144,19 @@ class ToolCallChunkEvent(BaseEvent):
     parent_message_id: Optional[str] = None
     delta: Optional[str] = None
 
+class ThinkingStartEvent(BaseEvent):
+    """
+    Event indicating the start of a thinking step event.
+    """
+    type: Literal[EventType.THINKING_START]
+    title: Optional[str] = 'Thinking...'
+
+class ThinkingEndEvent(BaseEvent):
+    """
+    Event indicating the end of a thinking step event.
+    """
+    type: Literal[EventType.THINKING_END]
+
 class StateSnapshotEvent(BaseEvent):
     """
     Event containing a snapshot of the state.

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

@@ -18,10 +18,12 @@ import {
   LangGraphEventTypes,
   State,
   MessagesInProgressRecord,
+  ThinkingInProgress,
   SchemaKeys,
   MessageInProgress,
   RunMetadata,
   PredictStateTool,
+  LangGraphReasoning
 } from "./types";
 import {
   AbstractAgent,
@@ -44,6 +46,11 @@ import {
   ToolCallArgsEvent,
   ToolCallEndEvent,
   ToolCallStartEvent,
+  ThinkingTextMessageStartEvent,
+  ThinkingTextMessageContentEvent,
+  ThinkingTextMessageEndEvent,
+  ThinkingStartEvent,
+  ThinkingEndEvent,
 } from "@ag-ui/client";
 import { RunsStreamPayload } from "@langchain/langgraph-sdk/dist/types";
 import {
@@ -52,15 +59,22 @@ import {
   filterObjectBySchemaKeys,
   getStreamPayloadInput,
   langchainMessagesToAgui,
+  resolveMessageContent,
+  resolveReasoningContent
 } from "@/utils";
 
 export type ProcessedEvents =
   | TextMessageStartEvent
   | TextMessageContentEvent
   | TextMessageEndEvent
+  | ThinkingTextMessageStartEvent
+  | ThinkingTextMessageContentEvent
+  | ThinkingTextMessageEndEvent
   | ToolCallStartEvent
   | ToolCallArgsEvent
   | ToolCallEndEvent
+  | ThinkingStartEvent
+  | ThinkingEndEvent
   | StateSnapshotEvent
   | StateDeltaEvent
   | MessagesSnapshotEvent
@@ -98,6 +112,7 @@ export class LangGraphAgent extends AbstractAgent {
   graphId: string;
   assistant?: Assistant;
   messagesInProcess: MessagesInProgressRecord;
+  thinkingProcess: null | ThinkingInProgress;
   activeRun?: RunMetadata;
   // @ts-expect-error no need to initialize subscriber right now
   subscriber: Subscriber<ProcessedEvents>;
@@ -108,6 +123,7 @@ export class LangGraphAgent extends AbstractAgent {
     this.agentName = config.agentName;
     this.graphId = config.graphId;
     this.assistantConfig = config.assistantConfig;
+    this.thinkingProcess = null
     this.client =
       config?.client ??
       new LangGraphClient({
@@ -389,7 +405,7 @@ export class LangGraphAgent extends AbstractAgent {
         let shouldEmitToolCalls = event.metadata["emit-tool-calls"] ?? true;
 
         if (event.data.chunk.response_metadata.finish_reason) return;
-        const currentStream = this.getMessageInProgress(this.activeRun!.id);
+        let currentStream = this.getMessageInProgress(this.activeRun!.id);
         const hasCurrentStream = Boolean(currentStream?.id);
         const toolCallData = event.data.chunk.tool_call_chunks?.[0];
         const toolCallUsedToPredictState = event.metadata["predict_state"]?.some(
@@ -401,11 +417,28 @@ export class LangGraphAgent extends AbstractAgent {
           hasCurrentStream && currentStream?.toolCallId && toolCallData.args;
         const isToolCallEndEvent = hasCurrentStream && currentStream?.toolCallId && !toolCallData;
 
-        const isMessageStartEvent = !hasCurrentStream && !toolCallData;
-        const isMessageContentEvent = hasCurrentStream && !toolCallData;
+        const reasoningData = resolveReasoningContent(event.data);
+        const messageContent = resolveMessageContent(event.data.chunk.content);
+        const isMessageContentEvent = Boolean(!toolCallData && messageContent);
+
         const isMessageEndEvent =
           hasCurrentStream && !currentStream?.toolCallId && !isMessageContentEvent;
 
+        if (reasoningData) {
+          this.handleThinkingEvent(reasoningData)
+          break;
+        }
+
+        if (!reasoningData && this.thinkingProcess) {
+          this.dispatchEvent({
+            type: EventType.THINKING_TEXT_MESSAGE_END,
+          })
+          this.dispatchEvent({
+            type: EventType.THINKING_END,
+          })
+          this.thinkingProcess = null;
+        }
+
         if (toolCallUsedToPredictState) {
           this.dispatchEvent({
             type: EventType.CUSTOM,
@@ -417,7 +450,7 @@ export class LangGraphAgent extends AbstractAgent {
         if (isToolCallEndEvent) {
           const resolved = this.dispatchEvent({
             type: EventType.TOOL_CALL_END,
-            toolCallId: currentStream.toolCallId!,
+            toolCallId: currentStream?.toolCallId!,
             rawEvent: event,
           });
           if (resolved) {
@@ -460,36 +493,35 @@ export class LangGraphAgent extends AbstractAgent {
         if (isToolCallArgsEvent && shouldEmitToolCalls) {
           this.dispatchEvent({
             type: EventType.TOOL_CALL_ARGS,
-            toolCallId: currentStream.toolCallId!,
+            toolCallId: currentStream?.toolCallId!,
             delta: toolCallData.args,
             rawEvent: event,
           });
           break;
         }
 
-        // Message started: emit TextMessageStart
-        if (isMessageStartEvent && shouldEmitMessages) {
-          const resolved = this.dispatchEvent({
-            type: EventType.TEXT_MESSAGE_START,
-            role: "assistant",
-            messageId: event.data.chunk.id,
-            rawEvent: event,
-          });
-          if (resolved) {
+        // Message content: emit TextMessageContent
+        if (isMessageContentEvent && shouldEmitMessages) {
+          // No existing message yet, also init the message
+          if (!currentStream) {
+            this.dispatchEvent({
+              type: EventType.TEXT_MESSAGE_START,
+              role: "assistant",
+              messageId: event.data.chunk.id,
+              rawEvent: event,
+            });
             this.setMessageInProgress(this.activeRun!.id, {
               id: event.data.chunk.id,
               toolCallId: null,
               toolCallName: null,
             });
+            currentStream = this.getMessageInProgress(this.activeRun!.id);
           }
-          break;
-        }
-        // Message content: emit TextMessageContent
-        if (isMessageContentEvent && shouldEmitMessages) {
+
           this.dispatchEvent({
             type: EventType.TEXT_MESSAGE_CONTENT,
             messageId: currentStream!.id,
-            delta: event.data.chunk.content,
+            delta: messageContent!,
             rawEvent: event,
           });
           break;
@@ -583,6 +615,51 @@ export class LangGraphAgent extends AbstractAgent {
     }
   }
 
+  handleThinkingEvent(reasoningData: LangGraphReasoning) {
+    if (!reasoningData || !reasoningData.type || !reasoningData.text) {
+      return;
+    }
+
+    const thinkingStepIndex = reasoningData.index;
+
+    if (this.thinkingProcess?.index && this.thinkingProcess.index !== thinkingStepIndex) {
+      if (this.thinkingProcess.type) {
+        this.dispatchEvent({
+          type: EventType.THINKING_TEXT_MESSAGE_END,
+        })
+      }
+      this.dispatchEvent({
+        type: EventType.THINKING_END,
+      })
+      this.thinkingProcess = null;
+    }
+
+    if (!this.thinkingProcess) {
+      // No thinking step yet. Start a new one
+      this.dispatchEvent({
+        type: EventType.THINKING_START,
+      })
+      this.thinkingProcess = {
+        index: thinkingStepIndex,
+      };
+    }
+
+
+    if (this.thinkingProcess.type !== reasoningData.type) {
+      this.dispatchEvent({
+        type: EventType.THINKING_TEXT_MESSAGE_START,
+      })
+      this.thinkingProcess.type = reasoningData.type
+    }
+
+    if (this.thinkingProcess.type) {
+      this.dispatchEvent({
+        type: EventType.THINKING_TEXT_MESSAGE_CONTENT,
+        delta: reasoningData.text
+      })
+    }
+  }
+
   getStateSnapshot(state: State) {
     const schemaKeys = this.activeRun!.schemaKeys!;
     // Do not emit state keys that are not part of the output schema

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

@@ -28,6 +28,11 @@ export type MessageInProgress = {
   toolCallName?: string | null;
 };
 
+export type ThinkingInProgress = {
+  index: number;
+  type?: LangGraphReasoning['type'];
+}
+
 export interface RunMetadata {
   id: string;
   schemaKeys?: SchemaKeys;
@@ -89,3 +94,9 @@ export interface PredictStateTool {
   state_key: string;
   tool_argument: string;
 }
+
+export interface LangGraphReasoning {
+  type: 'text';
+  text: string;
+  index: number
+}