improve ui streaming

ai/tools/github_repo.tsx

@@ -44,7 +44,7 @@ export const githubTool = new DynamicStructuredTool({
     "A tool to fetch details of a Github repository. Given owner and repo names, this tool will return the repo description, stars, and primary language.",
   schema: githubRepoToolSchema,
   func: async (input, config) => {
-    const stream = createRunnableUI(config, <GithubLoading />);
+    const stream = await createRunnableUI(config, <GithubLoading />);
     const result = await githubRepoTool(input);
     if (typeof result === "string") {
       // Failed to parse, return error message

ai/tools/invoice.tsx

@@ -57,7 +57,7 @@ export const invoiceTool = new DynamicStructuredTool({
     "A tool to fetch the invoice from an order. This should only be called if a user uploads an image/receipt of an order.",
   schema: InvoiceSchema,
   func: async (input, config) => {
-    const stream = createRunnableUI(config, <InvoiceLoading />);
+    const stream = await createRunnableUI(config, <InvoiceLoading />);
     stream.done(<Invoice {...input} />);
     return JSON.stringify(input, null);
   },

ai/tools/weather.tsx

@@ -65,7 +65,7 @@ export const weatherTool = new DynamicStructuredTool({
     "A tool to fetch the current weather, given a city and state. If the city/state is not provided, ask the user for both the city and state.",
   schema: weatherSchema,
   func: async (input, config) => {
-    const stream = createRunnableUI(config, <CurrentWeatherLoading />);
+    const stream = await createRunnableUI(config, <CurrentWeatherLoading />);
     const data = await weatherData(input);
     stream.done(<CurrentWeather {...data} />);
     return JSON.stringify(data, null);

components/prebuilt/chat.tsx

@@ -84,10 +84,13 @@ export default function Chat() {
           setHistory((prev) => [
             ...prev,
             ["user", input],
-            ["assistant", `Tool result: ${JSON.stringify(lastEvent["invokeTools"]["toolResult"], null)}`],
+            [
+              "assistant",
+              `Tool result: ${JSON.stringify(lastEvent["invokeTools"]["toolResult"], null)}`,
+            ],
           ]);
         } else {
-          console.log("ELSE!", lastEvent)
+          console.log("ELSE!", lastEvent);
         }
       }
     })();

utils/server.tsx

@@ -2,7 +2,11 @@ import "server-only";
 
 import { ReactNode, isValidElement } from "react";
 import { createStreamableUI, createStreamableValue } from "ai/rsc";
-import { Runnable } from "@langchain/core/runnables";
+import {
+  Runnable,
+  RunnableConfig,
+  RunnableLambda,
+} from "@langchain/core/runnables";
 import {
   CallbackManagerForToolRun,
   CallbackManagerForRetrieverRun,
@@ -18,6 +22,8 @@ import { AIProvider } from "./client";
 import { AIMessage } from "../ai/message";
 import { CompiledStateGraph } from "@langchain/langgraph";
 
+const STREAM_UI_RUN_NAME = "stream_ui_lambda";
+
 /**
  * Executes `streamEvents` method on a runnable
  * and converts the generator to a RSC friendly stream
@@ -47,10 +53,17 @@ export function streamRunnableUI<RunInput, RunOutput>(
     ).streamEvents(inputs, {
       version: "v1",
     })) {
+      if (
+        streamEvent.name === STREAM_UI_RUN_NAME &&
+        streamEvent.event === "on_chain_end"
+      ) {
+        if (isValidElement(streamEvent.data.output.value)) {
+          ui.append(streamEvent.data.output.value);
+        }
+      }
       const [kind, type] = streamEvent.event.split("_").slice(1);
       if (type === "stream" && kind !== "chain") {
         const chunk = streamEvent.data.chunk;
-
         if (isValidElement(chunk)) {
           ui.append(chunk);
         } else if ("text" in chunk && typeof chunk.text === "string") {
@@ -88,44 +101,24 @@ export function streamRunnableUI<RunInput, RunOutput>(
  * @param initialValue Initial React node to be sent to the client
  * @returns Vercel AI RSC compatible streamable UI
  */
-export const createRunnableUI = (
+export const createRunnableUI = async (
   config:
     | CallbackManagerForToolRun
     | CallbackManagerForRetrieverRun
     | CallbackManagerForChainRun
-    | CallbackManagerForLLMRun
     | undefined,
   initialValue?: React.ReactNode,
-): ReturnType<typeof createStreamableUI> => {
+): Promise<ReturnType<typeof createStreamableUI>> => {
   if (!config) throw new Error("No config provided");
 
-  const logStreamTracer = config.handlers.find(
-    (i): i is LogStreamCallbackHandler => i.name === "log_stream_tracer",
-  );
-
-  const ui = createStreamableUI(initialValue);
-
-  if (!logStreamTracer) throw new Error("No log stream tracer found");
-  // @ts-expect-error Private field
-  const runName = logStreamTracer.keyMapByRunId[config.runId];
-  if (!runName) {
-    console.log("No name found for", config.runId);
-    throw new Error("No run name found");
-  }
-
-  logStreamTracer.writer.write(
-    new RunLogPatch({
-      ops: [
-        {
-          op: "add",
-          path: `/logs/${runName}/streamed_output/-`,
-          value: ui.value,
-        },
-      ],
-    }),
-  );
-
-  return ui;
+  const lambda = RunnableLambda.from(
+    (init: React.ReactNode, config?: RunnableConfig) => {
+      const ui = createStreamableUI(init);
+      return ui;
+    },
+  ).withConfig({ runName: STREAM_UI_RUN_NAME });
+
+  return lambda.invoke(initialValue, { callbacks: config.getChild() });
 };
 
 /**