Merge pull request #1 from langchain-ai/brace/assistant-ui

Brace/assistant UI
2025-03-03 12:51:32 -08:00
parent 7448b33075 49e54bc64c
commit 5256efb23f
41 changed files with 7050 additions and 2034 deletions
--- a/agent/agent.tsx
+++ b/agent/agent.tsx
@@ -1,89 +1,96 @@
-/* eslint-disable @typescript-eslint/no-explicit-any */
-import {
-  StateGraph,
-  MessagesAnnotation,
-  START,
-  Annotation,
-} from "@langchain/langgraph";
-import { SystemMessage } from "@langchain/core/messages";
+import { StateGraph, START, END } from "@langchain/langgraph";
+import { ChatGoogleGenerativeAI } from "@langchain/google-genai";
+import { z } from "zod";
+import { GenerativeUIAnnotation, GenerativeUIState } from "./types";
+import { stockbrokerGraph } from "./stockbroker";
 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 { z, ZodTypeAny } from "zod";

-// const llm = new ChatOllama({ model: "deepseek-r1" });
-const llm = new ChatOpenAI({ model: "gpt-4o-mini", temperature: 0 });
+async function router(
+  state: GenerativeUIState,
+): Promise<Partial<GenerativeUIState>> {
+  const routerDescription = `The route to take based on the user's input.
+- stockbroker: can fetch the price of a ticker, purchase/sell a ticker, or get the user's portfolio
+- weather: can fetch the current weather conditions for a location
+- generalInput: handles all other cases where the above tools don't apply
+`;
+  const routerSchema = z.object({
+    route: z
+      .enum(["stockbroker", "weather", "generalInput"])
+      .describe(routerDescription),
+  });
+  const routerTool = {
+    name: "router",
+    description: "A tool to route the user's query to the appropriate tool.",
+    schema: routerSchema,
+  };

-interface ToolCall {
-  name: string;
-  args: Record<string, any>;
-  id?: string;
-  type?: "tool_call";
-}
-
-function findToolCall<Name extends string>(name: Name) {
-  return <Args extends ZodTypeAny>(
-    x: ToolCall
-  ): x is { name: Name; args: z.infer<Args> } => x.name === name;
-}
-
-const builder = new StateGraph(
-  Annotation.Root({
-    messages: MessagesAnnotation.spec["messages"],
-    ui: Annotation({ default: () => [], reducer: uiMessageReducer }),
-    timestamp: Annotation<number>,
+  const llm = new ChatGoogleGenerativeAI({
+    model: "gemini-2.0-flash",
+    temperature: 0,
  })
-)
-  .addNode("agent", async (state, config) => {
-    const ui = typedUi<typeof ComponentMap>(config);
+    .bindTools([routerTool], { tool_choice: "router" })
+    .withConfig({ tags: ["langsmith:nostream"] });

-    // const result = ui.interrupt("react-component", {
-    //   instruction: "Hello world",
-    // });
+  const prompt = `You're a highly helpful AI assistant, tasked with routing the user's query to the appropriate tool.
+You should analyze the user's input, and choose the appropriate tool to use.`;

-    // // 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([
-        {
-          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"
-        ),
-        ...state.messages,
-      ]);
-
-    const stockbrokerToolCall = message.tool_calls?.find(
-      findToolCall("stockbroker")<typeof stockbrokerSchema>
+  const recentHumanMessage = state.messages.findLast(
+    (m) => m.getType() === "human",
  );

-    if (stockbrokerToolCall) {
-      const instruction = `The stock price of ${
-        stockbrokerToolCall.args.company
-      } is ${Math.random() * 100}`;
-
-      ui.write("react-component", { instruction, logo: "hey" });
+  if (!recentHumanMessage) {
+    throw new Error("No human message found in state");
  }

-    return { messages: message, ui: ui.collect, timestamp: Date.now() };
+  const response = await llm.invoke([
+    { role: "system", content: prompt },
+    recentHumanMessage,
+  ]);
+
+  const toolCall = response.tool_calls?.[0]?.args as
+    | z.infer<typeof routerSchema>
+    | undefined;
+  if (!toolCall) {
+    throw new Error("No tool call found in response");
+  }
+
+  return {
+    next: toolCall.route,
+  };
+}
+
+function handleRoute(
+  state: GenerativeUIState,
+): "stockbroker" | "weather" | "generalInput" {
+  return state.next;
+}
+
+async function handleGeneralInput(state: GenerativeUIState) {
+  const llm = new ChatOpenAI({ model: "gpt-4o-mini", temperature: 0 });
+  const response = await llm.invoke(state.messages);
+
+  return {
+    messages: [response],
+  };
+}
+
+const builder = new StateGraph(GenerativeUIAnnotation)
+  .addNode("router", router)
+  .addNode("stockbroker", stockbrokerGraph)
+  .addNode("weather", () => {
+    throw new Error("Weather not implemented");
  })
-  .addEdge(START, "agent");
+  .addNode("generalInput", handleGeneralInput)
+
+  .addConditionalEdges("router", handleRoute, [
+    "stockbroker",
+    "weather",
+    "generalInput",
+  ])
+  .addEdge(START, "router")
+  .addEdge("stockbroker", END)
+  .addEdge("weather", END)
+  .addEdge("generalInput", END);

 export const graph = builder.compile();
-
-// event handler of evetns ˇtypes)
-// event handler for specific node -> handle node
-
-// TODO:
-// - Send run ID & additional metadata for the client to properly use messages (maybe we even have a config)
-// - Store that run ID in messages
+graph.name = "Generative UI Agent";
--- a/agent/find-tool-call.ts
+++ b/agent/find-tool-call.ts
@@ -0,0 +1,14 @@
+import { z, ZodTypeAny } from "zod";
+
+interface ToolCall {
+  name: string;
+  args: Record<string, any>;
+  id?: string;
+  type?: "tool_call";
+}
+
+export function findToolCall<Name extends string>(name: Name) {
+  return <Args extends ZodTypeAny>(
+    x: ToolCall,
+  ): x is { name: Name; args: z.infer<Args> } => x.name === name;
+}
--- a/agent/stockbroker/index.tsx
+++ b/agent/stockbroker/index.tsx
@@ -0,0 +1,11 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+import { StateGraph, START } from "@langchain/langgraph";
+import { StockbrokerAnnotation } from "./types";
+import { callTools } from "./nodes/tools";
+
+const builder = new StateGraph(StockbrokerAnnotation)
+  .addNode("agent", callTools)
+  .addEdge(START, "agent");
+
+export const stockbrokerGraph = builder.compile();
+stockbrokerGraph.name = "Stockbroker";
--- a/agent/stockbroker/nodes/tools.tsx
+++ b/agent/stockbroker/nodes/tools.tsx
@@ -0,0 +1,72 @@
+import { StockbrokerState } from "../types";
+import { ChatOpenAI } from "@langchain/openai";
+import { typedUi } from "@langchain/langgraph-sdk/react-ui/server";
+import type ComponentMap from "../../uis/index";
+import { z } from "zod";
+import { LangGraphRunnableConfig } from "@langchain/langgraph";
+import { findToolCall } from "../../find-tool-call";
+
+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: "stock-price",
+    description: "A tool to get the stock price of a company",
+    schema: getStockPriceSchema,
+  },
+  {
+    name: "portfolio",
+    description:
+      "A tool to get the user's portfolio details. Only call this tool if the user requests their portfolio details.",
+    schema: getPortfolioSchema,
+  },
+];
+
+export async function callTools(
+  state: StockbrokerState,
+  config: LangGraphRunnableConfig,
+): Promise<Partial<StockbrokerState>> {
+  const ui = typedUi<typeof ComponentMap>(config);
+
+  const message = await llm.bindTools(STOCKBROKER_TOOLS).invoke([
+    {
+      role: "system",
+      content:
+        "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("stock-price")<typeof getStockPriceSchema>,
+  );
+  const portfolioToolCall = message.tool_calls?.find(
+    findToolCall("portfolio")<typeof getStockPriceSchema>,
+  );
+
+  if (stockbrokerToolCall) {
+    const instruction = `The stock price of ${
+      stockbrokerToolCall.args.ticker
+    } is ${Math.random() * 100}`;
+
+    ui.write("stock-price", { instruction, logo: "hey" });
+  }
+
+  if (portfolioToolCall) {
+    ui.write("portfolio", {});
+  }
+
+  return {
+    messages: [message],
+    // TODO: Fix the ui return type.
+    ui: ui.collect as any[],
+    timestamp: Date.now(),
+  };
+}
--- a/agent/stockbroker/types.ts
+++ b/agent/stockbroker/types.ts
@@ -0,0 +1,11 @@
+import { Annotation } from "@langchain/langgraph";
+import { GenerativeUIAnnotation } from "../types";
+
+export const StockbrokerAnnotation = Annotation.Root({
+  messages: GenerativeUIAnnotation.spec.messages,
+  ui: GenerativeUIAnnotation.spec.ui,
+  timestamp: GenerativeUIAnnotation.spec.timestamp,
+  next: Annotation<"stockbroker" | "weather">(),
+});
+
+export type StockbrokerState = typeof StockbrokerAnnotation.State;
--- a/agent/types.ts
+++ b/agent/types.ts
@@ -0,0 +1,11 @@
+import { MessagesAnnotation, Annotation } from "@langchain/langgraph";
+import { uiMessageReducer } from "@langchain/langgraph-sdk/react-ui/types";
+
+export const GenerativeUIAnnotation = Annotation.Root({
+  messages: MessagesAnnotation.spec["messages"],
+  ui: Annotation({ default: () => [], reducer: uiMessageReducer }),
+  timestamp: Annotation<number>,
+  next: Annotation<"stockbroker" | "weather" | "generalInput">(),
+});
+
+export type GenerativeUIState = typeof GenerativeUIAnnotation.State;
--- a/agent/uis/index.tsx
+++ b/agent/uis/index.tsx
@@ -0,0 +1,8 @@
+import StockPrice from "./stock-price";
+import PortfolioView from "./portfolio-view";
+
+const ComponentMap = {
+  "stock-price": StockPrice,
+  portfolio: PortfolioView,
+} as const;
+export default ComponentMap;
--- a/agent/uis/portfolio-view/index.css
+++ b/agent/uis/portfolio-view/index.css
--- a/agent/uis/portfolio-view/index.tsx
+++ b/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>
+  );
+}
--- a/agent/uis/stock-price/index.css
+++ b/agent/uis/stock-price/index.css
@@ -0,0 +1 @@
+@import "tailwindcss";
--- a/agent/uis/stock-price/index.tsx
+++ b/agent/uis/stock-price/index.tsx
@@ -1,23 +1,28 @@
-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
-  const thread = useStream<{ messages: Message[] }, { messages: Message[] }>({
+  const thread = useStream<{ messages: Message[] }>({
    assistantId: "assistant_123",
    apiUrl: "http://localhost:3123",
  });

-  const aiTool = thread.messages
+  const messagesCopy = thread.messages;
+
+  const aiTool = messagesCopy
    .slice()
    .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 +57,3 @@ function ReactComponent(props: { instruction: string; logo: string }) {
    </div>
  );
 }
-
-const ComponentMap = { "react-component": ReactComponent } as const;
-export default ComponentMap;
--- a/components.json
+++ b/components.json
@@ -0,0 +1,21 @@
+{
+  "$schema": "https://ui.shadcn.com/schema.json",
+  "style": "new-york",
+  "rsc": false,
+  "tsx": true,
+  "tailwind": {
+    "config": "tailwind.config.js",
+    "css": "src/index.css",
+    "baseColor": "neutral",
+    "cssVariables": true,
+    "prefix": ""
+  },
+  "aliases": {
+    "components": "@/components",
+    "utils": "@/lib/utils",
+    "ui": "@/components/ui",
+    "lib": "@/lib",
+    "hooks": "@/hooks"
+  },
+  "iconLibrary": "lucide"
+}
--- a/eslint.config.js
+++ b/eslint.config.js
@@ -1,28 +1,28 @@
-import js from '@eslint/js'
-import globals from 'globals'
-import reactHooks from 'eslint-plugin-react-hooks'
-import reactRefresh from 'eslint-plugin-react-refresh'
-import tseslint from 'typescript-eslint'
+import js from "@eslint/js";
+import globals from "globals";
+import reactHooks from "eslint-plugin-react-hooks";
+import reactRefresh from "eslint-plugin-react-refresh";
+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-refresh': reactRefresh,
+      "react-hooks": reactHooks,
+      "react-refresh": reactRefresh,
    },
    rules: {
      ...reactHooks.configs.recommended.rules,
-      'react-refresh/only-export-components': [
-        'warn',
+      "react-refresh/only-export-components": [
+        "warn",
        { allowConstantExport: true },
      ],
    },
  },
-)
+);
--- a/index.html
+++ b/index.html
@@ -5,6 +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" />
  </head>
  <body>
    <div id="root"></div>
