feat(langgraph-py): add reasoning infrastructure

Author: ranst91

python-sdk/ag_ui/integrations/langgraph/agent/agent.py

@@ -1,6 +1,6 @@
 import uuid
 import json
-from typing import Optional, List, Any, Union, AsyncGenerator
+from typing import Optional, List, Any, Union, AsyncGenerator, Generator
 
 from fastapi.responses import StreamingResponse
 
@@ -10,10 +10,52 @@
 from langchain_core.messages import HumanMessage
 from langgraph.types import Command
 
-from .types import State, LangGraphPlatformMessage, MessagesInProgressRecord, SchemaKeys, MessageInProgress, RunMetadata, LangGraphEventTypes, CustomEventNames
-from .utils import agui_messages_to_langchain, DEFAULT_SCHEMA_KEYS, filter_object_by_schema_keys, get_stream_payload_input, langchain_messages_to_agui
-
-from ag_ui.core import EventType, CustomEvent, MessagesSnapshotEvent, RawEvent, RunAgentInput, RunErrorEvent, RunFinishedEvent, RunStartedEvent, StateDeltaEvent, StateSnapshotEvent, StepFinishedEvent, StepStartedEvent, TextMessageContentEvent, TextMessageEndEvent, TextMessageStartEvent, ToolCallArgsEvent, ToolCallEndEvent, ToolCallStartEvent
+from .types import (
+    State,
+    LangGraphPlatformMessage,
+    MessagesInProgressRecord,
+    SchemaKeys,
+    MessageInProgress,
+    RunMetadata,
+    LangGraphEventTypes,
+    CustomEventNames,
+    LangGraphReasoning
+)
+from .utils import (
+    agui_messages_to_langchain,
+    DEFAULT_SCHEMA_KEYS,
+    filter_object_by_schema_keys,
+    get_stream_payload_input,
+    langchain_messages_to_agui,
+    resolve_reasoning_content,
+    resolve_message_content
+)
+
+from ag_ui.core import (
+    EventType,
+    CustomEvent,
+    MessagesSnapshotEvent,
+    RawEvent,
+    RunAgentInput,
+    RunErrorEvent,
+    RunFinishedEvent,
+    RunStartedEvent,
+    StateDeltaEvent,
+    StateSnapshotEvent,
+    StepFinishedEvent,
+    StepStartedEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    TextMessageStartEvent,
+    ToolCallArgsEvent,
+    ToolCallEndEvent,
+    ToolCallStartEvent,
+    ThinkingTextMessageStartEvent,
+    ThinkingTextMessageContentEvent,
+    ThinkingTextMessageEndEvent,
+    ThinkingStartEvent,
+    ThinkingEndEvent,
+)
 from ag_ui.encoder import EventEncoder
 
 ProcessedEvents = Union[
@@ -60,6 +102,7 @@ async def _handle_stream_events(self, input: RunAgentInput) -> AsyncGenerator[st
         self.active_run = {
             "id": input.run_id,
             "thread_id": thread_id,
+            "thinking_process": None,
         }
 
         messages = input.messages or []
@@ -390,10 +433,28 @@ async def _handle_single_event(self, event: Any, state: State) -> AsyncGenerator
             is_tool_call_args_event = has_current_stream and current_stream.get("tool_call_id") and tool_call_data and tool_call_data.get("args")
             is_tool_call_end_event = has_current_stream and current_stream.get("tool_call_id") and not tool_call_data
 
-            is_message_start_event = not has_current_stream and not tool_call_data
-            is_message_content_event = has_current_stream and not tool_call_data
+            reasoning_data = resolve_reasoning_content(event["data"]["chunk"]) if event["data"]["chunk"] else None
+            message_content = resolve_message_content(event["data"]["chunk"].content) if event["data"]["chunk"] and event["data"]["chunk"].content else None
+            is_message_content_event = tool_call_data is None and message_content
             is_message_end_event = has_current_stream and not current_stream.get("tool_call_id") and not is_message_content_event
 
+            if reasoning_data:
+                self.handle_thinking_event(reasoning_data)
+                return
+
+            if reasoning_data is None and self.active_run.get('thinking_process', None) is not None:
+                yield self._dispatch_event(
+                    ThinkingTextMessageEndEvent(
+                        type=EventType.THINKING_TEXT_MESSAGE_END,
+                    )
+                )
+                yield self._dispatch_event(
+                    ThinkingEndEvent(
+                        type=EventType.THINKING_END,
+                    )
+                )
+                self.active_run["thinking_process"] = None
+
             if tool_call_used_to_predict_state:
                 yield self._dispatch_event(
                     CustomEvent(
@@ -442,27 +503,35 @@ async def _handle_single_event(self, event: Any, state: State) -> AsyncGenerator
 
             if is_tool_call_args_event and should_emit_tool_calls:
                 yield self._dispatch_event(
-                    ToolCallArgsEvent(type=EventType.TOOL_CALL_ARGS, tool_call_id=current_stream["tool_call_id"], delta=tool_call_data["args"], raw_event=event)
+                    ToolCallArgsEvent(
+                        type=EventType.TOOL_CALL_ARGS,
+                        tool_call_id=current_stream["tool_call_id"],
+                        delta=tool_call_data["args"],
+                        raw_event=event
+                    )
                 )
                 return
 
-            if is_message_start_event and should_emit_messages:
-                resolved = self._dispatch_event(
-                    TextMessageStartEvent(
-                        type=EventType.TEXT_MESSAGE_START,
-                        role="assistant",
-                        message_id=event["data"]["chunk"].id,
-                        raw_event=event,
+            if is_message_content_event and should_emit_messages:
+                if bool(current_stream and current_stream.get("id")) == False:
+                    yield self._dispatch_event(
+                        TextMessageStartEvent(
+                            type=EventType.TEXT_MESSAGE_START,
+                            role="assistant",
+                            message_id=event["data"]["chunk"].id,
+                            raw_event=event,
+                        )
                     )
-                )
-                if resolved:
                     self.set_message_in_progress(
-                        self.active_run["id"], MessageInProgress(id=event["data"]["chunk"].id)
+                        self.active_run["id"],
+                        MessageInProgress(
+                            id=event["data"]["chunk"].id,
+                            tool_call_id=None,
+                            tool_call_name=None
+                        )
                     )
-                yield resolved
-                return
+                    current_stream = self.get_message_in_progress(self.active_run["id"])
 
-            if is_message_content_event and should_emit_messages:
                 yield self._dispatch_event(
                     TextMessageContentEvent(
                         type=EventType.TEXT_MESSAGE_CONTENT,
@@ -533,6 +602,55 @@ async def _handle_single_event(self, event: Any, state: State) -> AsyncGenerator
                 CustomEvent(type=EventType.CUSTOM, name=event["name"], value=event["data"], raw_event=event)
             )
 
+    def handle_thinking_event(self, reasoning_data: LangGraphReasoning) -> Generator[str, Any, str | None]:
+        if not reasoning_data or "type" not in reasoning_data or "text" not in reasoning_data:
+            return ""
+
+        thinking_step_index = reasoning_data.get("index")
+
+        if (self.active_run.get("thinking_process") and
+                self.active_run["thinking_process"].get("index") and
+                self.active_run["thinking_process"]["index"] != thinking_step_index):
+
+            if self.active_run["thinking_process"].get("type"):
+                yield self._dispatch_event(
+                    ThinkingTextMessageEndEvent(
+                        type=EventType.THINKING_TEXT_MESSAGE_END,
+                    )
+                )
+            yield self._dispatch_event(
+                ThinkingEndEvent(
+                    type=EventType.THINKING_END,
+                )
+            )
+            self.active_run["thinking_process"] = None
+
+        if not self.active_run.get("thinking_process"):
+            yield self._dispatch_event(
+                ThinkingStartEvent(
+                    type=EventType.THINKING_START,
+                )
+            )
+            self.active_run["thinking_process"] = {
+                "index": thinking_step_index
+            }
+
+        if self.active_run["thinking_process"].get("type") != reasoning_data["type"]:
+            yield self._dispatch_event(
+                ThinkingTextMessageStartEvent(
+                    type=EventType.THINKING_TEXT_MESSAGE_START,
+                )
+            )
+            self.active_run["thinking_process"]["type"] = reasoning_data["type"]
+
+        if self.active_run["thinking_process"].get("type"):
+            yield self._dispatch_event(
+                ThinkingTextMessageContentEvent(
+                    type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
+                    delta=reasoning_data["text"]
+                )
+            )
+
     async def get_checkpoint_before_message(self, message_id: str, thread_id: str):
         if not thread_id:
             raise ValueError("Missing thread_id in config")

python-sdk/ag_ui/integrations/langgraph/agent/types.py

@@ -1,4 +1,4 @@
-from typing import TypedDict, Optional, List, Any, Dict, Union
+from typing import TypedDict, Optional, List, Any, Dict, Union, Literal
 from typing_extensions import NotRequired
 from enum import Enum
 
@@ -28,6 +28,11 @@ class CustomEventNames(str, Enum):
     "config": NotRequired[Optional[List[str]]]
 })
 
+ThinkingProcess = TypedDict("ThinkingProcess", {
+    "index": int,
+    "type": NotRequired[Optional[Literal['text']]],
+})
+
 MessageInProgress = TypedDict("MessageInProgress", {
     "id": str,
     "tool_call_id": NotRequired[Optional[str]],
@@ -41,7 +46,8 @@ class CustomEventNames(str, Enum):
     "prev_node_name": NotRequired[Optional[str]],
     "exiting_node": NotRequired[bool],
     "manually_emitted_state": NotRequired[Optional[State]],
-    "thread_id": NotRequired[Optional[str]]
+    "thread_id": NotRequired[Optional[ThinkingProcess]],
+    "thinking_process": NotRequired[Optional[str]]
 })
 
 MessagesInProgressRecord = Dict[str, Optional[MessageInProgress]]
@@ -77,3 +83,9 @@ class LangGraphPlatformActionExecutionMessage(BaseLangGraphPlatformMessage):
     "state_key": str,
     "tool_argument": str
 })
+
+LangGraphReasoning = TypedDict("LangGraphReasoning", {
+    "type": str,
+    "text": str,
+    "index": int
+})

python-sdk/ag_ui/integrations/langgraph/agent/utils.py

@@ -11,7 +11,7 @@
     ToolCall as AGUIToolCall,
     FunctionCall as AGUIFunctionCall,
 )
-from .types import State, SchemaKeys
+from .types import State, SchemaKeys, LangGraphReasoning
 
 DEFAULT_SCHEMA_KEYS = ["tools"]
 
@@ -43,7 +43,7 @@ def langchain_messages_to_agui(messages: List[BaseMessage]) -> List[AGUIMessage]
             agui_messages.append(AGUIUserMessage(
                 id=str(message.id),
                 role="user",
-                content=stringify_if_needed(message.content),
+                content=stringify_if_needed(resolve_message_content(message.content)),
                 name=message.name,
             ))
         elif isinstance(message, AIMessage):
@@ -63,22 +63,22 @@ def langchain_messages_to_agui(messages: List[BaseMessage]) -> List[AGUIMessage]
             agui_messages.append(AGUIAssistantMessage(
                 id=str(message.id),
                 role="assistant",
-                content=stringify_if_needed(message.content),
+                content=stringify_if_needed(resolve_message_content(message.content)),
                 tool_calls=tool_calls,
                 name=message.name,
             ))
         elif isinstance(message, SystemMessage):
             agui_messages.append(AGUISystemMessage(
                 id=str(message.id),
                 role="system",
-                content=stringify_if_needed(message.content),
+                content=stringify_if_needed(resolve_message_content(message.content)),
                 name=message.name,
             ))
         elif isinstance(message, ToolMessage):
             agui_messages.append(AGUIToolMessage(
                 id=str(message.id),
                 role="tool",
-                content=stringify_if_needed(message.content),
+                content=stringify_if_needed(resolve_message_content(message.content)),
                 tool_call_id=message.tool_call_id,
             ))
         else:
@@ -125,4 +125,56 @@ def agui_messages_to_langchain(messages: List[AGUIMessage]) -> List[BaseMessage]
             ))
         else:
             raise ValueError(f"Unsupported message role: {role}")
