improved agent

agent/agent.tsx

@@ -9,12 +9,32 @@ import { SystemMessage } from "@langchain/core/messages";
 import { ChatOpenAI } from "@langchain/openai";
 import { typedUi } from "@langchain/langgraph-sdk/react-ui/server";
 import { uiMessageReducer } from "@langchain/langgraph-sdk/react-ui/types";
-import type ComponentMap from "./ui";
+import type ComponentMap from "./uis/index";
 import { z, ZodTypeAny } from "zod";
 
-// const llm = new ChatOllama({ model: "deepseek-r1" });
 const llm = new ChatOpenAI({ model: "gpt-4o-mini", temperature: 0 });
 
+const getStockPriceSchema = z.object({
+  ticker: z.string().describe("The ticker symbol of the company"),
+});
+const getPortfolioSchema = z.object({
+  get_portfolio: z.boolean().describe("Should be true."),
+});
+
+const STOCKBROKER_TOOLS = [
+  {
+    name: "get_stock_price",
+    description: "A tool to get the stock price of a company",
+    schema: getStockPriceSchema,
+  },
+  {
+    name: "get_portfolio",
+    description:
+      "A tool to get the user's portfolio details. Only call this tool if the user requests their portfolio details.",
+    schema: getPortfolioSchema,
+  },
+];
+
 interface ToolCall {
   name: string;
   args: Record<string, any>;
@@ -24,7 +44,7 @@ interface ToolCall {
 
 function findToolCall<Name extends string>(name: Name) {
   return <Args extends ZodTypeAny>(
-    x: ToolCall
+    x: ToolCall,
   ): x is { name: Name; args: z.infer<Args> } => x.name === name;
 }
 
@@ -33,46 +53,37 @@ const builder = new StateGraph(
     messages: MessagesAnnotation.spec["messages"],
     ui: Annotation({ default: () => [], reducer: uiMessageReducer }),
     timestamp: Annotation<number>,
-  })
+  }),
 )
   .addNode("agent", async (state, config) => {
     const ui = typedUi<typeof ComponentMap>(config);
 
-    // const result = ui.interrupt("react-component", {
-    //   instruction: "Hello world",
-    // });
-
-    // // throw new Error("Random error");
-    // // stream custom events
-    // for (let count = 0; count < 10; count++) config.writer?.({ count });
-
-    // How do I properly assign
-    const stockbrokerSchema = z.object({ company: z.string() });
     const message = await llm
-      .bindTools([
+      .bindTools(STOCKBROKER_TOOLS)
-        {
-          name: "stockbroker",
-          description: "A tool to get the stock price of a company",
-          schema: stockbrokerSchema,
-        },
-      ])
       .invoke([
         new SystemMessage(
-          "You are a stockbroker agent that uses tools to get the stock price of a company"
+          "You are a stockbroker agent that uses tools to get the stock price of a company",
         ),
         ...state.messages,
       ]);
 
     const stockbrokerToolCall = message.tool_calls?.find(
-      findToolCall("stockbroker")<typeof stockbrokerSchema>
+      findToolCall("get_stock_price")<typeof getStockPriceSchema>,
+    );
+    const portfolioToolCall = message.tool_calls?.find(
+      findToolCall("get_portfolio")<typeof getStockPriceSchema>,
     );
 
     if (stockbrokerToolCall) {
       const instruction = `The stock price of ${
-        stockbrokerToolCall.args.company
+        stockbrokerToolCall.args.ticker
       } is ${Math.random() * 100}`;
 
-      ui.write("react-component", { instruction, logo: "hey" });
+      ui.write("stock-price", { instruction, logo: "hey" });
+    }
+
+    if (portfolioToolCall) {
+      ui.write("portfolio-view", {});
     }
 
     return { messages: message, ui: ui.collect, timestamp: Date.now() };

agent/uis/index.tsx

@@ -0,0 +1,8 @@
+import StockPrice from "./stock-price";
+import PortfolioView from "./portfolio-view";
+
+const ComponentMap = {
+  "stock-price": StockPrice,
+  "portfolio-view": PortfolioView,
+} as const;
+export default ComponentMap;

agent/uis/portfolio-view/index.tsx

@@ -0,0 +1,9 @@
+import "./index.css";
+
+export default function PortfolioView() {
+  return (
+    <div className="flex flex-col gap-2 border border-solid border-slate-500 p-4 rounded-md">
+      Portfolio View
+    </div>
+  );
+}

agent/uis/stock-price/index.css

@@ -0,0 +1 @@
+@import "tailwindcss";

agent/uis/stock-price/index.tsx

@@ -1,9 +1,12 @@
-import "./ui.css";
+import "./index.css";
 import { useStream } from "@langchain/langgraph-sdk/react";
 import type { AIMessage, Message } from "@langchain/langgraph-sdk";
 import { useState } from "react";
 
-function ReactComponent(props: { instruction: string; logo: string }) {
+export default function StockPrice(props: {
+  instruction: string;
+  logo: string;
+}) {
   const [counter, setCounter] = useState(0);
 
   // useStream should be able to be infered from context
@@ -17,7 +20,7 @@ function ReactComponent(props: { instruction: string; logo: string }) {
     .reverse()
     .find(
       (message): message is AIMessage =>
-        message.type === "ai" && !!message.tool_calls?.length
+        message.type === "ai" && !!message.tool_calls?.length,
     );
 
   const toolCallId = aiTool?.tool_calls?.[0]?.id;
@@ -52,6 +55,3 @@ function ReactComponent(props: { instruction: string; logo: string }) {
     </div>
   );
 }
-
-const ComponentMap = { "react-component": ReactComponent } as const;
-export default ComponentMap;

eslint.config.js

@@ -1,28 +1,28 @@
-import js from '@eslint/js'
+import js from "@eslint/js";
-import globals from 'globals'
+import globals from "globals";
-import reactHooks from 'eslint-plugin-react-hooks'
+import reactHooks from "eslint-plugin-react-hooks";
-import reactRefresh from 'eslint-plugin-react-refresh'
+import reactRefresh from "eslint-plugin-react-refresh";
-import tseslint from 'typescript-eslint'
+import tseslint from "typescript-eslint";
 
 export default tseslint.config(
-  { ignores: ['dist'] },
+  { ignores: ["dist"] },
   {
     extends: [js.configs.recommended, ...tseslint.configs.recommended],
-    files: ['**/*.{ts,tsx}'],
+    files: ["**/*.{ts,tsx}"],
     languageOptions: {
       ecmaVersion: 2020,
       globals: globals.browser,
     },
     plugins: {
-      'react-hooks': reactHooks,
+      "react-hooks": reactHooks,
-      'react-refresh': reactRefresh,
+      "react-refresh": reactRefresh,
     },
     rules: {
       ...reactHooks.configs.recommended.rules,
-      'react-refresh/only-export-components': [
+      "react-refresh/only-export-components": [
-        'warn',
+        "warn",
         { allowConstantExport: true },
       ],
     },
   },
-)
+);

index.html

@@ -5,7 +5,7 @@
     <link rel="icon" type="image/svg+xml" href="/vite.svg" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <title>Vite + React + TS</title>
-    <link href="/src/styles.css" rel="stylesheet">
+    <link href="/src/styles.css" rel="stylesheet" />
   </head>
   <body>
     <div id="root"></div>

package.json

@@ -8,13 +8,14 @@
     "dev": "vite",
     "build": "tsc -b && vite build",
     "lint": "eslint .",
+    "format": "prettier --write .",
     "preview": "vite preview"
   },
   "dependencies": {
     "@assistant-ui/react": "^0.8.0",
     "@assistant-ui/react-markdown": "^0.8.0",
-    "@langchain/core": "^0.3.40",
+    "@langchain/core": "^0.3.41",
-    "@langchain/langgraph": "^0.2.46",
+    "@langchain/langgraph": "^0.2.49",
     "@langchain/langgraph-api": "*",
     "@langchain/langgraph-cli": "*",
     "@langchain/langgraph-sdk": "*",
@@ -29,6 +30,7 @@
     "esbuild": "^0.25.0",
     "esbuild-plugin-tailwindcss": "^2.0.1",
     "lucide-react": "^0.476.0",
+    "prettier": "^3.5.2",
     "react": "^19.0.0",
     "react-dom": "^19.0.0",
     "remark-gfm": "^4.0.1",

src/App.tsx

@@ -1,81 +1,12 @@
 import "./App.css";
-import { useStream } from "@langchain/langgraph-sdk/react";
-import type { Message } from "@langchain/langgraph-sdk";
-import type {
-  UIMessage,
-  RemoveUIMessage,
-} from "@langchain/langgraph-sdk/react-ui/types";
-import { LoadExternalComponent } from "@langchain/langgraph-sdk/react-ui/client";
 import { Thread } from "@/components/assistant-ui/thread";
 
 function App() {
-  const thread = useStream<
+  return (
-    { messages: Message[]; ui: UIMessage[] },
-    {
-      messages?: Message[] | Message | string;
-      ui?: (UIMessage | RemoveUIMessage)[] | UIMessage | RemoveUIMessage;
-    },
-    UIMessage | RemoveUIMessage
-  >({
-    apiUrl: "http://localhost:2024",
-    assistantId: "agent",
-  });
-
-   return (
     <div className="h-full">
       <Thread />
     </div>
   );
-
-  return (
-    <>
-      <div className="grid grid-cols-2 gap-2">
-        {thread.messages.map((message, idx) => {
-          const meta = thread.getMessagesMetadata(message, idx);
-          const seenState = meta?.firstSeenState;
-          const customComponent = seenState?.values.ui
-            .slice()
-            .reverse()
-            .find(
-              ({ additional_kwargs }) =>
-                additional_kwargs.run_id === seenState.metadata?.run_id
-            );
-
-          return (
-            <div key={message.id}>
-              <pre>{JSON.stringify(message, null, 2)}</pre>
-              {customComponent && (
-                <LoadExternalComponent
-                  assistantId="agent"
-                  stream={thread}
-                  message={customComponent}
-                />
-              )}
-            </div>
-          );
-        })}
-      </div>
-      <form
-        onSubmit={(e) => {
-          e.preventDefault();
-
-          const formData = new FormData(e.target as HTMLFormElement);
-          const message = formData.get("message");
-          if (typeof message !== "string") return;
-          thread.submit({
-            messages: [{ type: "human", content: message }],
-          });
-        }}
-      >
-        <input
-          type="text"
-          name="message"
-          defaultValue="What's the price of AAPL?"
-        />
-        <button type="submit">Send</button>
-      </form>
-    </>
-  );
 }
 
 export default App;

src/components/assistant-ui/markdown-text.tsx

@@ -66,64 +66,151 @@ const useCopyToClipboard = ({
 
 const defaultComponents = memoizeMarkdownComponents({
   h1: ({ className, ...props }) => (
-    <h1 className={cn("mb-8 scroll-m-20 text-4xl font-extrabold tracking-tight last:mb-0", className)} {...props} />
+    <h1
+      className={cn(
+        "mb-8 scroll-m-20 text-4xl font-extrabold tracking-tight last:mb-0",
+        className,
+      )}
+      {...props}
+    />
   ),
   h2: ({ className, ...props }) => (
-    <h2 className={cn("mb-4 mt-8 scroll-m-20 text-3xl font-semibold tracking-tight first:mt-0 last:mb-0", className)} {...props} />
+    <h2
+      className={cn(
+        "mb-4 mt-8 scroll-m-20 text-3xl font-semibold tracking-tight first:mt-0 last:mb-0",
+        className,
+      )}
+      {...props}
+    />
   ),
   h3: ({ className, ...props }) => (
-    <h3 className={cn("mb-4 mt-6 scroll-m-20 text-2xl font-semibold tracking-tight first:mt-0 last:mb-0", className)} {...props} />
+    <h3
+      className={cn(
+        "mb-4 mt-6 scroll-m-20 text-2xl font-semibold tracking-tight first:mt-0 last:mb-0",
+        className,
+      )}
+      {...props}
+    />
   ),
   h4: ({ className, ...props }) => (
-    <h4 className={cn("mb-4 mt-6 scroll-m-20 text-xl font-semibold tracking-tight first:mt-0 last:mb-0", className)} {...props} />
+    <h4
+      className={cn(
+        "mb-4 mt-6 scroll-m-20 text-xl font-semibold tracking-tight first:mt-0 last:mb-0",
+        className,
+      )}
+      {...props}
+    />
   ),
   h5: ({ className, ...props }) => (
-    <h5 className={cn("my-4 text-lg font-semibold first:mt-0 last:mb-0", className)} {...props} />
+    <h5
+      className={cn(
+        "my-4 text-lg font-semibold first:mt-0 last:mb-0",
+        className,
+      )}
+      {...props}
+    />
   ),
   h6: ({ className, ...props }) => (
-    <h6 className={cn("my-4 font-semibold first:mt-0 last:mb-0", className)} {...props} />
+    <h6
+      className={cn("my-4 font-semibold first:mt-0 last:mb-0", className)}
+      {...props}
+    />
   ),
   p: ({ className, ...props }) => (
-    <p className={cn("mb-5 mt-5 leading-7 first:mt-0 last:mb-0", className)} {...props} />
+    <p
+      className={cn("mb-5 mt-5 leading-7 first:mt-0 last:mb-0", className)}
+      {...props}
+    />
   ),
   a: ({ className, ...props }) => (
-    <a className={cn("text-primary font-medium underline underline-offset-4", className)} {...props} />
+    <a
+      className={cn(
+        "text-primary font-medium underline underline-offset-4",
+        className,
+      )}
+      {...props}
+    />
   ),
   blockquote: ({ className, ...props }) => (
-    <blockquote className={cn("border-l-2 pl-6 italic", className)} {...props} />
+    <blockquote
+      className={cn("border-l-2 pl-6 italic", className)}
+      {...props}
+    />
   ),
   ul: ({ className, ...props }) => (
-    <ul className={cn("my-5 ml-6 list-disc [&>li]:mt-2", className)} {...props} />
+    <ul
+      className={cn("my-5 ml-6 list-disc [&>li]:mt-2", className)}
+      {...props}
+    />
   ),
   ol: ({ className, ...props }) => (
-    <ol className={cn("my-5 ml-6 list-decimal [&>li]:mt-2", className)} {...props} />
+    <ol
+      className={cn("my-5 ml-6 list-decimal [&>li]:mt-2", className)}
+      {...props}
+    />
   ),
   hr: ({ className, ...props }) => (
     <hr className={cn("my-5 border-b", className)} {...props} />
   ),
   table: ({ className, ...props }) => (
-    <table className={cn("my-5 w-full border-separate border-spacing-0 overflow-y-auto", className)} {...props} />
+    <table
+      className={cn(
+        "my-5 w-full border-separate border-spacing-0 overflow-y-auto",
+        className,
+      )}
+      {...props}
+    />
   ),
   th: ({ className, ...props }) => (
-    <th className={cn("bg-muted px-4 py-2 text-left font-bold first:rounded-tl-lg last:rounded-tr-lg [&[align=center]]:text-center [&[align=right]]:text-right", className)} {...props} />
+    <th
+      className={cn(
+        "bg-muted px-4 py-2 text-left font-bold first:rounded-tl-lg last:rounded-tr-lg [&[align=center]]:text-center [&[align=right]]:text-right",
+        className,
+      )}
+      {...props}
+    />
   ),
   td: ({ className, ...props }) => (
-    <td className={cn("border-b border-l px-4 py-2 text-left last:border-r [&[align=center]]:text-center [&[align=right]]:text-right", className)} {...props} />
+    <td
+      className={cn(
+        "border-b border-l px-4 py-2 text-left last:border-r [&[align=center]]:text-center [&[align=right]]:text-right",
+        className,
+      )}
+      {...props}
+    />
   ),
   tr: ({ className, ...props }) => (
-    <tr className={cn("m-0 border-b p-0 first:border-t [&:last-child>td:first-child]:rounded-bl-lg [&:last-child>td:last-child]:rounded-br-lg", className)} {...props} />
+    <tr
+      className={cn(
+        "m-0 border-b p-0 first:border-t [&:last-child>td:first-child]:rounded-bl-lg [&:last-child>td:last-child]:rounded-br-lg",
+        className,
+      )}
+      {...props}
+    />
   ),
   sup: ({ className, ...props }) => (
-    <sup className={cn("[&>a]:text-xs [&>a]:no-underline", className)} {...props} />
+    <sup
+      className={cn("[&>a]:text-xs [&>a]:no-underline", className)}
+      {...props}
+    />
   ),
   pre: ({ className, ...props }) => (
-    <pre className={cn("overflow-x-auto rounded-b-lg bg-black p-4 text-white", className)} {...props} />
+    <pre
+      className={cn(
+        "overflow-x-auto rounded-b-lg bg-black p-4 text-white",
+        className,
+      )}
+      {...props}
+    />
   ),
   code: function Code({ className, ...props }) {
     const isCodeBlock = useIsMarkdownCodeBlock();
     return (
       <code
-        className={cn(!isCodeBlock && "bg-muted rounded border font-semibold", className)}
+        className={cn(
+          !isCodeBlock && "bg-muted rounded border font-semibold",
+          className,
+        )}
         {...props}
       />
     );

src/components/assistant-ui/thread.tsx

@@ -2,8 +2,10 @@ import {
   ActionBarPrimitive,
   BranchPickerPrimitive,
   ComposerPrimitive,
+  getExternalStoreMessages,
   MessagePrimitive,
   ThreadPrimitive,
+  useMessage,
 } from "@assistant-ui/react";
 import type { FC } from "react";
 import {
@@ -17,12 +19,14 @@ import {
   SendHorizontalIcon,
 } from "lucide-react";
 import { cn } from "@/lib/utils";
+import { LoadExternalComponent } from "@langchain/langgraph-sdk/react-ui/client";
 
 import { Avatar, AvatarFallback } from "@/components/ui/avatar";
 import { Button } from "@/components/ui/button";
 import { MarkdownText } from "@/components/assistant-ui/markdown-text";
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
-
+import { Message } from "@langchain/langgraph-sdk";
+import { useStreamContext } from "@/providers/Stream";
 
 export const Thread: FC = () => {
   return (
@@ -78,9 +82,7 @@ const ThreadWelcome: FC = () => {
           <Avatar>
             <AvatarFallback>C</AvatarFallback>
           </Avatar>
-          <p className="mt-4 font-medium">
+          <p className="mt-4 font-medium">How can I help you today?</p>
-            How can I help you today?
-          </p>
         </div>
         <ThreadWelcomeSuggestions />
       </div>
@@ -93,12 +95,12 @@ const ThreadWelcomeSuggestions: FC = () => {
     <div className="mt-3 flex w-full items-stretch justify-center gap-4">
       <ThreadPrimitive.Suggestion
         className="hover:bg-muted/80 flex max-w-sm grow basis-0 flex-col items-center justify-center rounded-lg border p-3 transition-colors ease-in"
-        prompt="What is the weather in Tokyo?"
+        prompt="What's the current price of $APPL?"
         method="replace"
         autoSend
       >
         <span className="line-clamp-2 text-ellipsis text-sm font-semibold">
-          What is the weather in Tokyo?
+          What's the current price of $APPL?
         </span>
       </ThreadPrimitive.Suggestion>
       <ThreadPrimitive.Suggestion
@@ -108,7 +110,7 @@ const ThreadWelcomeSuggestions: FC = () => {
         autoSend
       >
         <span className="line-clamp-2 text-ellipsis text-sm font-semibold">
-          What is assistant-ui?
+          What's the weather like in San Francisco Today?
         </span>
       </ThreadPrimitive.Suggestion>
     </div>
@@ -205,13 +207,67 @@ const EditComposer: FC = () => {
   );
 };
 
+function CustomComponent({
+  message,
+  idx,
+  thread,
+}: {
+  message: Message;
+  idx: number;
+  thread: ReturnType<typeof useStreamContext>;
+}) {
+  const meta = thread.getMessagesMetadata(message, idx);
+  const seenState = meta?.firstSeenState;
+  const customComponent = seenState?.values.ui
+    .slice()
+    .reverse()
+    .find(
+      ({ additional_kwargs }) =>
+        additional_kwargs.run_id === seenState.metadata?.run_id,
+    );
+
+  return (
+    <div key={message.id}>
+      <pre>{JSON.stringify(message, null, 2)}</pre>
+      {customComponent && (
+        <LoadExternalComponent
+          assistantId="agent"
+          stream={thread}
+          message={customComponent}
+        />
+      )}
+    </div>
+  );
+}
+
 const AssistantMessage: FC = () => {
+  const thread = useStreamContext();
+  const assistantMsgs = useMessage((m) => {
+    const langchainMessage = getExternalStoreMessages<Message>(m);
+    return langchainMessage;
+  })?.[0];
+  let threadMsgIdx: number | undefined = undefined;
+  const threadMsg = thread.messages.find((m, idx) => {
+    if (m.id === assistantMsgs?.id) {
+      threadMsgIdx = idx;
+      return true;
+    }
+  });
+
   return (
     <MessagePrimitive.Root className="grid grid-cols-[auto_auto_1fr] grid-rows-[auto_1fr] relative w-full max-w-[var(--thread-max-width)] py-4">
       <Avatar className="col-start-1 row-span-full row-start-1 mr-4">
         <AvatarFallback>A</AvatarFallback>
       </Avatar>
 
+      {threadMsg && threadMsgIdx !== undefined && (
+        <CustomComponent
+          message={threadMsg}
+          idx={threadMsgIdx}
+          thread={thread}
+        />
+      )}
+
       <div className="text-foreground max-w-[calc(var(--thread-max-width)*0.8)] break-words leading-7 col-span-2 col-start-2 row-start-1 my-1.5">
         <MessagePrimitive.Content components={{ Text: MarkdownText }} />
       </div>
@@ -271,7 +327,10 @@ const BranchPicker: FC<BranchPickerPrimitive.Root.Props> = ({
   return (
     <BranchPickerPrimitive.Root
       hideWhenSingleBranch
-      className={cn("text-muted-foreground inline-flex items-center text-xs", className)}
+      className={cn(
+        "text-muted-foreground inline-flex items-center text-xs",
+        className,
+      )}
       {...rest}
     >
       <BranchPickerPrimitive.Previous asChild>

src/components/ui/avatar.tsx

@@ -1,7 +1,7 @@
-import * as React from "react"
+import * as React from "react";
-import * as AvatarPrimitive from "@radix-ui/react-avatar"
+import * as AvatarPrimitive from "@radix-ui/react-avatar";
 
-import { cn } from "@/lib/utils"
+import { cn } from "@/lib/utils";
 
 function Avatar({
   className,
@@ -12,11 +12,11 @@ function Avatar({
       data-slot="avatar"
       className={cn(
         "relative flex size-8 shrink-0 overflow-hidden rounded-full",
-        className
+        className,
       )}
       {...props}
     />
-  )
+  );
 }
 
 function AvatarImage({
@@ -29,7 +29,7 @@ function AvatarImage({
       className={cn("aspect-square size-full", className)}
       {...props}
     />
-  )
+  );
 }
 
 function AvatarFallback({
@@ -41,11 +41,11 @@ function AvatarFallback({
       data-slot="avatar-fallback"
       className={cn(
         "bg-muted flex size-full items-center justify-center rounded-full",
-        className
+        className,
       )}
       {...props}
     />
-  )
+  );
 }
 
-export { Avatar, AvatarImage, AvatarFallback }
+export { Avatar, AvatarImage, AvatarFallback };

src/components/ui/button.tsx

@@ -1,8 +1,8 @@
-import * as React from "react"
+import * as React from "react";
-import { Slot } from "@radix-ui/react-slot"
+import { Slot } from "@radix-ui/react-slot";
-import { cva, type VariantProps } from "class-variance-authority"
+import { cva, type VariantProps } from "class-variance-authority";
 
-import { cn } from "@/lib/utils"
+import { cn } from "@/lib/utils";
 
 const buttonVariants = cva(
   "inline-flex items-center justify-center gap-2 whitespace-nowrap rounded-md text-sm font-medium transition-[color,box-shadow] disabled:pointer-events-none disabled:opacity-50 [&_svg]:pointer-events-none [&_svg:not([class*='size-'])]:size-4 [&_svg]:shrink-0 outline-none focus-visible:border-ring focus-visible:ring-ring/50 focus-visible:ring-[3px] aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive",
@@ -31,8 +31,8 @@ const buttonVariants = cva(
       variant: "default",
       size: "default",
     },
-  }
+  },
-)
+);
 
 function Button({
   className,
@@ -42,9 +42,9 @@ function Button({
   ...props
 }: React.ComponentProps<"button"> &
   VariantProps<typeof buttonVariants> & {
-    asChild?: boolean
+    asChild?: boolean;
   }) {
-  const Comp = asChild ? Slot : "button"
+  const Comp = asChild ? Slot : "button";
 
   return (
     <Comp
@@ -52,7 +52,7 @@ function Button({
       className={cn(buttonVariants({ variant, size, className }))}
       {...props}
     />
-  )
+  );
 }
 
-export { Button, buttonVariants }
+export { Button, buttonVariants };

src/components/ui/tooltip.tsx

@@ -1,7 +1,7 @@
-import * as React from "react"
+import * as React from "react";
-import * as TooltipPrimitive from "@radix-ui/react-tooltip"
+import * as TooltipPrimitive from "@radix-ui/react-tooltip";
 
-import { cn } from "@/lib/utils"
+import { cn } from "@/lib/utils";
 
 function TooltipProvider({
   delayDuration = 0,
@@ -13,7 +13,7 @@ function TooltipProvider({
       delayDuration={delayDuration}
       {...props}
     />
-  )
+  );
 }
 
 function Tooltip({
@@ -23,13 +23,13 @@ function Tooltip({
     <TooltipProvider>
       <TooltipPrimitive.Root data-slot="tooltip" {...props} />
     </TooltipProvider>
-  )
+  );
 }
 
 function TooltipTrigger({
   ...props
 }: React.ComponentProps<typeof TooltipPrimitive.Trigger>) {
-  return <TooltipPrimitive.Trigger data-slot="tooltip-trigger" {...props} />
+  return <TooltipPrimitive.Trigger data-slot="tooltip-trigger" {...props} />;
 }
 
 function TooltipContent({
@@ -45,7 +45,7 @@ function TooltipContent({
         sideOffset={sideOffset}
         className={cn(
           "bg-primary text-primary-foreground animate-in fade-in-0 zoom-in-95 data-[state=closed]:animate-out data-[state=closed]:fade-out-0 data-[state=closed]:zoom-out-95 data-[side=bottom]:slide-in-from-top-2 data-[side=left]:slide-in-from-right-2 data-[side=right]:slide-in-from-left-2 data-[side=top]:slide-in-from-bottom-2 z-50 w-fit rounded-md px-3 py-1.5 text-xs text-balance",
-          className
+          className,
         )}
         {...props}
       >
@@ -53,7 +53,7 @@ function TooltipContent({
         <TooltipPrimitive.Arrow className="bg-primary fill-primary z-50 size-2.5 translate-y-[calc(-50%_-_2px)] rotate-45 rounded-[2px]" />
       </TooltipPrimitive.Content>
     </TooltipPrimitive.Portal>
-  )
+  );
 }
 
-export { Tooltip, TooltipTrigger, TooltipContent, TooltipProvider }
+export { Tooltip, TooltipTrigger, TooltipContent, TooltipProvider };

src/lib/utils.ts

@@ -1,6 +1,6 @@
-import { clsx, type ClassValue } from "clsx"
+import { clsx, type ClassValue } from "clsx";
-import { twMerge } from "tailwind-merge"
+import { twMerge } from "tailwind-merge";
 
 export function cn(...inputs: ClassValue[]) {
-  return twMerge(clsx(inputs))
+  return twMerge(clsx(inputs));
 }

src/main.tsx

@@ -2,9 +2,12 @@ import { createRoot } from "react-dom/client";
 import "./index.css";
 import App from "./App.tsx";
 import { RuntimeProvider } from "./providers/Runtime.tsx";
+import { StreamProvider } from "./providers/Stream.tsx";
 
-createRoot(document.getElementById("root")!).render((
+createRoot(document.getElementById("root")!).render(
-  <RuntimeProvider>
+  <StreamProvider>
-    <App />
+    <RuntimeProvider>
-  </RuntimeProvider>
+      <App />
-));
+    </RuntimeProvider>
+  </StreamProvider>,
+);

src/providers/Runtime.tsx

@@ -1,87 +1,40 @@
-import { useState, ReactNode } from "react";
+import { ReactNode } from "react";
 import {
   useExternalStoreRuntime,
-  ThreadMessageLike,
   AppendMessage,
   AssistantRuntimeProvider,
 } from "@assistant-ui/react";
-import { Message } from "@langchain/langgraph-sdk";
+import { HumanMessage } from "@langchain/langgraph-sdk";
-
+import { useStreamContext } from "./Stream";
-function langChainRoleToAssistantRole(role: Message["type"]): "system" | "assistant" | "user" {
+import { convertLangChainMessages } from "./convert-messages";
-  if (role === "ai") return "assistant";
-  if (role === "system") return "system";
-  if (["human", "tool", "function"].includes(role)) return "user";
-  throw new Error(`Unknown role: ${role}`);
-}
-
-function langChainContentToAssistantContent(content: Message["content"]): ThreadMessageLike["content"] {
-  if (!content) return [];
-  
-  if (typeof content === "string") return content;
-
-  if (typeof content === "object") {
-    if ("text" in content) {
-      return [{
-        type: "text",
-        text: content.text as string,
-      }]
-    }
-
-    if ("thinking" in content) {
-      return [{
-        type: "reasoning",
-        text: content.thinking as string,
-      }]
-    }
-  }
-
-  throw new Error(`Unknown content: ${content}`);
-}
- 
-const convertMessage = (message: Message): ThreadMessageLike => {
-  return {
-    role: langChainRoleToAssistantRole(message.type),
-    content: langChainContentToAssistantContent(message.content),
-  };
-};
-
-async function sleep(ms: number) {
-  return new Promise((resolve) => setTimeout(resolve, ms));
-}
 
 export function RuntimeProvider({
   children,
 }: Readonly<{
   children: ReactNode;
 }>) {
-  const [isRunning, setIsRunning] = useState(false);
+  const stream = useStreamContext();
-  const [messages, setMessages] = useState<Message[]>([]);
 
   const onNew = async (message: AppendMessage) => {
     if (message.content[0]?.type !== "text")
       throw new Error("Only text messages are supported");
 
     const input = message.content[0].text;
-    setMessages((currentConversation) => [
+    const humanMessage: HumanMessage = { type: "human", content: input };
-      ...currentConversation,
+    // TODO: I dont think I need to do this, since we're passing stream.messages into the state hook, and it should update when we call `submit`
-      { type: "human", content: input },
+    // setMessages((currentConversation) => [
-    ]);
+    //   ...currentConversation,
+    //   humanMessage,
+    // ]);
 
-    setIsRunning(true);
+    stream.submit({ messages: [humanMessage] });
-    // CALL API HERE
+    console.log("Sent message", humanMessage);
-    // const assistantMessage = await backendApi(input);
-    await sleep(2000);
-    setMessages((currentConversation) => [
-      ...currentConversation,
-      { type: "ai", content: [{ type: "text", text: "This is an assistant message" }] },
-    ]);
-    setIsRunning(false);
   };
 
   const runtime = useExternalStoreRuntime({
-    isRunning,
+    isRunning: stream.isLoading,
-    messages,
+    messages: stream.messages,
-    convertMessage,
+    convertMessage: convertLangChainMessages,
     onNew,
   });
 

src/providers/Stream.tsx

@@ -0,0 +1,58 @@
+import React, { createContext, useContext, ReactNode } from "react";
+import { useStream } from "@langchain/langgraph-sdk/react";
+import type { Message } from "@langchain/langgraph-sdk";
+import type {
+  UIMessage,
+  RemoveUIMessage,
+} from "@langchain/langgraph-sdk/react-ui/types";
+
+// Define the type for the context value
+type StreamContextType = ReturnType<
+  typeof useStream<
+    { messages: Message[]; ui: UIMessage[] },
+    {
+      messages?: Message[] | Message | string;
+      ui?: (UIMessage | RemoveUIMessage)[] | UIMessage | RemoveUIMessage;
+    },
+    UIMessage | RemoveUIMessage
+  >
+>;
+
+// Create the context with a default undefined value
+const StreamContext = createContext<StreamContextType | undefined>(undefined);
+
+// Create a provider component
+export const StreamProvider: React.FC<{ children: ReactNode }> = ({
+  children,
+}) => {
+  const streamValue = useStream<
+    { messages: Message[]; ui: UIMessage[] },
+    {
+      messages?: Message[] | Message | string;
+      ui?: (UIMessage | RemoveUIMessage)[] | UIMessage | RemoveUIMessage;
+    },
+    UIMessage | RemoveUIMessage
+  >({
+    apiUrl: "http://localhost:2024",
+    assistantId: "agent",
+  });
+
+  console.log("StreamProvider", streamValue);
+
+  return (
+    <StreamContext.Provider value={streamValue}>
+      {children}
+    </StreamContext.Provider>
+  );
+};
+
+// Create a custom hook to use the context
+export const useStreamContext = (): StreamContextType => {
+  const context = useContext(StreamContext);
+  if (context === undefined) {
+    throw new Error("useStreamContext must be used within a StreamProvider");
+  }
+  return context;
+};
+
+export default StreamContext;

src/providers/convert-messages.ts

@@ -0,0 +1,140 @@
+import { ThreadMessageLike, ToolCallContentPart } from "@assistant-ui/react";
+import { Message, AIMessage, ToolMessage } from "@langchain/langgraph-sdk";
+
+export const getMessageType = (message: Record<string, any>): string => {
+  if (Array.isArray(message.id)) {
+    const lastItem = message.id[message.id.length - 1];
+    if (lastItem.startsWith("HumanMessage")) {
+      return "human";
+    } else if (lastItem.startsWith("AIMessage")) {
+      return "ai";
+    } else if (lastItem.startsWith("ToolMessage")) {
+      return "tool";
+    } else if (
+      lastItem.startsWith("BaseMessage") ||
+      lastItem.startsWith("SystemMessage")
+    ) {
+      return "system";
+    }
+  }
+
+  if ("getType" in message && typeof message.getType === "function") {
+    return message.getType();
+  } else if ("_getType" in message && typeof message._getType === "function") {
+    return message._getType();
+  } else if ("type" in message) {
+    return message.type as string;
+  } else {
+    console.error(message);
+    throw new Error("Unsupported message type");
+  }
+};
+
+function getMessageContentOrThrow(message: unknown): string {
+  if (typeof message !== "object" || message === null) {
+    return "";
+  }
+
+  const castMsg = message as Record<string, any>;
+
+  if (
+    typeof castMsg?.content !== "string" &&
+    (!Array.isArray(castMsg.content) || castMsg.content[0]?.type !== "text") &&
+    (!castMsg.kwargs ||
+      !castMsg.kwargs?.content ||
+      typeof castMsg.kwargs?.content !== "string")
+  ) {
+    console.error(castMsg);
+    throw new Error("Only text messages are supported");
+  }
+
+  let content = "";
+  if (Array.isArray(castMsg.content) && castMsg.content[0]?.type === "text") {
+    content = castMsg.content[0].text;
+  } else if (typeof castMsg.content === "string") {
+    content = castMsg.content;
+  } else if (
+    castMsg?.kwargs &&
+    castMsg.kwargs?.content &&
+    typeof castMsg.kwargs?.content === "string"
+  ) {
+    content = castMsg.kwargs.content;
+  }
+
+  return content;
+}
+
+export function convertLangChainMessages(message: Message): ThreadMessageLike {
+  const content = getMessageContentOrThrow(message);
+
+  switch (getMessageType(message)) {
+    case "system":
+      return {
+        role: "system",
+        id: message.id,
+        content: [{ type: "text", text: content }],
+      };
+    case "human":
+      return {
+        role: "user",
+        id: message.id,
+        content: [{ type: "text", text: content }],
+        // ...(message.additional_kwargs
+        //   ? {
+        //       metadata: {
+        //         custom: {
+        //           ...message.additional_kwargs,
+        //         },
+        //       },
+        //     }
+        //   : {}),
+      };
+    case "ai":
+      const aiMsg = message as AIMessage;
+      const toolCallsContent: ToolCallContentPart[] = aiMsg.tool_calls?.length
+        ? aiMsg.tool_calls.map((tc) => ({
+            type: "tool-call" as const,
+            toolCallId: tc.id ?? "",
+            toolName: tc.name,
+            args: tc.args,
+            argsText: JSON.stringify(tc.args),
+          }))
+        : [];
+      return {
+        role: "assistant",
+        id: message.id,
+        content: [
+          ...toolCallsContent,
+          {
+            type: "text",
+            text: content,
+          },
+        ],
+        // ...(message.additional_kwargs
+        //   ? {
+        //       metadata: {
+        //         custom: {
+        //           ...message.additional_kwargs,
+        //         },
+        //       },
+        //     }
+        //   : {}),
+      };
+    case "tool":
+      const toolMsg = message as ToolMessage;
+      return {
+        role: "user",
+        content: [
+          {
+            type: "tool-call",
+            toolName: toolMsg.name ?? "ToolCall",
+            toolCallId: toolMsg.tool_call_id,
+            result: content,
+          },
+        ],
+      };
+    default:
+      console.error(message);
+      throw new Error(`Unsupported message type: ${getMessageType(message)}`);
+  }
+}

tailwind.config.js

@@ -1,57 +1,57 @@
 /** @type {import('tailwindcss').Config} */
 module.exports = {
-    darkMode: ["class"],
+  darkMode: ["class"],
-    content: ["./index.html", "./src/**/*.{ts,tsx,js,jsx}"],
+  content: ["./index.html", "./src/**/*.{ts,tsx,js,jsx}"],
   theme: {
-  	extend: {
+    extend: {
-  		borderRadius: {
+      borderRadius: {
-  			lg: 'var(--radius)',
+        lg: "var(--radius)",
-  			md: 'calc(var(--radius) - 2px)',
+        md: "calc(var(--radius) - 2px)",
-  			sm: 'calc(var(--radius) - 4px)'
+        sm: "calc(var(--radius) - 4px)",
-  		},
+      },
-  		colors: {
+      colors: {
-  			background: 'hsl(var(--background))',
+        background: "hsl(var(--background))",
-  			foreground: 'hsl(var(--foreground))',
+        foreground: "hsl(var(--foreground))",
-  			card: {
+        card: {
-  				DEFAULT: 'hsl(var(--card))',
+          DEFAULT: "hsl(var(--card))",
-  				foreground: 'hsl(var(--card-foreground))'
+          foreground: "hsl(var(--card-foreground))",
-  			},
+        },
-  			popover: {
+        popover: {
-  				DEFAULT: 'hsl(var(--popover))',
+          DEFAULT: "hsl(var(--popover))",
-  				foreground: 'hsl(var(--popover-foreground))'
+          foreground: "hsl(var(--popover-foreground))",
-  			},
+        },
-  			primary: {
+        primary: {
-  				DEFAULT: 'hsl(var(--primary))',
+          DEFAULT: "hsl(var(--primary))",
-  				foreground: 'hsl(var(--primary-foreground))'
+          foreground: "hsl(var(--primary-foreground))",
-  			},
+        },
-  			secondary: {
+        secondary: {
-  				DEFAULT: 'hsl(var(--secondary))',
+          DEFAULT: "hsl(var(--secondary))",
-  				foreground: 'hsl(var(--secondary-foreground))'
+          foreground: "hsl(var(--secondary-foreground))",
-  			},
+        },
-  			muted: {
+        muted: {
-  				DEFAULT: 'hsl(var(--muted))',
+          DEFAULT: "hsl(var(--muted))",
-  				foreground: 'hsl(var(--muted-foreground))'
+          foreground: "hsl(var(--muted-foreground))",
-  			},
+        },
-  			accent: {
+        accent: {
-  				DEFAULT: 'hsl(var(--accent))',
+          DEFAULT: "hsl(var(--accent))",
-  				foreground: 'hsl(var(--accent-foreground))'
+          foreground: "hsl(var(--accent-foreground))",
-  			},
+        },
-  			destructive: {
+        destructive: {
-  				DEFAULT: 'hsl(var(--destructive))',
+          DEFAULT: "hsl(var(--destructive))",
-  				foreground: 'hsl(var(--destructive-foreground))'
+          foreground: "hsl(var(--destructive-foreground))",
-  			},
+        },
-  			border: 'hsl(var(--border))',
+        border: "hsl(var(--border))",
-  			input: 'hsl(var(--input))',
+        input: "hsl(var(--input))",
-  			ring: 'hsl(var(--ring))',
+        ring: "hsl(var(--ring))",
-  			chart: {
+        chart: {
-  				'1': 'hsl(var(--chart-1))',
+          1: "hsl(var(--chart-1))",
-  				'2': 'hsl(var(--chart-2))',
+          2: "hsl(var(--chart-2))",
-  				'3': 'hsl(var(--chart-3))',
+          3: "hsl(var(--chart-3))",
-  				'4': 'hsl(var(--chart-4))',
+          4: "hsl(var(--chart-4))",
-  				'5': 'hsl(var(--chart-5))'
+          5: "hsl(var(--chart-5))",
-  			}
+        },
-  		}
+      },
-  	}
+    },
   },
   plugins: [require("tailwindcss-animate")],
-}
+};

vite.config.ts

@@ -1,7 +1,7 @@
-import path from "path"
+import path from "path";
-import { defineConfig } from 'vite'
+import { defineConfig } from "vite";
-import react from '@vitejs/plugin-react'
+import react from "@vitejs/plugin-react";
-import tailwindcss from '@tailwindcss/vite'
+import tailwindcss from "@tailwindcss/vite";
 
 // https://vite.dev/config/
 export default defineConfig({
@@ -11,4 +11,4 @@ export default defineConfig({
       "@": path.resolve(__dirname, "./src"),
     },
   },
-})
+});