feat(ui): add merge option to UI messages (#4473)

- Add docs about (partial) streaming UI components from LLMs
- Add missing support for "nostream" in LangGraph
This commit is contained in:
David Duong
2025-04-30 20:15:57 +02:00
committed by GitHub
7 changed files with 238 additions and 27 deletions
+170 -14
View File
@@ -207,18 +207,6 @@ Behind the scenes, `LoadExternalComponent` will fetch the JS and CSS for the UI
## How-to guides
### Show loading UI when components are loading
You can provide a fallback UI to be rendered when the components are loading.
```tsx
<LoadExternalComponent
stream={thread}
message={ui}
fallback={<div>Loading...</div>}
/>
```
### Provide custom components on the client side
If you already have the components loaded in your client application, you can provide a map of such components to be rendered directly without fetching the UI code from LangGraph Platform.
@@ -235,6 +223,18 @@ const clientComponents = {
/>;
```
### Show loading UI when components are loading
You can provide a fallback UI to be rendered when the components are loading.
```tsx
<LoadExternalComponent
stream={thread}
message={ui}
fallback={<div>Loading...</div>}
/>
```
### Customise the namespace of UI components.
By default `LoadExternalComponent` will use the `assistantId` from `useStream()` hook to fetch the code for UI components. You can customise this by providing a `namespace` prop to the `LoadExternalComponent` component.
@@ -316,9 +316,9 @@ const WeatherComponent = (props: { city: string }) => {
};
```
### Streaming UI updates before the node execution is finished
### Streaming UI messages from the server
You can stream UI updates before the node execution is finished by using the `onCustomEvent` callback of the `useStream()` hook.
You can stream UI messages before the node execution is finished by using the `onCustomEvent` callback of the `useStream()` hook. This is especially useful when updating the UI component as the LLM is generating the response.
```tsx
import { uiMessageReducer } from "@langchain/langgraph-sdk/react-ui";
@@ -335,6 +335,162 @@ const { thread, submit } = useStream({
});
```
Then you can pushing updates to the UI component by calling `ui.push()` / `push_ui_message()` with the same ID as the UI message you wish to update.
=== "Python"
```python
from typing import Annotated, Sequence, TypedDict
from langchain_anthropic import ChatAnthropic
from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage
from langgraph.graph import StateGraph
from langgraph.graph.message import add_messages
from langgraph.graph.ui import AnyUIMessage, push_ui_message, ui_message_reducer
class AgentState(TypedDict): # noqa: D101
messages: Annotated[Sequence[BaseMessage], add_messages]
ui: Annotated[Sequence[AnyUIMessage], ui_message_reducer]
class CreateTextDocument(TypedDict):
"""Prepare a document heading for the user."""
title: str
async def writer_node(state: AgentState):
model = ChatAnthropic(model="claude-3-5-sonnet-latest")
message: AIMessage = await model.bind_tools(
tools=[CreateTextDocument],
tool_choice={"type": "tool", "name": "CreateTextDocument"},
).ainvoke(state["messages"])
tool_call = next(
(x["args"] for x in message.tool_calls if x["name"] == "CreateTextDocument"),
None,
)
if tool_call:
ui_message = push_ui_message("writer", tool_call, message=message)
ui_message_id = ui_message["id"]
# We're already streaming the LLM response to the client through UI messages
# so we don't need to stream it again to the `messages` stream mode.
content_stream = model.with_config({"tags": ["nostream"]}).astream(
f"Create a document with the title: {tool_call['title']}"
)
content: AIMessageChunk | None = None
async for chunk in content_stream:
content = content + chunk if content else chunk
push_ui_message(
"writer",
{"content": content.text()},
id=ui_message_id,
message=message,
# Use `merge=rue` to merge props with the existing UI message
merge=True,
)
return {"messages": [message]}
```
=== "JS"
```tsx
import {
Annotation,
MessagesAnnotation,
type LangGraphRunnableConfig,
} from "@langchain/langgraph";
import { z } from "zod";
import { ChatAnthropic } from "@langchain/anthropic";
import {
typedUi,
uiMessageReducer,
} from "@langchain/langgraph-sdk/react-ui/server";
import type { AIMessageChunk } from "@langchain/core/messages";
import type ComponentMap from "./ui";
const AgentState = Annotation.Root({
...MessagesAnnotation.spec,
ui: Annotation({ reducer: uiMessageReducer, default: () => [] }),
});
async function writerNode(
state: typeof AgentState.State,
config: LangGraphRunnableConfig
): Promise<typeof AgentState.Update> {
const ui = typedUi<typeof ComponentMap>(config);
const model = new ChatAnthropic({ model: "claude-3-5-sonnet-latest" });
const message = await model
.bindTools(
[
{
name: "create_text_document",
description: "Prepare a document heading for the user.",
schema: z.object({ title: z.string() }),
},
],
{ tool_choice: { type: "tool", name: "create_text_document" } }
)
.invoke(state.messages);
type ToolCall = { name: "create_text_document"; args: { title: string } };
const toolCall = message.tool_calls?.find(
(tool): tool is ToolCall => tool.name === "create_text_document"
);
if (toolCall) {
const { id, name } = ui.push(
{ name: "writer", props: { title: toolCall.args.title } },
{ message }
);
const contentStream = await model
// We're already streaming the LLM response to the client through UI messages
// so we don't need to stream it again to the `messages` stream mode.
.withConfig({ tags: ["nostream"] })
.stream(`Create a short poem with the topic: ${message.text}`);
let content: AIMessageChunk | undefined;
for await (const chunk of contentStream) {
content = content?.concat(chunk) ?? chunk;
ui.push(
{ id, name, props: { content: content?.text } },
// Use `merge: true` to merge props with the existing UI message
{ message, merge: true }
);
}
}
return { messages: [message] };
}
```
=== "`ui.tsx`"
```tsx
function WriterComponent(props: { title: string; content?: string }) {
return (
<article>
<h2>{props.title}</h2>
<p style={{ whiteSpace: "pre-wrap" }}>{props.content}</p>
</article>
);
}
export default {
weather: WriterComponent,
};
```
### 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 `remove_ui_message` / `ui.delete` with the ID of the UI message.
+3 -1
View File
@@ -14,8 +14,10 @@ EMPTY_SEQ: tuple[str, ...] = tuple()
MISSING = object()
# --- Public constants ---
TAG_NOSTREAM = sys.intern("langsmith:nostream")
TAG_NOSTREAM = sys.intern("nostream")
"""Tag to disable streaming for a chat model."""
TAG_NOSTREAM_ALT = sys.intern("langsmith:nostream")
"""Tag to disable streaming for a chat model. (Deprecated in favour of "nostream")"""
TAG_HIDDEN = sys.intern("langsmith:hidden")
"""Tag to hide a node/edge from certain tracing/streaming environments."""
START = sys.intern("__start__")
+9 -1
View File
@@ -1,4 +1,4 @@
from typing import Any, Literal, Optional, Union
from typing import Any, Literal, Optional, Union, cast
from uuid import uuid4
from langchain_core.messages import AnyMessage
@@ -55,6 +55,7 @@ def push_ui_message(
metadata: Optional[dict[str, Any]] = None,
message: Optional[AnyMessage] = None,
state_key: str = "ui",
merge: bool = False,
) -> UIMessage:
"""Push a new UI message to update the UI state.
@@ -100,6 +101,7 @@ def push_ui_message(
"name": name,
"props": props,
"metadata": {
"merge": merge,
"run_id": config.get("run_id", None),
"tags": config.get("tags", None),
"name": config.get("run_name", None),
@@ -191,6 +193,12 @@ def ui_message_reducer(
ids_to_remove.add(msg_id)
else:
ids_to_remove.discard(msg_id)
if cast(UIMessage, msg).get("metadata", {}).get("merge", False):
prev_msg = merged[existing_idx]
msg = msg.copy()
msg["props"] = {**prev_msg["props"], **msg["props"]}
merged[existing_idx] = msg
else:
if msg.get("type") == "remove-ui":
+4 -2
View File
@@ -13,7 +13,7 @@ from langchain_core.callbacks import BaseCallbackHandler
from langchain_core.messages import BaseMessage
from langchain_core.outputs import ChatGenerationChunk, LLMResult
from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM
from langgraph.constants import NS_SEP, TAG_HIDDEN, TAG_NOSTREAM, TAG_NOSTREAM_ALT
from langgraph.types import Command, StreamChunk
try:
@@ -93,7 +93,9 @@ class StreamMessagesHandler(BaseCallbackHandler, _StreamingCallbackHandler):
metadata: Optional[dict[str, Any]] = None,
**kwargs: Any,
) -> Any:
if metadata and (not tags or TAG_NOSTREAM not in tags):
if metadata and (
not tags or (TAG_NOSTREAM not in tags and TAG_NOSTREAM_ALT not in tags)
):
self.metadata[run_id] = (
tuple(cast(str, metadata["langgraph_checkpoint_ns"]).split(NS_SEP)),
metadata,
+2
View File
@@ -8,6 +8,8 @@ export {
} from "./client.js";
export {
uiMessageReducer,
isUIMessage,
isRemoveUIMessage,
type UIMessage,
type RemoveUIMessage,
} from "./types.js";
+26 -5
View File
@@ -38,21 +38,42 @@ export const typedUi = <Decl extends Record<string, ElementType>>(
const runId = (config.metadata?.run_id as string | undefined) ?? config.runId;
if (!runId) throw new Error("run_id is required");
const handlePush = <K extends keyof PropMap & string>(
function handlePush<K extends keyof PropMap & string>(
message: {
id?: string;
name: K;
props: PropMap[K];
metadata?: Record<string, unknown>;
},
options?: { message?: MessageLike },
): UIMessage => {
const evt: UIMessage = {
options?: { message?: MessageLike; merge?: boolean },
): UIMessage<K, PropMap[K]>;
function handlePush<K extends keyof PropMap & string>(
message: {
id?: string;
name: K;
props: Partial<PropMap[K]>;
metadata?: Record<string, unknown>;
},
options: { message?: MessageLike; merge: true },
): UIMessage<K, Partial<PropMap[K]>>;
function handlePush<K extends keyof PropMap & string>(
message: {
id?: string;
name: K;
props: PropMap[K] | Partial<PropMap[K]>;
metadata?: Record<string, unknown>;
},
options?: { message?: MessageLike; merge?: boolean },
): UIMessage<K, PropMap[K] | Partial<PropMap[K]>> {
const evt: UIMessage<K, PropMap[K] | Partial<PropMap[K]>> = {
type: "ui" as const,
id: message?.id ?? uuidv4(),
name: message?.name,
props: message?.props,
metadata: {
merge: options?.merge || undefined,
run_id: runId,
tags: config.tags,
name: config.runName,
@@ -64,7 +85,7 @@ export const typedUi = <Decl extends Record<string, ElementType>>(
config.writer?.(evt);
config.configurable?.__pregel_send?.([[stateKey, evt]]);
return evt;
};
}
const handleDelete = (id: string): RemoveUIMessage => {
const evt: RemoveUIMessage = { type: "remove-ui", id };
+24 -4
View File
@@ -1,10 +1,14 @@
export interface UIMessage {
export interface UIMessage<
TName extends string = string,
TProps extends Record<string, unknown> = Record<string, unknown>,
> {
type: "ui";
id: string;
name: string;
props: Record<string, unknown>;
name: TName;
props: TProps;
metadata: {
merge?: boolean;
run_id?: string;
name?: string;
tags?: string[];
@@ -18,6 +22,20 @@ export interface RemoveUIMessage {
id: string;
}
export function isUIMessage(message: unknown): message is UIMessage {
if (typeof message !== "object" || message == null) return false;
if (!("type" in message)) return false;
return message.type === "ui";
}
export function isRemoveUIMessage(
message: unknown,
): message is RemoveUIMessage {
if (typeof message !== "object" || message == null) return false;
if (!("type" in message)) return false;
return message.type === "remove-ui";
}
export function uiMessageReducer(
state: UIMessage[],
update: UIMessage | RemoveUIMessage | (UIMessage | RemoveUIMessage)[],
@@ -33,7 +51,9 @@ export function uiMessageReducer(
const index = state.findIndex((ui) => ui.id === event.id);
if (index !== -1) {
newState[index] = event;
newState[index] = event.metadata.merge
? { ...event, props: { ...state[index].props, ...event.props } }
: event;
} else {
newState.push(event);
}