From 1b6e12ef1426592f71330b9859975a81a852c875 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 30 Apr 2025 11:52:47 +0200 Subject: [PATCH 1/4] feat(ui): add `merge` option to UI messages --- libs/langgraph/langgraph/graph/ui.py | 8 +++++++ libs/sdk-js/src/react-ui/index.ts | 2 ++ libs/sdk-js/src/react-ui/server/server.ts | 29 +++++++++++++++++++---- libs/sdk-js/src/react-ui/types.ts | 19 ++++++++++++++- 4 files changed, 53 insertions(+), 5 deletions(-) diff --git a/libs/langgraph/langgraph/graph/ui.py b/libs/langgraph/langgraph/graph/ui.py index aac9d2044..871b6c108 100644 --- a/libs/langgraph/langgraph/graph/ui.py +++ b/libs/langgraph/langgraph/graph/ui.py @@ -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 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": diff --git a/libs/sdk-js/src/react-ui/index.ts b/libs/sdk-js/src/react-ui/index.ts index e54ce73e6..e63e79d16 100644 --- a/libs/sdk-js/src/react-ui/index.ts +++ b/libs/sdk-js/src/react-ui/index.ts @@ -8,6 +8,8 @@ export { } from "./client.js"; export { uiMessageReducer, + isUIMessage, + isRemoveUIMessage, type UIMessage, type RemoveUIMessage, } from "./types.js"; diff --git a/libs/sdk-js/src/react-ui/server/server.ts b/libs/sdk-js/src/react-ui/server/server.ts index 708385694..8349121d0 100644 --- a/libs/sdk-js/src/react-ui/server/server.ts +++ b/libs/sdk-js/src/react-ui/server/server.ts @@ -38,21 +38,42 @@ export const typedUi = >( const runId = (config.metadata?.run_id as string | undefined) ?? config.runId; if (!runId) throw new Error("run_id is required"); - const handlePush = ( + function handlePush( message: { id?: string; name: K; props: PropMap[K]; metadata?: Record; }, - options?: { message?: MessageLike }, - ): UIMessage => { + options?: { message?: MessageLike; merge?: boolean }, + ): UIMessage; + + function handlePush( + message: { + id?: string; + name: K; + props: Partial; + metadata?: Record; + }, + options: { message?: MessageLike; merge: true }, + ): UIMessage; + + function handlePush( + message: { + id?: string; + name: K; + props: PropMap[K] | Partial; + metadata?: Record; + }, + options?: { message?: MessageLike; merge?: boolean }, + ): UIMessage { const evt: UIMessage = { 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 = >( config.writer?.(evt); config.configurable?.__pregel_send?.([[stateKey, evt]]); return evt; - }; + } const handleDelete = (id: string): RemoveUIMessage => { const evt: RemoveUIMessage = { type: "remove-ui", id }; diff --git a/libs/sdk-js/src/react-ui/types.ts b/libs/sdk-js/src/react-ui/types.ts index d60d81728..65efaa2b9 100644 --- a/libs/sdk-js/src/react-ui/types.ts +++ b/libs/sdk-js/src/react-ui/types.ts @@ -5,6 +5,7 @@ export interface UIMessage { name: string; props: Record; metadata: { + merge?: boolean; run_id?: string; name?: string; tags?: string[]; @@ -18,6 +19,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 +48,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); } From 5804e788d8fb7a0c7e464578f5ffc2a046ebcf58 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 30 Apr 2025 19:46:14 +0200 Subject: [PATCH 2/4] Add docs --- .../docs/cloud/how-tos/generative_ui_react.md | 184 ++++++++++++++++-- libs/sdk-js/src/react-ui/server/server.ts | 8 +- libs/sdk-js/src/react-ui/types.ts | 9 +- 3 files changed, 180 insertions(+), 21 deletions(-) diff --git a/docs/docs/cloud/how-tos/generative_ui_react.md b/docs/docs/cloud/how-tos/generative_ui_react.md index 0c7a2e0cf..ab96c4a80 100644 --- a/docs/docs/cloud/how-tos/generative_ui_react.md +++ b/docs/docs/cloud/how-tos/generative_ui_react.md @@ -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 -Loading...} -/> -``` - ### 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 +Loading...} +/> +``` + ### 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 { + const ui = typedUi(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 ( +
+

{props.title}

+

{props.content}

