feat(sdk-js): automatically write ui messages (#3833)

docs/docs/cloud/how-tos/generative_ui_react.md

@@ -103,7 +103,8 @@ const AgentState = Annotation.Root({
 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.
+    // 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" })
@@ -120,7 +121,7 @@ export const graph = new StateGraph(AgentState)
     // Emit UI elements with associated AI message
     ui.push({ name: "weather", props: weather }, { message: response });
 
-    return { messages: [response], ui: ui.items };
+    return { messages: [response] };
   })
   .addEdge("__start__", "weather")
   .compile();

libs/sdk-js/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@langchain/langgraph-sdk",
-  "version": "0.0.56",
+  "version": "0.0.57",
   "description": "Client library for interacting with the LangGraph API",
   "type": "module",
   "packageManager": "yarn@1.22.19",

libs/sdk-js/src/react-ui/server/server.ts

@@ -6,15 +6,34 @@ interface MessageLike {
   id?: string;
 }
 
-export const typedUi = <Decl extends Record<string, ElementType>>(config: {
-  writer?: (chunk: unknown) => void;
-  runId?: string;
-  metadata?: Record<string, unknown>;
-  tags?: string[];
-  runName?: string;
-}) => {
+/**
+ * Helper to send and persist UI messages. Accepts a map of component names to React components
+ * as type argument to provide type safety. Will also write to the `options?.stateKey` state.
+ *
+ * @param config LangGraphRunnableConfig
+ * @param options
+ * @returns
+ */
+export const typedUi = <Decl extends Record<string, ElementType>>(
+  config: {
+    writer?: (chunk: unknown) => void;
+    runId?: string;
+    metadata?: Record<string, unknown>;
+    tags?: string[];
+    runName?: string;
+    configurable?: {
+      __pregel_send?: (writes_: [string, unknown][]) => void;
+      [key: string]: unknown;
+    };
+  },
+  options?: {
+    /** The key to write the UI messages to. Defaults to `ui`. */
+    stateKey?: string;
+  },
+) => {
   type PropMap = { [K in keyof Decl]: ComponentPropsWithoutRef<Decl[K]> };
   let items: (UIMessage | RemoveUIMessage)[] = [];
+  const stateKey = options?.stateKey ?? "ui";
 
   const runId = (config.metadata?.run_id as string | undefined) ?? config.runId;
   if (!runId) throw new Error("run_id is required");
@@ -48,6 +67,7 @@ export const typedUi = <Decl extends Record<string, ElementType>>(config: {
     };
     items.push(evt);
     config.writer?.(evt);
+    config.configurable?.__pregel_send?.([[stateKey, evt]]);
     return evt;
   };
 
@@ -55,6 +75,7 @@ export const typedUi = <Decl extends Record<string, ElementType>>(config: {
     const evt: RemoveUIMessage = { type: "remove-ui", id };
     items.push(evt);
     config.writer?.(evt);
+    config.configurable?.__pregel_send?.([[stateKey, evt]]);
     return evt;
   };