feat(langgraph): Add UI messages API (#4157)

Sample usage:

```python
from typing import Annotated, Sequence, TypedDict

from langchain_core.messages import BaseMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages

from langgraph.graph.ui import AnyUIMessage, ui_message_reducer, push_ui_message


class AgentState(TypedDict):  # noqa: D101
    messages: Annotated[Sequence[BaseMessage], add_messages]
    ui: Annotated[Sequence[AnyUIMessage], ui_message_reducer]


async def agent(state: AgentState):  # noqa: D103
    message = await ChatOpenAI(model="gpt-4o-mini", temperature=0).ainvoke(
        state["messages"]
    )

    # Also directly writes the message to `ui`
    push_ui_message("simple", {"hello": "123"}, message=message, state_key="ui")

    return {"messages": [message]}

```
This commit is contained in:
David Duong
2025-04-03 22:10:18 +02:00
committed by GitHub
2 changed files with 315 additions and 55 deletions
+109 -55
View File
@@ -12,10 +12,6 @@ Generative user interfaces (Generative UI) allows agents to go beyond text and g
LangGraph Platform supports colocating your React components with your graph code. This allows you to focus on building specific UI components for your graph while easily plugging into existing chat interfaces such as [Agent Chat](https://agentchat.vercel.app) and loading the code only when actually needed.
!!! warning "LangGraph.js only"
Currently only LangGraph.js supports Generative UI. Support for Python is coming soon.
## Tutorial
### 1. Define and configure UI components
@@ -74,58 +70,105 @@ CSS and Tailwind 4.x is also supported out of the box, so you can freely use Tai
### 2. Send the UI components in your graph
Use the `typedUi` utility to emit UI elements from your agent nodes:
=== "Python"
```typescript title="src/agent/index.ts"
import {
typedUi,
uiMessageReducer,
} from "@langchain/langgraph-sdk/react-ui/server";
```python title="src/agent.py"
import uuid
from typing import Annotated, Sequence, TypedDict
import { ChatOpenAI } from "@langchain/openai";
import { v4 as uuidv4 } from "uuid";
import { z } from "zod";
from langchain_core.messages import AIMessage, BaseMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langgraph.graph.ui import AnyUIMessage, ui_message_reducer, push_ui_message
import type ComponentMap from "./ui.js";
import {
Annotation,
MessagesAnnotation,
StateGraph,
type LangGraphRunnableConfig,
} from "@langchain/langgraph";
class AgentState(TypedDict): # noqa: D101
messages: Annotated[Sequence[BaseMessage], add_messages]
ui: Annotated[Sequence[AnyUIMessage], ui_message_reducer]
const AgentState = Annotation.Root({
...MessagesAnnotation.spec,
ui: Annotation({ reducer: uiMessageReducer, default: () => [] }),
});
export const graph = new StateGraph(AgentState)
.addNode("weather", async (state, config) => {
// Provide the type of the component map to ensure
// type safety of `ui.push()` calls as well as
// pushing the messages to the `ui` and sending a custom event as well.
const ui = typedUi<typeof ComponentMap>(config);
async def weather(state: AgentState):
class WeatherOutput(TypedDict):
city: str
const weather = await new ChatOpenAI({ model: "gpt-4o-mini" })
.withStructuredOutput(z.object({ city: z.string() }))
.withConfig({ tags: ["langsmith:nostream"] })
.invoke(state.messages);
weather: WeatherOutput = (
await ChatOpenAI(model="gpt-4o-mini")
.with_structured_output(WeatherOutput)
.with_config({"tags": ["nostream"]})
.ainvoke(state["messages"])
)
const response = {
id: uuidv4(),
type: "ai",
content: `Here's the weather for ${weather.city}`,
};
message = AIMessage(
id=str(uuid.uuid4()),
content=f"Here's the weather for {weather['city']}",
)
// Emit UI elements with associated AI message
ui.push({ name: "weather", props: weather }, { message: response });
# Emit UI elements associated with the message
push_ui_message("weather", weather, message=message)
return {"messages": [message]}
return { messages: [response] };
})
.addEdge("__start__", "weather")
.compile();
```
workflow = StateGraph(AgentState)
workflow.add_node(weather)
workflow.add_edge("__start__", "weather")
graph = workflow.compile()
```
=== "JS"
Use the `typedUi` utility to emit UI elements from your agent nodes:
```typescript title="src/agent/index.ts"
import {
typedUi,
uiMessageReducer,
} from "@langchain/langgraph-sdk/react-ui/server";
import { ChatOpenAI } from "@langchain/openai";
import { v4 as uuidv4 } from "uuid";
import { z } from "zod";
import type ComponentMap from "./ui.js";
import {
Annotation,
MessagesAnnotation,
StateGraph,
type LangGraphRunnableConfig,
} from "@langchain/langgraph";
const AgentState = Annotation.Root({
...MessagesAnnotation.spec,
ui: Annotation({ reducer: uiMessageReducer, default: () => [] }),
});
export const graph = new StateGraph(AgentState)
.addNode("weather", async (state, config) => {
// Provide the type of the component map to ensure
// type safety of `ui.push()` calls as well as
// pushing the messages to the `ui` and sending a custom event as well.
const ui = typedUi<typeof ComponentMap>(config);
const weather = await new ChatOpenAI({ model: "gpt-4o-mini" })
.withStructuredOutput(z.object({ city: z.string() }))
.withConfig({ tags: ["nostream"] })
.invoke(state.messages);
const response = {
id: uuidv4(),
type: "ai",
content: `Here's the weather for ${weather.city}`,
};
// Emit UI elements associated with the AI message
ui.push({ name: "weather", props: weather }, { message: response });
return { messages: [response] };
})
.addEdge("__start__", "weather")
.compile();
```
### 3. Handle UI elements in your React application
@@ -294,18 +337,29 @@ const { thread, submit } = useStream({
### Remove UI messages from state
Similar to how messages can be removed from the state by appending a RemoveMessage you can remove an UI message from the state by calling `ui.delete` with the ID of the UI message.
Similar to how messages can be removed from the state by appending a RemoveMessage you can remove an UI message from the state by calling `remove_ui_message` / `ui.delete` with the ID of the UI message.
```tsx
// pushed message
const message = ui.push({ name: "weather", props: { city: "London" } });
=== "Python"
// remove said message
ui.delete(message.id);
```python
from langgraph.graph.ui import push_ui_message, delete_ui_message
// return new state to persist changes
return { ui: ui.items };
```
# push message
message = push_ui_message("weather", {"city": "London"})
# remove said message
delete_ui_message(message["id"])
```
=== "JS"
```tsx
// push message
const message = ui.push({ name: "weather", props: { city: "London" } });
// remove said message
ui.delete(message.id);
```
## Learn more
+206
View File
@@ -0,0 +1,206 @@
from typing import Any, Literal, Optional, Union
from uuid import uuid4
from langchain_core.messages import AnyMessage
from typing_extensions import TypedDict
from langgraph.constants import CONF, CONFIG_KEY_SEND
from langgraph.utils.config import get_config, get_stream_writer
class UIMessage(TypedDict):
"""A message type for UI updates in LangGraph.
This TypedDict represents a UI message that can be sent to update the UI state.
It contains information about the UI component to render and its properties.
Attributes:
type: Literal type indicating this is a UI message.
id: Unique identifier for the UI message.
name: Name of the UI component to render.
props: Properties to pass to the UI component.
metadata: Additional metadata about the UI message.
"""
type: Literal["ui"]
id: str
name: str
props: dict[str, Any]
metadata: dict[str, Any]
class RemoveUIMessage(TypedDict):
"""A message type for removing UI components in LangGraph.
This TypedDict represents a message that can be sent to remove a UI component
from the current state.
Attributes:
type: Literal type indicating this is a remove-ui message.
id: Unique identifier of the UI message to remove.
"""
type: Literal["remove-ui"]
id: str
AnyUIMessage = Union[UIMessage, RemoveUIMessage]
def push_ui_message(
name: str,
props: dict[str, Any],
*,
id: Optional[str] = None,
metadata: Optional[dict[str, Any]] = None,
message: Optional[AnyMessage] = None,
state_key: str = "ui",
) -> UIMessage:
"""Push a new UI message to update the UI state.
This function creates and sends a UI message that will be rendered in the UI.
It also updates the graph state with the new UI message.
Args:
name: Name of the UI component to render.
props: Properties to pass to the UI component.
id: Optional unique identifier for the UI message.
If not provided, a random UUID will be generated.
metadata: Optional additional metadata about the UI message.
message: Optional message object to associate with the UI message.
state_key: Key in the graph state where the UI messages are stored.
Defaults to "ui".
Returns:
The created UI message.
Example:
.. code-block:: python
push_ui_message(
name="component-name",
props={"content": "Hello world"},
)
"""
writer = get_stream_writer()
config = get_config()
message_id = None
if message:
if isinstance(message, dict) and "id" in message:
message_id = message.get("id")
elif hasattr(message, "id"):
message_id = message.id
evt: UIMessage = {
"type": "ui",
"id": id or str(uuid4()),
"name": name,
"props": props,
"metadata": {
**(config.get("metadata") or {}),
"tags": config.get("tags", None),
"name": config.get("run_name", None),
"run_id": config.get("run_id", None),
**(metadata or {}),
**({"message_id": message_id} if message_id else {}),
},
}
writer(evt)
config[CONF][CONFIG_KEY_SEND]([(state_key, evt)])
return evt
def delete_ui_message(id: str, *, state_key: str = "ui") -> RemoveUIMessage:
"""Delete a UI message by ID from the UI state.
This function creates and sends a message to remove a UI component from the current state.
It also updates the graph state to remove the UI message.
Args:
id: Unique identifier of the UI component to remove.
state_key: Key in the graph state where the UI messages are stored. Defaults to "ui".
Returns:
The remove UI message.
Example:
.. code-block:: python
delete_ui_message("message-123")
"""
writer = get_stream_writer()
config = get_config()
evt: RemoveUIMessage = {"type": "remove-ui", "id": id}
writer(evt)
config[CONF][CONFIG_KEY_SEND]([(state_key, evt)])
return evt
def ui_message_reducer(
left: Union[list[AnyUIMessage], AnyUIMessage],
right: Union[list[AnyUIMessage], AnyUIMessage],
) -> list[AnyUIMessage]:
"""Merge two lists of UI messages, supporting removing UI messages.
This function combines two lists of UI messages, handling both regular UI messages
and `remove-ui` messages. When a `remove-ui` message is encountered, it removes any
UI message with the matching ID from the current state.
Args:
left: First list of UI messages or single UI message.
right: Second list of UI messages or single UI message.
Returns:
Combined list of UI messages with removals applied.
Example:
.. code-block:: python
messages = ui_message_reducer(
[{"type": "ui", "id": "1", "name": "Chat", "props": {}}],
{"type": "remove-ui", "id": "1"}
)
"""
if not isinstance(left, list):
left = [left]
if not isinstance(right, list):
right = [right]
# merge messages
merged = left.copy()
merged_by_id = {m.get("id"): i for i, m in enumerate(merged)}
ids_to_remove = set()
for msg in right:
msg_id = msg.get("id")
if (existing_idx := merged_by_id.get(msg_id)) is not None:
if msg.get("type") == "remove-ui":
ids_to_remove.add(msg_id)
else:
ids_to_remove.discard(msg_id)
merged[existing_idx] = msg
else:
if msg.get("type") == "remove-ui":
raise ValueError(
f"Attempting to delete an UI message with an ID that doesn't exist ('{msg_id}')"
)
merged_by_id[msg_id] = len(merged)
merged.append(msg)
merged = [m for m in merged if m.get("id") not in ids_to_remove]
return merged