diff --git a/integrations/langgraph/python/ag_ui_langgraph/agent.py b/integrations/langgraph/python/ag_ui_langgraph/agent.py index a67deae3b..5751af75f 100644 --- a/integrations/langgraph/python/ag_ui_langgraph/agent.py +++ b/integrations/langgraph/python/ag_ui_langgraph/agent.py @@ -6,10 +6,10 @@ from langgraph.graph.state import CompiledStateGraph try: - from langchain.schema import BaseMessage, SystemMessage + from langchain.schema import BaseMessage, SystemMessage, ToolMessage except ImportError: # Langchain >= 1.0.0 - from langchain_core.messages import BaseMessage, SystemMessage + from langchain_core.messages import BaseMessage, SystemMessage, ToolMessage from langchain_core.runnables import RunnableConfig, ensure_config from langchain_core.messages import HumanMessage @@ -682,6 +682,51 @@ async def _handle_single_event(self, event: Any, state: State) -> AsyncGenerator elif event_type == LangGraphEventTypes.OnToolEnd: tool_call_output = event["data"]["output"] + + if isinstance(tool_call_output, Command): + # Extract ToolMessages from Command.update + messages = tool_call_output.update.get('messages', []) + tool_messages = [m for m in messages if isinstance(m, ToolMessage)] + + # Process each tool message + for tool_msg in tool_messages: + if not self.active_run["has_function_streaming"]: + yield self._dispatch_event( + ToolCallStartEvent( + type=EventType.TOOL_CALL_START, + tool_call_id=tool_msg.tool_call_id, + tool_call_name=tool_msg.name, + parent_message_id=tool_msg.id, + raw_event=event, + ) + ) + yield self._dispatch_event( + ToolCallArgsEvent( + type=EventType.TOOL_CALL_ARGS, + tool_call_id=tool_msg.tool_call_id, + delta=json.dumps(event["data"].get("input", {})), + raw_event=event + ) + ) + yield self._dispatch_event( + ToolCallEndEvent( + type=EventType.TOOL_CALL_END, + tool_call_id=tool_msg.tool_call_id, + raw_event=event + ) + ) + + yield self._dispatch_event( + ToolCallResultEvent( + type=EventType.TOOL_CALL_RESULT, + tool_call_id=tool_msg.tool_call_id, + message_id=str(uuid.uuid4()), + content=tool_msg.content, + role="tool" + ) + ) + return + if not self.active_run["has_function_streaming"]: yield self._dispatch_event( ToolCallStartEvent(