--- a/langgraph.json
+++ b/langgraph.json
@@ -3,7 +3,7 @@
    "agent": "./agent/agent.tsx:graph"
  },
  "ui": {
-    "agent": "./agent/ui.tsx"
+    "agent": "./agent/uis/index.tsx"
  },
  "_INTERNAL_docker_tag": "20",
  "env": ".env"
--- a/package.json
+++ b/package.json
@@ -8,39 +8,60 @@
    "dev": "vite",
    "build": "tsc -b && vite build",
    "lint": "eslint .",
+    "format": "prettier --write .",
    "preview": "vite preview"
  },
  "dependencies": {
-    "react": "^19.0.0",
-    "react-dom": "^19.0.0",
-    "@langchain/core": "^0.3.40",
-    "@langchain/langgraph": "^0.2.46",
-    "@langchain/openai": "^0.4.4",
+    "@assistant-ui/react": "^0.8.0",
+    "@assistant-ui/react-markdown": "^0.8.0",
+    "@langchain/core": "^0.3.41",
+    "@langchain/google-genai": "^0.1.10",
+    "@langchain/langgraph": "^0.2.49",
    "@langchain/langgraph-api": "*",
    "@langchain/langgraph-cli": "*",
    "@langchain/langgraph-sdk": "*",
+    "@langchain/openai": "^0.4.4",
+    "@radix-ui/react-avatar": "^1.1.3",
+    "@radix-ui/react-slot": "^1.1.2",
+    "@radix-ui/react-tooltip": "^1.1.8",
+    "@tailwindcss/postcss": "^4.0.9",
+    "@tailwindcss/vite": "^4.0.9",
+    "class-variance-authority": "^0.7.1",
+    "clsx": "^2.1.1",
    "esbuild": "^0.25.0",
    "esbuild-plugin-tailwindcss": "^2.0.1",
-    "tailwindcss": "^4.0.6",
+    "framer-motion": "^12.4.9",
+    "lucide-react": "^0.476.0",
+    "prettier": "^3.5.2",
+    "react": "^19.0.0",
+    "react-dom": "^19.0.0",
+    "react-markdown": "^10.0.1",
+    "remark-gfm": "^4.0.1",
+    "tailwind-merge": "^3.0.2",
+    "tailwindcss-animate": "^1.0.7",
    "uuid": "^11.0.5",
    "zod": "^3.24.2"
  },
  "resolutions": {
-    "@langchain/langgraph-api": "http://localhost:3123/17/@langchain/langgraph-api",
-    "@langchain/langgraph-cli": "http://localhost:3123/17/@langchain/langgraph-cli",
-    "@langchain/langgraph-sdk": "http://localhost:3123/17/@langchain/langgraph-sdk"
+    "@langchain/langgraph-api": "next",
+    "@langchain/langgraph-cli": "next",
+    "@langchain/langgraph-sdk": "next"
  },
  "devDependencies": {
    "@eslint/js": "^9.19.0",
+    "@types/node": "^22.13.5",
    "@types/react": "^19.0.8",
    "@types/react-dom": "^19.0.3",
    "@vitejs/plugin-react": "^4.3.4",
+    "autoprefixer": "^10.4.20",
    "eslint": "^9.19.0",
    "eslint-plugin-react-hooks": "^5.0.0",
    "eslint-plugin-react-refresh": "^0.4.18",
    "globals": "^15.14.0",
+    "tailwindcss": "^4.0.6",
    "typescript": "~5.7.2",
    "typescript-eslint": "^8.22.0",
    "vite": "^6.1.0"
-  }
+  },
+  "packageManager": "pnpm@10.5.1+sha512.c424c076bd25c1a5b188c37bb1ca56cc1e136fbf530d98bcb3289982a08fd25527b8c9c4ec113be5e3393c39af04521dd647bcf1d0801eaf8ac6a7b14da313af"
 }
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
--- a/src/App.tsx
+++ b/src/App.tsx
@@ -1,74 +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/thread";

 function App() {
-  const thread = useStream<
-    { messages: Message[]; ui: UIMessage[] },
-    {
-      messages?: Message[] | Message | string;
-      ui?: (UIMessage | RemoveUIMessage)[] | UIMessage | RemoveUIMessage;
-    },
-    UIMessage | RemoveUIMessage
-  >({
-    apiUrl: "http://localhost:2024",
-    assistantId: "agent",
-  });
-
  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>
+      <Thread />
    </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;
--- a/src/components/icons/langgraph.tsx
+++ b/src/components/icons/langgraph.tsx
--- a/src/components/thread/index.tsx
+++ b/src/components/thread/index.tsx
@@ -0,0 +1,166 @@
+import { v4 as uuidv4 } from "uuid";
+import { useEffect, useRef } from "react";
+import { cn } from "@/lib/utils";
+import { useStreamContext } from "@/providers/Stream";
+import { useState, FormEvent } from "react";
+import { Input } from "../ui/input";
+import { Button } from "../ui/button";
+import { Message } from "@langchain/langgraph-sdk";
+import { AssistantMessage, AssistantMessageLoading } from "./messages/ai";
+import { HumanMessage } from "./messages/human";
+import {
+  DO_NOT_RENDER_ID_PREFIX,
+  ensureToolCallsHaveResponses,
+} from "@/lib/ensure-tool-responses";
+import { LangGraphLogoSVG } from "../icons/langgraph";
+
+// const dummyMessages = [
+//   { type: "human", content: "Hi! What can you do?" },
+//   {
+//     type: "ai",
+//     content: `Hello! I can assist you with a variety of tasks, including:
+
+// 1. **Answering Questions**: I can provide information on a wide range of topics, from science and history to technology and culture.
+// 2. **Writing Assistance**: I can help you draft emails, essays, reports, and creative writing pieces.
+// 3. **Learning Support**: I can explain concepts, help with homework, and provide study tips.
+// 4. **Language Help**: I can assist with translations, grammar, and vocabulary in multiple languages.
+// 5. **Recommendations**: I can suggest books, movies, recipes, and more based on your interests.
+// 6. **General Advice**: I can offer tips on various subjects, including productivity, wellness, and personal development.
+
+// If you have something specific in mind, feel free to ask!`,
+//   },
+// ];
+
+function Title({ className }: { className?: string }) {
+  return (
+    <div className={cn("flex gap-2 items-center", className)}>
+      <LangGraphLogoSVG width={32} height={32} />
+      <h1 className="text-xl font-medium">LangGraph Chat</h1>
+    </div>
+  );
+}
+
+export function Thread() {
+  const [input, setInput] = useState("");
+  const [firstTokenReceived, setFirstTokenReceived] = useState(false);
+  const stream = useStreamContext();
+  // const messages = [...dummyMessages, ...stream.messages];
+  const messages = stream.messages;
+  const isLoading = stream.isLoading;
+  const prevMessageLength = useRef(0);
+
+  useEffect(() => {
+    if (
+      messages.length !== prevMessageLength.current &&
+      messages?.length &&
+      messages[messages.length - 1].type === "ai"
+    ) {
+      setFirstTokenReceived(true);
+      prevMessageLength.current = messages.length;
+    }
+  }, [messages]);
+
+  const handleSubmit = (e: FormEvent) => {
+    e.preventDefault();
+    if (!input.trim() || isLoading) return;
+    setFirstTokenReceived(false);
+
+    const newHumanMessage: Message = {
+      id: uuidv4(),
+      type: "human",
+      content: input,
+    };
+
+    stream.submit(
+      {
+        messages: [
+          ...ensureToolCallsHaveResponses(stream.messages),
+          newHumanMessage,
+        ],
+      },
+      {
+        streamMode: ["values"],
+      },
+    );
+
+    setInput("");
+  };
+
+  const chatStarted = isLoading || messages.length > 0;
+  const renderMessages = messages.filter(
+    (m) => !m.id?.startsWith(DO_NOT_RENDER_ID_PREFIX),
+  );
+
+  return (
+    <div
+      className={cn(
+        "flex flex-col w-full h-full",
+        chatStarted ? "relative" : "",
+      )}
+    >
+      <div className={cn("flex-1 px-4", chatStarted ? "pb-28" : "mt-64")}>
+        {!chatStarted && (
+          <div className="flex justify-center">
+            <Title className="mb-12" />
+          </div>
+        )}
+        {chatStarted && (
+          <div className="hidden md:flex absolute top-4 right-4">
+            <Title />
+          </div>
+        )}
+
+        <div
+          className={cn(
+            "flex flex-col gap-4 max-w-4xl w-full mx-auto mt-12 overflow-y-auto",
+            !chatStarted && "hidden",
+          )}
+        >
+          {renderMessages.map((message, index) =>
+            message.type === "human" ? (
+              <HumanMessage
+                key={"id" in message ? message.id : `${message.type}-${index}`}
+                message={message as Message}
+                isLoading={isLoading}
+              />
+            ) : (
+              <AssistantMessage
+                key={"id" in message ? message.id : `${message.type}-${index}`}
+                message={message as Message}
+                isLoading={isLoading}
+              />
+            ),
+          )}
+          {isLoading && !firstTokenReceived && <AssistantMessageLoading />}
+        </div>
+      </div>
+
+      <div
+        className={cn(
+          "bg-white rounded-2xl border-[1px] border-gray-200 shadow-md p-3 mx-auto w-full max-w-5xl",
+          chatStarted ? "fixed bottom-6 left-0 right-0" : "",
+        )}
+      >
+        <form
+          onSubmit={handleSubmit}
+          className="flex w-full gap-2 max-w-5xl mx-auto"
+        >
+          <Input
+            type="text"
+            value={input}
+            onChange={(e) => setInput(e.target.value)}
+            placeholder="Type your message..."
+            className="p-5 border-[0px] shadow-none ring-0 outline-none focus:outline-none focus:ring-0"
+          />
+          <Button
+            type="submit"
+            className="p-5"
+            disabled={isLoading || !input.trim()}
+          >
+            Send
+          </Button>
+        </form>
+      </div>
+    </div>
+  );
+}
--- a/src/components/thread/markdown-text.tsx
+++ b/src/components/thread/markdown-text.tsx
@@ -0,0 +1,217 @@
+"use client";
+
+import "@assistant-ui/react-markdown/styles/dot.css";
+
+import {
+  CodeHeaderProps,
+  unstable_memoizeMarkdownComponents as memoizeMarkdownComponents,
+  useIsMarkdownCodeBlock,
+} from "@assistant-ui/react-markdown";
+import ReactMarkdown from "react-markdown";
+import remarkGfm from "remark-gfm";
+import { FC, memo, useState } from "react";
+import { CheckIcon, CopyIcon } from "lucide-react";
+
+import { TooltipIconButton } from "@/components/thread/tooltip-icon-button";
+import { cn } from "@/lib/utils";
+
+const MarkdownTextImpl = ({ children }: { children: string }) => {
+  return (
+    <ReactMarkdown remarkPlugins={[remarkGfm]} components={defaultComponents}>
+      {children}
+    </ReactMarkdown>
+  );
+};
+
+export const MarkdownText = memo(MarkdownTextImpl);
+
+const CodeHeader: FC<CodeHeaderProps> = ({ language, code }) => {
+  const { isCopied, copyToClipboard } = useCopyToClipboard();
+  const onCopy = () => {
+    if (!code || isCopied) return;
+    copyToClipboard(code);
+  };
+
+  return (
+    <div className="flex items-center justify-between gap-4 rounded-t-lg bg-zinc-900 px-4 py-2 text-sm font-semibold text-white">
+      <span className="lowercase [&>span]:text-xs">{language}</span>
+      <TooltipIconButton tooltip="Copy" onClick={onCopy}>
+        {!isCopied && <CopyIcon />}
+        {isCopied && <CheckIcon />}
+      </TooltipIconButton>
+    </div>
+  );
+};
+
+const useCopyToClipboard = ({
+  copiedDuration = 3000,
+}: {
+  copiedDuration?: number;
+} = {}) => {
+  const [isCopied, setIsCopied] = useState<boolean>(false);
+
+  const copyToClipboard = (value: string) => {
+    if (!value) return;
+
+    navigator.clipboard.writeText(value).then(() => {
+      setIsCopied(true);
+      setTimeout(() => setIsCopied(false), copiedDuration);
+    });
+  };
+
+  return { isCopied, copyToClipboard };
+};
+
+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}
+    />
+  ),
+  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}
+    />
+  ),
+  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}
+    />
+  ),
+  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}
+    />
+  ),
+  h5: ({ 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}
+    />
+  ),
+  p: ({ 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}
+    />
+  ),
+  blockquote: ({ 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}
+    />
+  ),
+  ol: ({ 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}
+    />
+  ),
+  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}
+    />
+  ),
+  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}
+    />
+  ),
+  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}
+    />
+  ),
+  sup: ({ 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}
+    />
+  ),
+  code: function Code({ className, ...props }) {
+    const isCodeBlock = useIsMarkdownCodeBlock();
+    return (
+      <code
+        className={cn(
+          !isCodeBlock && "bg-muted rounded border font-semibold",
+          className,
+        )}
+        {...props}
+      />
+    );
+  },
+  CodeHeader,
+});
--- a/src/components/thread/messages/ai.tsx
+++ b/src/components/thread/messages/ai.tsx
@@ -0,0 +1,104 @@
+import { useStreamContext } from "@/providers/Stream";
+import { Message } from "@langchain/langgraph-sdk";
+import { getContentString } from "../utils";
+import { BranchSwitcher, CommandBar } from "./shared";
+import { Avatar, AvatarFallback } from "@/components/ui/avatar";
+import { MarkdownText } from "../markdown-text";
+import { LoadExternalComponent } from "@langchain/langgraph-sdk/react-ui/client";
+
+function CustomComponent({
+  message,
+  thread,
+}: {
+  message: Message;
+  thread: ReturnType<typeof useStreamContext>;
+}) {
+  const meta = thread.getMessagesMetadata(message);
+  const seenState = meta?.firstSeenState;
+  const customComponent = seenState?.values.ui
+    .slice()
+    .reverse()
+    .find(
+      ({ additional_kwargs }) =>
+        additional_kwargs.run_id === seenState.metadata?.run_id,
+    );
+
+  if (!customComponent) {
+    return null;
+  }
+
+  return (
+    <div key={message.id}>
+      {customComponent && (
+        <LoadExternalComponent
+          assistantId="agent"
+          stream={thread}
+          message={customComponent}
+        />
+      )}
+    </div>
+  );
+}
+
+export function AssistantMessage({
+  message,
+  isLoading,
+}: {
+  message: Message;
+  isLoading: boolean;
+}) {
+  const thread = useStreamContext();
+  const meta = thread.getMessagesMetadata(message);
+  const parentCheckpoint = meta?.firstSeenState?.parent_checkpoint;
+
+  const contentString = getContentString(message.content);
+
+  const handleRegenerate = () => {
+    thread.submit(undefined, { checkpoint: parentCheckpoint });
+  };
+
+  return (
+    <div className="flex items-start mr-auto gap-2 group">
+      <Avatar>
+        <AvatarFallback>A</AvatarFallback>
+      </Avatar>
+      <div className="flex flex-col gap-2">
+        <CustomComponent message={message} thread={thread} />
+        {contentString.length > 0 && (
+          <div className="rounded-2xl bg-muted px-4 py-2">
+            <MarkdownText>{contentString}</MarkdownText>
+          </div>
+        )}
+        <div className="flex gap-2 items-center mr-auto opacity-0 group-hover:opacity-100 transition-opacity">
+          <BranchSwitcher
+            branch={meta?.branch}
+            branchOptions={meta?.branchOptions}
+            onSelect={(branch) => thread.setBranch(branch)}
+            isLoading={isLoading}
+          />
+          <CommandBar
+            content={contentString}
+            isLoading={isLoading}
+            isAiMessage={true}
+            handleRegenerate={handleRegenerate}
+          />
+        </div>
+      </div>
+    </div>
+  );
+}
+
+export function AssistantMessageLoading() {
+  return (
+    <div className="flex items-start mr-auto gap-2">
+      <Avatar>
+        <AvatarFallback>A</AvatarFallback>
+      </Avatar>
+      <div className="flex items-center gap-1 rounded-2xl bg-muted px-4 py-2 h-8">
+        <div className="w-1.5 h-1.5 rounded-full bg-foreground/50 animate-[pulse_1.5s_ease-in-out_infinite]"></div>
+        <div className="w-1.5 h-1.5 rounded-full bg-foreground/50 animate-[pulse_1.5s_ease-in-out_0.5s_infinite]"></div>
+        <div className="w-1.5 h-1.5 rounded-full bg-foreground/50 animate-[pulse_1.5s_ease-in-out_1s_infinite]"></div>
+      </div>
+    </div>
+  );
+}
--- a/src/components/thread/messages/human.tsx
+++ b/src/components/thread/messages/human.tsx
@@ -0,0 +1,107 @@
+import { useStreamContext } from "@/providers/Stream";
+import { Message } from "@langchain/langgraph-sdk";
+import { useState } from "react";
+import { getContentString } from "../utils";
+import { cn } from "@/lib/utils";
+import { Textarea } from "@/components/ui/textarea";
+import { BranchSwitcher, CommandBar } from "./shared";
+
+function EditableContent({
+  value,
+  setValue,
+  onSubmit,
+}: {
+  value: string;
+  setValue: React.Dispatch<React.SetStateAction<string>>;
+  onSubmit: () => void;
+}) {
+  const handleKeyDown = (e: React.KeyboardEvent) => {
+    if ((e.metaKey || e.ctrlKey) && e.key === "Enter") {
+      e.preventDefault();
+      onSubmit();
+    }
+  };
+
+  return (
+    <Textarea
+      value={value}
+      onChange={(e) => setValue(e.target.value)}
+      onKeyDown={handleKeyDown}
+    />
+  );
+}
+
+export function HumanMessage({
+  message,
+  isLoading,
+}: {
+  message: Message;
+  isLoading: boolean;
+}) {
+  const thread = useStreamContext();
+  const meta = thread.getMessagesMetadata(message);
+  const parentCheckpoint = meta?.firstSeenState?.parent_checkpoint;
+
+  const [isEditing, setIsEditing] = useState(false);
+  const [value, setValue] = useState("");
+  const contentString = getContentString(message.content);
+
+  const handleSubmitEdit = () => {
+    setIsEditing(false);
+    thread.submit(
+      {
+        messages: [
+          {
+            ...message,
+            content: value,
+          },
+        ],
+      },
+      {
+        checkpoint: parentCheckpoint,
+      },
+    );
+  };
+
+  return (
+    <div
+      className={cn(
+        "flex items-center ml-auto gap-2 px-4 py-2 group",
+        isEditing && "w-full max-w-xl",
+      )}
+    >
+      <div className={cn("flex flex-col gap-2", isEditing && "w-full")}>
+        {isEditing ? (
+          <EditableContent
+            value={value}
+            setValue={setValue}
+            onSubmit={handleSubmitEdit}
+          />
+        ) : (
+          <p>{contentString}</p>
+        )}
+        <div className="flex gap-2 items-center ml-auto opacity-0 group-hover:opacity-100 transition-opacity">
+          <BranchSwitcher
+            branch={meta?.branch}
+            branchOptions={meta?.branchOptions}
+            onSelect={(branch) => thread.setBranch(branch)}
+            isLoading={isLoading}
+          />
+          <CommandBar
+            isLoading={isLoading}
+            content={contentString}
+            isEditing={isEditing}
+            setIsEditing={(c) => {
+              if (c) {
+                setValue(contentString);
+              }
+              setIsEditing(c);
+            }}
+            handleSubmitEdit={handleSubmitEdit}
+            isHumanMessage={true}
+          />
+        </div>
+      </div>
+    </div>
+  );
+}
--- a/src/components/thread/messages/shared.tsx
+++ b/src/components/thread/messages/shared.tsx
@@ -0,0 +1,213 @@
+import {
+  XIcon,
+  SendHorizontal,
+  RefreshCcw,
+  Pencil,
+  Copy,
+  CopyCheck,
+  ChevronLeft,
+  ChevronRight,
+} from "lucide-react";
+import { TooltipIconButton } from "../tooltip-icon-button";
+import { AnimatePresence, motion } from "framer-motion";
+import { useState } from "react";
+import { Button } from "@/components/ui/button";
+
+function ContentCopyable({
+  content,
+  disabled,
+}: {
+  content: string;
+  disabled: boolean;
+}) {
+  const [copied, setCopied] = useState(false);
+
+  const handleCopy = (e: React.MouseEvent<HTMLButtonElement, MouseEvent>) => {
+    e.stopPropagation();
+    navigator.clipboard.writeText(content);
+    setCopied(true);
+    setTimeout(() => setCopied(false), 2000);
+  };
+
+  return (
+    <TooltipIconButton
+      onClick={(e: any) => handleCopy(e)}
+      variant="ghost"
+      tooltip="Copy content"
+      disabled={disabled}
+    >
+      <AnimatePresence mode="wait" initial={false}>
+        {copied ? (
+          <motion.div
+            key="check"
+            initial={{ opacity: 0, scale: 0.8 }}
+            animate={{ opacity: 1, scale: 1 }}
+            exit={{ opacity: 0, scale: 0.8 }}
+            transition={{ duration: 0.15 }}
+          >
+            <CopyCheck className="text-green-500" />
+          </motion.div>
+        ) : (
+          <motion.div
+            key="copy"
+            initial={{ opacity: 0, scale: 0.8 }}
+            animate={{ opacity: 1, scale: 1 }}
+            exit={{ opacity: 0, scale: 0.8 }}
+            transition={{ duration: 0.15 }}
+          >
+            <Copy />
+          </motion.div>
+        )}
+      </AnimatePresence>
+    </TooltipIconButton>
+  );
+}
+
+export function BranchSwitcher({
+  branch,
+  branchOptions,
+  onSelect,
+  isLoading,
+}: {
+  branch: string | undefined;
+  branchOptions: string[] | undefined;
+  onSelect: (branch: string) => void;
+  isLoading: boolean;
+}) {
+  if (!branchOptions || !branch) return null;
+  const index = branchOptions.indexOf(branch);
+
+  return (
+    <div className="flex items-center gap-2">
+      <Button
+        variant="ghost"
+        size="icon"
+        onClick={() => {
+          const prevBranch = branchOptions[index - 1];
+          if (!prevBranch) return;
+          onSelect(prevBranch);
+        }}
+        disabled={isLoading}
+      >
+        <ChevronLeft />
+      </Button>
+      <span className="text-sm">
+        {index + 1} / {branchOptions.length}
+      </span>
+      <Button
+        variant="ghost"
+        size="icon"
+        onClick={() => {
+          const nextBranch = branchOptions[index + 1];
+          if (!nextBranch) return;
+          onSelect(nextBranch);
+        }}
+        disabled={isLoading}
+      >
+        <ChevronRight />
+      </Button>
+    </div>
+  );
+}
+
+export function CommandBar({
+  content,
+  isHumanMessage,
+  isAiMessage,
+  isEditing,
+  setIsEditing,
+  handleSubmitEdit,
+  handleRegenerate,
+  isLoading,
+}: {
+  content: string;
+  isHumanMessage?: boolean;
+  isAiMessage?: boolean;
+  isEditing?: boolean;
+  setIsEditing?: React.Dispatch<React.SetStateAction<boolean>>;
+  handleSubmitEdit?: () => void;
+  handleRegenerate?: () => void;
+  isLoading: boolean;
+}) {
+  if (isHumanMessage && isAiMessage) {
+    throw new Error(
+      "Can only set one of isHumanMessage or isAiMessage to true, not both.",
+    );
+  }
+
+  if (!isHumanMessage && !isAiMessage) {
+    throw new Error(
+      "One of isHumanMessage or isAiMessage must be set to true.",
+    );
+  }
+
+  if (
+    isHumanMessage &&
+    (isEditing === undefined ||
+      setIsEditing === undefined ||
+      handleSubmitEdit === undefined)
+  ) {
+    throw new Error(
+      "If isHumanMessage is true, all of isEditing, setIsEditing, and handleSubmitEdit must be set.",
+    );
+  }
+
+  const showEdit =
+    isHumanMessage &&
+    isEditing !== undefined &&
+    !!setIsEditing &&
+    !!handleSubmitEdit;
+
+  if (isHumanMessage && isEditing && !!setIsEditing && !!handleSubmitEdit) {
+    return (
+      <div className="flex items-center gap-2">
+        <TooltipIconButton
+          disabled={isLoading}
+          tooltip="Cancel edit"
+          variant="ghost"
+          onClick={() => {
+            setIsEditing(false);
+          }}
+        >
+          <XIcon />
+        </TooltipIconButton>
+        <TooltipIconButton
+          disabled={isLoading}
+          tooltip="Submit"
+          variant="secondary"
+          onClick={handleSubmitEdit}
+        >
+          <SendHorizontal />
+        </TooltipIconButton>
+      </div>
+    );
+  }
+
+  return (
+    <div className="flex items-center gap-2">
+      <ContentCopyable content={content} disabled={isLoading} />
+      {isAiMessage && !!handleRegenerate && (
+        <TooltipIconButton
+          disabled={isLoading}
+          tooltip="Refresh"
+          variant="ghost"
+          onClick={handleRegenerate}
+        >
+          <RefreshCcw />
+        </TooltipIconButton>
+      )}
+      {showEdit && (
+        <TooltipIconButton
+          disabled={isLoading}
+          tooltip="Edit"
+          variant="ghost"
+          onClick={() => {
+            setIsEditing?.(true);
+          }}
+        >
+          <Pencil />
+        </TooltipIconButton>
+      )}
+    </div>
+  );
+}
--- a/src/components/thread/tooltip-icon-button.tsx
+++ b/src/components/thread/tooltip-icon-button.tsx
@@ -0,0 +1,44 @@
+"use client";
+
+import { forwardRef } from "react";
+
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipProvider,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { Button, ButtonProps } from "@/components/ui/button";
+import { cn } from "@/lib/utils";
+
+export type TooltipIconButtonProps = ButtonProps & {
+  tooltip: string;
+  side?: "top" | "bottom" | "left" | "right";
+};
+
+export const TooltipIconButton = forwardRef<
+  HTMLButtonElement,
+  TooltipIconButtonProps
+>(({ children, tooltip, side = "bottom", className, ...rest }, ref) => {
+  return (
+    <TooltipProvider>
+      <Tooltip>
+        <TooltipTrigger asChild>
+          <Button
+            variant="ghost"
+            size="icon"
+            {...rest}
+            className={cn("size-6 p-1", className)}
+            ref={ref}
+          >
+            {children}
+            <span className="sr-only">{tooltip}</span>
+          </Button>
+        </TooltipTrigger>
+        <TooltipContent side={side}>{tooltip}</TooltipContent>
+      </Tooltip>
+    </TooltipProvider>
+  );
+});
+
+TooltipIconButton.displayName = "TooltipIconButton";
--- a/src/components/thread/utils.ts
+++ b/src/components/thread/utils.ts
@@ -0,0 +1,9 @@
+import { MessageContent } from "@langchain/core/messages";
+
+export function getContentString(content: MessageContent): string {
+  if (typeof content === "string") return content;
+  const texts = content
+    .filter((c): c is { type: "text"; text: string } => c.type === "text")
+    .map((c) => c.text);
+  return texts.join(" ");
+}
--- a/src/components/ui/avatar.tsx
+++ b/src/components/ui/avatar.tsx
@@ -0,0 +1,51 @@
+import * as React from "react";
+import * as AvatarPrimitive from "@radix-ui/react-avatar";
+
+import { cn } from "@/lib/utils";
+
+function Avatar({
+  className,
+  ...props
+}: React.ComponentProps<typeof AvatarPrimitive.Root>) {
+  return (
+    <AvatarPrimitive.Root
+      data-slot="avatar"
+      className={cn(
+        "relative flex size-8 shrink-0 overflow-hidden rounded-full",
+        className,
+      )}
+      {...props}
+    />
+  );
+}
+
+function AvatarImage({
+  className,
+  ...props
+}: React.ComponentProps<typeof AvatarPrimitive.Image>) {
+  return (
+    <AvatarPrimitive.Image
+      data-slot="avatar-image"
+      className={cn("aspect-square size-full", className)}
+      {...props}
+    />
+  );
+}
+
+function AvatarFallback({
+  className,
+  ...props
+}: React.ComponentProps<typeof AvatarPrimitive.Fallback>) {
+  return (
+    <AvatarPrimitive.Fallback
+      data-slot="avatar-fallback"
+      className={cn(
+        "bg-muted flex size-full items-center justify-center rounded-full",
+        className,
+      )}
+      {...props}
+    />
+  );
+}
+
+export { Avatar, AvatarImage, AvatarFallback };
--- a/src/components/ui/button.tsx
+++ b/src/components/ui/button.tsx
@@ -0,0 +1,60 @@
+import * as React from "react";
+import { Slot } from "@radix-ui/react-slot";
+import { cva, type VariantProps } from "class-variance-authority";
+
+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",
+  {
+    variants: {
+      variant: {
+        default:
+          "bg-primary text-primary-foreground shadow-xs hover:bg-primary/90",
+        destructive:
+          "bg-destructive text-white shadow-xs hover:bg-destructive/90 focus-visible:ring-destructive/20 dark:focus-visible:ring-destructive/40",
+        outline:
+          "border border-input bg-background shadow-xs hover:bg-accent hover:text-accent-foreground",
+        secondary:
+          "bg-secondary text-secondary-foreground shadow-xs hover:bg-secondary/80",
+        ghost: "hover:bg-accent hover:text-accent-foreground",
+        link: "text-primary underline-offset-4 hover:underline",
+      },
+      size: {
+        default: "h-9 px-4 py-2 has-[>svg]:px-3",
+        sm: "h-8 rounded-md gap-1.5 px-3 has-[>svg]:px-2.5",
+        lg: "h-10 rounded-md px-6 has-[>svg]:px-4",
+        icon: "size-9",
+      },
+    },
+    defaultVariants: {
+      variant: "default",
+      size: "default",
+    },
+  },
+);
+
+type ButtonProps = React.ComponentProps<"button"> &
+  VariantProps<typeof buttonVariants> & {
+    asChild?: boolean;
+  };
+
+function Button({
+  className,
+  variant,
+  size,
+  asChild = false,
+  ...props
+}: ButtonProps) {
+  const Comp = asChild ? Slot : "button";
+
+  return (
+    <Comp
+      data-slot="button"
+      className={cn(buttonVariants({ variant, size, className }))}
+      {...props}
+    />
+  );
+}
+
+export { Button, buttonVariants, type ButtonProps };
--- a/src/components/ui/input.tsx
+++ b/src/components/ui/input.tsx
@@ -0,0 +1,20 @@
+import * as React from "react";
+
+import { cn } from "@/lib/utils";
+
+function Input({ className, type, ...props }: React.ComponentProps<"input">) {
+  return (
+    <input
+      type={type}
+      data-slot="input"
+      className={cn(
+        "border-input file:text-foreground placeholder:text-muted-foreground selection:bg-primary selection:text-primary-foreground flex h-9 w-full min-w-0 rounded-md border bg-transparent px-3 py-1 text-base shadow-xs transition-[color,box-shadow] outline-none file:inline-flex file:h-7 file:border-0 file:bg-transparent file:text-sm file:font-medium disabled:pointer-events-none disabled:cursor-not-allowed disabled:opacity-50 md:text-sm",
+        "aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive",
+        className,
+      )}
+      {...props}
+    />
+  );
+}
+
+export { Input };
--- a/src/components/ui/textarea.tsx
+++ b/src/components/ui/textarea.tsx
@@ -0,0 +1,18 @@
+import * as React from "react";
+
+import { cn } from "@/lib/utils";
+
+function Textarea({ className, ...props }: React.ComponentProps<"textarea">) {
+  return (
+    <textarea
+      data-slot="textarea"
+      className={cn(
+        "border-input placeholder:text-muted-foreground focus-visible:border-ring focus-visible:ring-ring/50 aria-invalid:ring-destructive/20 dark:aria-invalid:ring-destructive/40 aria-invalid:border-destructive flex field-sizing-content min-h-16 w-full rounded-md border bg-transparent px-3 py-2 text-base shadow-xs transition-[color,box-shadow] outline-none focus-visible:ring-[3px] disabled:cursor-not-allowed disabled:opacity-50 md:text-sm",
+        className,
+      )}
+      {...props}
+    />
+  );
+}
+
+export { Textarea };
--- a/src/components/ui/tooltip.tsx
+++ b/src/components/ui/tooltip.tsx
@@ -0,0 +1,59 @@
+import * as React from "react";
+import * as TooltipPrimitive from "@radix-ui/react-tooltip";
+
+import { cn } from "@/lib/utils";
+
+function TooltipProvider({
+  delayDuration = 0,
+  ...props
+}: React.ComponentProps<typeof TooltipPrimitive.Provider>) {
+  return (
+    <TooltipPrimitive.Provider
+      data-slot="tooltip-provider"
+      delayDuration={delayDuration}
+      {...props}
+    />
+  );
+}
+
+function Tooltip({
+  ...props
+}: React.ComponentProps<typeof TooltipPrimitive.Root>) {
+  return (
+    <TooltipProvider>
+      <TooltipPrimitive.Root data-slot="tooltip" {...props} />
+    </TooltipProvider>
+  );
+}
+
+function TooltipTrigger({
+  ...props
+}: React.ComponentProps<typeof TooltipPrimitive.Trigger>) {
+  return <TooltipPrimitive.Trigger data-slot="tooltip-trigger" {...props} />;
+}
+
+function TooltipContent({
+  className,
+  sideOffset = 0,
+  children,
+  ...props
+}: React.ComponentProps<typeof TooltipPrimitive.Content>) {
+  return (
+    <TooltipPrimitive.Portal>
+      <TooltipPrimitive.Content
+        data-slot="tooltip-content"
+        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,
+        )}
+        {...props}
+      >
+        {children}
+        <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 };
--- a/src/index.css
+++ b/src/index.css
@@ -0,0 +1,124 @@
+@import "tailwindcss";
+
+@plugin "tailwindcss-animate";
+
+@custom-variant dark (&:is(.dark *));
+
+:root {
+  --background: oklch(1 0 0);
+  --foreground: oklch(0.145 0 0);
+  --card: oklch(1 0 0);
+  --card-foreground: oklch(0.145 0 0);
+  --popover: oklch(1 0 0);
+  --popover-foreground: oklch(0.145 0 0);
+  --primary: oklch(0.205 0 0);
+  --primary-foreground: oklch(0.985 0 0);
+  --secondary: oklch(0.97 0 0);
+  --secondary-foreground: oklch(0.205 0 0);
+  --muted: oklch(0.97 0 0);
+  --muted-foreground: oklch(0.556 0 0);
+  --accent: oklch(0.97 0 0);
+  --accent-foreground: oklch(0.205 0 0);
+  --destructive: oklch(0.577 0.245 27.325);
+  --destructive-foreground: oklch(0.577 0.245 27.325);
+  --border: oklch(0.922 0 0);
+  --input: oklch(0.922 0 0);
+  --ring: oklch(0.87 0 0);
+  --chart-1: oklch(0.646 0.222 41.116);
+  --chart-2: oklch(0.6 0.118 184.704);
+  --chart-3: oklch(0.398 0.07 227.392);
+  --chart-4: oklch(0.828 0.189 84.429);
+  --chart-5: oklch(0.769 0.188 70.08);
+  --radius: 0.625rem;
+  --sidebar: oklch(0.985 0 0);
+  --sidebar-foreground: oklch(0.145 0 0);
+  --sidebar-primary: oklch(0.205 0 0);
+  --sidebar-primary-foreground: oklch(0.985 0 0);
+  --sidebar-accent: oklch(0.97 0 0);
+  --sidebar-accent-foreground: oklch(0.205 0 0);
+  --sidebar-border: oklch(0.922 0 0);
+  --sidebar-ring: oklch(0.87 0 0);
+}
+
+.dark {
+  --background: oklch(0.145 0 0);
+  --foreground: oklch(0.985 0 0);
+  --card: oklch(0.145 0 0);
+  --card-foreground: oklch(0.985 0 0);
+  --popover: oklch(0.145 0 0);
+  --popover-foreground: oklch(0.985 0 0);
+  --primary: oklch(0.985 0 0);
+  --primary-foreground: oklch(0.205 0 0);
+  --secondary: oklch(0.269 0 0);
+  --secondary-foreground: oklch(0.985 0 0);
+  --muted: oklch(0.269 0 0);
+  --muted-foreground: oklch(0.708 0 0);
+  --accent: oklch(0.269 0 0);
+  --accent-foreground: oklch(0.985 0 0);
+  --destructive: oklch(0.396 0.141 25.723);
+  --destructive-foreground: oklch(0.637 0.237 25.331);
+  --border: oklch(0.269 0 0);
+  --input: oklch(0.269 0 0);
+  --ring: oklch(0.439 0 0);
+  --chart-1: oklch(0.488 0.243 264.376);
+  --chart-2: oklch(0.696 0.17 162.48);
+  --chart-3: oklch(0.769 0.188 70.08);
+  --chart-4: oklch(0.627 0.265 303.9);
+  --chart-5: oklch(0.645 0.246 16.439);
+  --sidebar: oklch(0.205 0 0);
+  --sidebar-foreground: oklch(0.985 0 0);
+  --sidebar-primary: oklch(0.488 0.243 264.376);
+  --sidebar-primary-foreground: oklch(0.985 0 0);
+  --sidebar-accent: oklch(0.269 0 0);
+  --sidebar-accent-foreground: oklch(0.985 0 0);
+  --sidebar-border: oklch(0.269 0 0);
+  --sidebar-ring: oklch(0.439 0 0);
+}
+
+@theme inline {
+  --color-background: var(--background);
+  --color-foreground: var(--foreground);
+  --color-card: var(--card);
+  --color-card-foreground: var(--card-foreground);
+  --color-popover: var(--popover);
+  --color-popover-foreground: var(--popover-foreground);
+  --color-primary: var(--primary);
+  --color-primary-foreground: var(--primary-foreground);
+  --color-secondary: var(--secondary);
+  --color-secondary-foreground: var(--secondary-foreground);
+  --color-muted: var(--muted);
+  --color-muted-foreground: var(--muted-foreground);
+  --color-accent: var(--accent);
+  --color-accent-foreground: var(--accent-foreground);
+  --color-destructive: var(--destructive);
+  --color-destructive-foreground: var(--destructive-foreground);
+  --color-border: var(--border);
+  --color-input: var(--input);
+  --color-ring: var(--ring);
+  --color-chart-1: var(--chart-1);
+  --color-chart-2: var(--chart-2);
+  --color-chart-3: var(--chart-3);
+  --color-chart-4: var(--chart-4);
+  --color-chart-5: var(--chart-5);
+  --radius-sm: calc(var(--radius) - 4px);
+  --radius-md: calc(var(--radius) - 2px);
+  --radius-lg: var(--radius);
+  --radius-xl: calc(var(--radius) + 4px);
+  --color-sidebar: var(--sidebar);
+  --color-sidebar-foreground: var(--sidebar-foreground);
+  --color-sidebar-primary: var(--sidebar-primary);
+  --color-sidebar-primary-foreground: var(--sidebar-primary-foreground);
+  --color-sidebar-accent: var(--sidebar-accent);
+  --color-sidebar-accent-foreground: var(--sidebar-accent-foreground);
+  --color-sidebar-border: var(--sidebar-border);
+  --color-sidebar-ring: var(--sidebar-ring);
+}
+
+@layer base {
+  * {
+    @apply border-border outline-ring/50;
+  }
+  body {
+    @apply bg-background text-foreground;
+  }
+}
--- a/src/lib/ensure-tool-responses.ts
+++ b/src/lib/ensure-tool-responses.ts
@@ -0,0 +1,34 @@
+import { v4 as uuidv4 } from "uuid";
+import { Message, ToolMessage } from "@langchain/langgraph-sdk";
+
+export const DO_NOT_RENDER_ID_PREFIX = "do-not-render-";
+
+export function ensureToolCallsHaveResponses(messages: Message[]): Message[] {
+  const newMessages: ToolMessage[] = [];
+
+  messages.forEach((message, index) => {
+    if (message.type !== "ai" || message.tool_calls?.length === 0) {
+      // If it's not an AI message, or it doesn't have tool calls, we can ignore.
+      return;
+    }
+    // If it has tool calls, ensure the message which follows this is a tool message
+    const followingMessage = messages[index + 1];
+    if (followingMessage && followingMessage.type === "tool") {
+      // Following message is a tool message, so we can ignore.
+      return;
+    }
+
+    // Since the following message is not a tool message, we must create a new tool message
+    newMessages.push(
+      ...(message.tool_calls?.map((tc) => ({
+        type: "tool" as const,
+        tool_call_id: tc.id ?? "",
+        id: `${DO_NOT_RENDER_ID_PREFIX}${uuidv4()}`,
+        name: tc.name,
+        content: "Successfully handled tool call.",
+      })) ?? []),
+    );
+  });
+
+  return newMessages;
+}
--- a/src/lib/utils.ts
+++ b/src/lib/utils.ts
@@ -0,0 +1,6 @@
+import { clsx, type ClassValue } from "clsx";
+import { twMerge } from "tailwind-merge";
+
+export function cn(...inputs: ClassValue[]) {
+  return twMerge(clsx(inputs));
+}
--- a/src/main.tsx
+++ b/src/main.tsx
@@ -1,5 +1,10 @@
 import { createRoot } from "react-dom/client";
 import "./index.css";
 import App from "./App.tsx";