-    return langchain_messages
+    return langchain_messages
+
+def resolve_reasoning_content(chunk: Any) -> LangGraphReasoning | None:
+    content = chunk.content
+    if not content:
+        return None
+
+    # Anthropic reasoning response
+    if isinstance(content, list) and content and content[0]:
+        if not content[0].get("thinking"):
+            return None
+        return LangGraphReasoning(
+            text=content[0]["thinking"],
+            type="text",
+            index=content[0].get("index", 0)
+        )
+
+    # OpenAI reasoning response
+    if hasattr(chunk, "additional_kwargs"):
+        reasoning = chunk.additional_kwargs.get("reasoning", {})
+        summary = reasoning.get("summary", [])
+        if summary:
+            data = summary[0]
+            if not data or not data.get("text"):
+                return None
+            return LangGraphReasoning(
+                type="text",
+                text=data["text"],
+                index=data.get("index", 0)
+            )
+
+    try:
+        parsed = json.loads(content)
+        return LangGraphReasoning(
+            type=parsed.get("type", "text"),
+            text=parsed.get("text", ""),
+            index=parsed.get("index", 0)
+        )
+    except json.JSONDecodeError:
+        return None
+
+def resolve_message_content(content: Any) -> str | None:
+    if not content:
+        return None
+
+    if isinstance(content, str):
+        return content
+
+    if isinstance(content, list) and content:
+        content_text = next((c.get("text") for c in content if isinstance(c, dict) and c.get("type") == "text"), None)
+        return content_text
+
+    return None