+
+ ); + } + + 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. diff --git a/libs/sdk-js/src/react-ui/server/server.ts b/libs/sdk-js/src/react-ui/server/server.ts index 8349121d0..8eb14d17f 100644 --- a/libs/sdk-js/src/react-ui/server/server.ts +++ b/libs/sdk-js/src/react-ui/server/server.ts @@ -46,7 +46,7 @@ export const typedUi = >( metadata?: Record; }, options?: { message?: MessageLike; merge?: boolean }, - ): UIMessage; + ): UIMessage; function handlePush( message: { @@ -56,7 +56,7 @@ export const typedUi = >( metadata?: Record; }, options: { message?: MessageLike; merge: true }, - ): UIMessage; + ): UIMessage>; function handlePush( message: { @@ -66,8 +66,8 @@ export const typedUi = >( metadata?: Record; }, options?: { message?: MessageLike; merge?: boolean }, - ): UIMessage { - const evt: UIMessage = { + ): UIMessage> { + const evt: UIMessage> = { type: "ui" as const, id: message?.id ?? uuidv4(), name: message?.name, diff --git a/libs/sdk-js/src/react-ui/types.ts b/libs/sdk-js/src/react-ui/types.ts index 65efaa2b9..32bfe4188 100644 --- a/libs/sdk-js/src/react-ui/types.ts +++ b/libs/sdk-js/src/react-ui/types.ts @@ -1,9 +1,12 @@ -export interface UIMessage { +export interface UIMessage< + TName extends string = string, + TProps extends Record = Record, +> { type: "ui"; id: string; - name: string; - props: Record; + name: TName; + props: TProps; metadata: { merge?: boolean; run_id?: string; From b3371a1d63c813fc6b192f90bb6cd3b3d43fe65b Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 30 Apr 2025 19:50:14 +0200 Subject: [PATCH 3/4] Add missing "nostream" support --- libs/langgraph/langgraph/constants.py | 4 +++- libs/langgraph/langgraph/pregel/messages.py | 6 ++++-- 2 files changed, 7 insertions(+), 3 deletions(-) diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 59e72e9c9..167c340f1 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -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__") diff --git a/libs/langgraph/langgraph/pregel/messages.py b/libs/langgraph/langgraph/pregel/messages.py index de6425309..ef82bd483 100644 --- a/libs/langgraph/langgraph/pregel/messages.py +++ b/libs/langgraph/langgraph/pregel/messages.py @@ -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, From be1af772f50597c96d8ba3f06b4b48774d5182ac Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Wed, 30 Apr 2025 20:04:48 +0200 Subject: [PATCH 4/4] Fix lint --- libs/langgraph/langgraph/graph/ui.py | 4 +-- libs/langgraph/tests/test_pregel.py | 37 +++++++++++----------------- 2 files changed, 17 insertions(+), 24 deletions(-) diff --git a/libs/langgraph/langgraph/graph/ui.py b/libs/langgraph/langgraph/graph/ui.py index 871b6c108..08dc6fac7 100644 --- a/libs/langgraph/langgraph/graph/ui.py +++ b/libs/langgraph/langgraph/graph/ui.py @@ -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 @@ -194,7 +194,7 @@ def ui_message_reducer( else: ids_to_remove.discard(msg_id) - if msg.get("metadata", {}).get("merge", False): + if cast(UIMessage, msg).get("metadata", {}).get("merge", False): prev_msg = merged[existing_idx] msg = msg.copy() msg["props"] = {**prev_msg["props"], **msg["props"]} diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 75d408f29..d15c037c8 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6351,9 +6351,7 @@ def test_double_interrupt_subgraph( @pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC) -def test_multi_resume( - request: pytest.FixtureRequest, checkpointer_name: str -) -> None: +def test_multi_resume(request: pytest.FixtureRequest, checkpointer_name: str) -> None: checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") class ChildState(TypedDict): @@ -6362,11 +6360,11 @@ def test_multi_resume( human_inputs: list[str] def get_human_input(state: ChildState): - human_input = interrupt(state['prompt']) + human_input = interrupt(state["prompt"]) return { - 'human_input': human_input, - 'human_inputs': [human_input], + "human_input": human_input, + "human_inputs": [human_input], } child_graph = ( @@ -6385,13 +6383,13 @@ def test_multi_resume( return [ Send( "child_graph", - {'prompt': prompt}, + {"prompt": prompt}, ) - for prompt in state['prompts'] + for prompt in state["prompts"] ] def cleanup(state: ParentState): - assert len(state['human_inputs']) == len(state["prompts"]) + assert len(state["human_inputs"]) == len(state["prompts"]) parent_graph = ( StateGraph(ParentState) @@ -6404,21 +6402,19 @@ def test_multi_resume( ) thread_config: RunnableConfig = { - 'configurable': { - 'thread_id': uuid.uuid4(), + "configurable": { + "thread_id": uuid.uuid4(), }, } - prompts = ['a', 'b', 'c', 'd', 'e'] + prompts = ["a", "b", "c", "d", "e"] events = parent_graph.invoke( - {'prompts': prompts}, - thread_config, - stream_mode='values' + {"prompts": prompts}, thread_config, stream_mode="values" ) - assert len(events['__interrupt__']) == len(prompts) - interrupt_values = {i.value for i in events['__interrupt__']} + assert len(events["__interrupt__"]) == len(prompts) + interrupt_values = {i.value for i in events["__interrupt__"]} assert interrupt_values == set(prompts) resume_map: dict[str, str] = { @@ -6428,11 +6424,8 @@ def test_multi_resume( result = parent_graph.invoke(Command(resume=resume_map), thread_config) assert result == { - 'prompts': prompts, - 'human_inputs': [ - f"human input for prompt {prompt}" - for prompt in prompts - ], + "prompts": prompts, + "human_inputs": [f"human input for prompt {prompt}" for prompt in prompts], }