+import { StreamProvider } from "./providers/Stream.tsx";

-createRoot(document.getElementById("root")!).render(<App />);
+createRoot(document.getElementById("root")!).render(
+  <StreamProvider>
+    <App />
+  </StreamProvider>,
+);
--- a/src/providers/Stream.tsx
+++ b/src/providers/Stream.tsx
@@ -0,0 +1,47 @@
+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";
+
+const useTypedStream = useStream<
+  { messages: Message[]; ui: UIMessage[] },
+  {
+    UpdateType: {
+      messages?: Message[] | Message | string;
+      ui?: (UIMessage | RemoveUIMessage)[] | UIMessage | RemoveUIMessage;
+    };
+    CustomUpdateType: UIMessage | RemoveUIMessage;
+  }
+>;
+
+type StreamContextType = ReturnType<typeof useTypedStream>;
+const StreamContext = createContext<StreamContextType | undefined>(undefined);
+
+export const StreamProvider: React.FC<{ children: ReactNode }> = ({
+  children,
+}) => {
+  const streamValue = useTypedStream({
+    apiUrl: "http://localhost:2024",
+    assistantId: "agent",
+  });
+
+  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;
--- a/tailwind.config.js
+++ b/tailwind.config.js
@@ -0,0 +1,57 @@
+/** @type {import('tailwindcss').Config} */
+module.exports = {
+  darkMode: ["class"],
+  content: ["./index.html", "./src/**/*.{ts,tsx,js,jsx}"],
+  theme: {
+    extend: {
+      borderRadius: {
+        lg: "var(--radius)",
+        md: "calc(var(--radius) - 2px)",
+        sm: "calc(var(--radius) - 4px)",
+      },
+      colors: {
+        background: "hsl(var(--background))",
+        foreground: "hsl(var(--foreground))",
+        card: {
+          DEFAULT: "hsl(var(--card))",
+          foreground: "hsl(var(--card-foreground))",
+        },
+        popover: {
+          DEFAULT: "hsl(var(--popover))",
+          foreground: "hsl(var(--popover-foreground))",
+        },
+        primary: {
+          DEFAULT: "hsl(var(--primary))",
+          foreground: "hsl(var(--primary-foreground))",
+        },
+        secondary: {
+          DEFAULT: "hsl(var(--secondary))",
+          foreground: "hsl(var(--secondary-foreground))",
+        },
+        muted: {
+          DEFAULT: "hsl(var(--muted))",
+          foreground: "hsl(var(--muted-foreground))",
+        },
+        accent: {
+          DEFAULT: "hsl(var(--accent))",
+          foreground: "hsl(var(--accent-foreground))",
+        },
+        destructive: {
+          DEFAULT: "hsl(var(--destructive))",
+          foreground: "hsl(var(--destructive-foreground))",
+        },
+        border: "hsl(var(--border))",
+        input: "hsl(var(--input))",
+        ring: "hsl(var(--ring))",
+        chart: {
+          1: "hsl(var(--chart-1))",
+          2: "hsl(var(--chart-2))",
+          3: "hsl(var(--chart-3))",
+          4: "hsl(var(--chart-4))",
+          5: "hsl(var(--chart-5))",
+        },
+      },
+    },
+  },
+  plugins: [require("tailwindcss-animate")],
+};
--- a/tsconfig.app.json
+++ b/tsconfig.app.json
@@ -1,10 +1,10 @@
 {
  "compilerOptions": {
    "tsBuildInfoFile": "./node_modules/.tmp/tsconfig.app.tsbuildinfo",
-    "target": "ES2020",
+    "target": "ES2023",
    "useDefineForClassFields": true,
-    "lib": ["ES2020", "DOM", "DOM.Iterable"],
-    "module": "ESNext",
+    "lib": ["ES2023", "DOM", "DOM.Iterable"],
+    "module": "ES2022",
    "skipLibCheck": true,

    /* Bundler mode */
@@ -20,7 +20,11 @@
    "noUnusedLocals": true,
    "noUnusedParameters": true,
    "noFallthroughCasesInSwitch": true,
-    "noUncheckedSideEffectImports": true
+    "noUncheckedSideEffectImports": true,
+    "baseUrl": ".",
+    "paths": {
+      "@/*": ["./src/*"]
+    }
  },
  "include": ["src", "agent"]
 }
--- a/tsconfig.json
+++ b/tsconfig.json
@@ -3,5 +3,11 @@
  "references": [
    { "path": "./tsconfig.app.json" },
    { "path": "./tsconfig.node.json" }
-  ]
+  ],
+  "compilerOptions": {
+    "baseUrl": ".",
+    "paths": {
+      "@/*": ["./src/*"]
+    }
+  }
 }
--- a/vite.config.ts
+++ b/vite.config.ts
@@ -1,7 +1,14 @@
-import { defineConfig } from 'vite'
-import react from '@vitejs/plugin-react'
+import path from "path";
+import { defineConfig } from "vite";
+import react from "@vitejs/plugin-react";
+import tailwindcss from "@tailwindcss/vite";

 // https://vite.dev/config/
 export default defineConfig({
-  plugins: [react()],
-})
+  plugins: [react(), tailwindcss()],
+  resolve: {
+    alias: {
+      "@": path.resolve(__dirname, "./src"),
+    },
+  },
+});