diff --git a/agent-demo/.gitignore b/agent-demo/.gitignore new file mode 100644 index 0000000000..5ef6a52078 --- /dev/null +++ b/agent-demo/.gitignore @@ -0,0 +1,41 @@ +# See https://help.github.com/articles/ignoring-files/ for more about ignoring files. + +# dependencies +/node_modules +/.pnp +.pnp.* +.yarn/* +!.yarn/patches +!.yarn/plugins +!.yarn/releases +!.yarn/versions + +# testing +/coverage + +# next.js +/.next/ +/out/ + +# production +/build + +# misc +.DS_Store +*.pem + +# debug +npm-debug.log* +yarn-debug.log* +yarn-error.log* +.pnpm-debug.log* + +# env files (can opt-in for committing if needed) +.env* + +# vercel +.vercel + +# typescript +*.tsbuildinfo +next-env.d.ts diff --git a/agent-demo/README.md b/agent-demo/README.md new file mode 100644 index 0000000000..3c065ca151 --- /dev/null +++ b/agent-demo/README.md @@ -0,0 +1,27 @@ +This is the [assistant-ui](https://github.com/Yonom/assistant-ui) starter project. + +## Getting Started + +First, add your OpenAI API key to `.env.local` file: + +``` +OPENAI_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx +``` + +Then, run the development server: + +```bash +npm run dev +# or +yarn dev +# or +pnpm dev +# or +bun dev +``` + +Open [http://localhost:3000](http://localhost:3000) with your browser to see the result. + +You can start editing the page by modifying `app/page.tsx`. The page auto-updates as you edit the file. + +test diff --git a/agent-demo/app/api/chat/route.ts b/agent-demo/app/api/chat/route.ts new file mode 100644 index 0000000000..77f931568e --- /dev/null +++ b/agent-demo/app/api/chat/route.ts @@ -0,0 +1,41 @@ +import { openai } from "@ai-sdk/openai"; +import { frontendTools } from "@assistant-ui/react-ai-sdk"; +import { + aiDocumentFormats, + injectDocumentStateMessages, +} from "@blocknote/xl-ai/server"; +import { convertToModelMessages, streamText, UIMessage } from "ai"; +import { JSONSchema7 } from "json-schema"; + +export async function POST(req: Request) { + const { + messages, + tools, + }: { + messages: UIMessage[]; + tools: Record; + } = await req.json(); + const result = streamText({ + system: aiDocumentFormats.html.systemPrompt, + model: openai("gpt-4o-2024-08-06"), // openai("gpt-5-nano"), + messages: convertToModelMessages(injectDocumentStateMessages(messages)), + tools: { + ...(frontendTools(tools) as any), // TODO: tools vs toolDefinitions + web_search: openai.tools.webSearch({}), + }, + // providerOptions: { + // openai: { + // reasoningEffort: "low", + // }, + // }, + // toolChoice: "required", // TODO: make configurable from client and make toolbar "required" + }); + + return result.toUIMessageStreamResponse(); +} + +// - the "getDocument" tool shows the document as an array of html blocks (the cursor is BETWEEN two blocks as indicated by cursor: true). +// - the "getDocumentSelection" tool shows the current user selection, if any. +// - when the document is empty, prefer updating the empty block before adding new blocks. Otherwise, prefer updating existing blocks over removing and adding (but this also depends on the user's question). +// Don't call "getDocument" or "getDocumentSelection" tools directly, the information is already available. +// When there is a selection, issue operations against the result of "getDocumentSelection". You can still use the result of "getDocument" (which includes the selection as well) to understand the context. diff --git a/agent-demo/app/assistant.tsx b/agent-demo/app/assistant.tsx new file mode 100644 index 0000000000..9736c44d02 --- /dev/null +++ b/agent-demo/app/assistant.tsx @@ -0,0 +1,81 @@ +"use client"; + +import { AssistantModal } from "@/components/assistant-ui/assistant-modal"; +import { ThreadListSidebar } from "@/components/assistant-ui/threadlist-sidebar"; +import { + Breadcrumb, + BreadcrumbItem, + BreadcrumbList, +} from "@/components/ui/breadcrumb"; +import { Separator } from "@/components/ui/separator"; +import { + SidebarInset, + SidebarProvider, + SidebarTrigger, +} from "@/components/ui/sidebar"; +import { AssistantRuntimeProvider } from "@assistant-ui/react"; +import { useAISDKRuntime } from "@assistant-ui/react-ai-sdk"; + +import { useChatContext } from "@/components/ChatContext"; +import { useChat } from "@ai-sdk/react"; +import dynamic from "next/dynamic"; +import { useEffect } from "react"; + +const Document = dynamic(() => import("./document"), { + ssr: false, +}); + +console.log("Assistant", Document); +export const Assistant = () => { + const ctx = useChatContext(); + const chat = useChat({ + chat: ctx.chat, + }); + + const runtime = useAISDKRuntime(chat); + + useEffect(() => { + // not documented! + (ctx.transport as any).setRuntime(runtime); + }, [runtime, ctx.transport]); + + return ( + + +
+ + +
+ + + + + + {/* */} + Agent demo + {/* */} + + {/* + + Starter Template + */} + + +
+
+ {/* */} +
+ +
+
+ +
+ + + + ); +}; diff --git a/agent-demo/app/document.tsx b/agent-demo/app/document.tsx new file mode 100644 index 0000000000..a5dff2567f --- /dev/null +++ b/agent-demo/app/document.tsx @@ -0,0 +1,254 @@ +"use client"; + +import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { + getMessageWithToolCallId, + useChatContext, +} from "@/components/ChatContext"; +import { Chat } from "@ai-sdk/react"; +import { + ActionBarPrimitive, + makeAssistantTool, + makeAssistantToolUI, + tool, + ToolCallMessagePartProps, +} from "@assistant-ui/react"; +import { BlockNoteEditor } from "@blocknote/core"; +import { filterSuggestionItems } from "@blocknote/core/extensions"; +import "@blocknote/core/fonts/inter.css"; +import { en } from "@blocknote/core/locales"; +import { BlockNoteView } from "@blocknote/mantine"; +import "@blocknote/mantine/style.css"; +import { + FormattingToolbar, + FormattingToolbarController, + getDefaultReactSlashMenuItems, + getFormattingToolbarItems, + SuggestionMenuController, + useCreateBlockNote, +} from "@blocknote/react"; +import { + aiDocumentFormats, + AIExtension, + AIMenuController, + AIToolbarButton, + createStreamToolsArraySchema, + getAISlashMenuItems, +} from "@blocknote/xl-ai"; +import { en as aiEn } from "@blocknote/xl-ai/locales"; +import "@blocknote/xl-ai/style.css"; +import { PencilIcon, UndoIcon } from "lucide-react"; +import { useCallback, useEffect, useMemo } from "react"; + +const UndoActionBar = () => { + return ( + + {/* */} + { + window.alert("Not implemented"); + }} + > + + + {/* */} + + ); +}; + +const BlockNoteToolUI = ( + props: ToolCallMessagePartProps & { + editor: BlockNoteEditor; + }, +) => { + const ctx = useChatContext(); + console.log("render"); + const onHover = useCallback(() => { + ctx.setPreviewDocument(props.toolCallId); + }, [props.toolCallId]); + + const onMouseLeave = useCallback(() => { + ctx.setPreviewDocument(undefined); + console.log("onMouseLeave"); + }, [ctx]); + + const message = getMessageWithToolCallId(ctx.chat, props.toolCallId); + + if (props.status.type === "running") { + return ( +
+ Updating document... +
+ ); + } else if (props.status.type === "complete") { + return ( + // TODO: should get rid of min-h-10, but otherwise layout shifts when undo is shown +
+ {" "} + + Updated document + + {/* */} + +
+ ); + } else { + throw new Error("Not implemented"); + } +}; + +export default function Document() { + const ctx = useChatContext(); + + // Creates a new editor instance. + const editor = useCreateBlockNote({ + dictionary: { + ...en, + ai: aiEn, // add default translations for the AI extension + }, + // Register the AI extension + extensions: [ + AIExtension({ + chatProvider: () => ctx.chat as Chat, + // transport: new DefaultChatTransport({ + // // URL to your backend API, see example source in `packages/xl-ai-server/src/routes/regular.ts` + // api: `${BASE_URL}/regular/streamText`, + // }), + }), + // createAIAutoCompleteExtension(), + ], + // We set some initial content for demo purposes + initialContent: [ + { + type: "heading", + props: { + level: 1, + }, + content: "Open source software", + }, + { + type: "paragraph", + content: + "Open source software refers to computer programs whose source code is made available to the public, allowing anyone to view, modify, and distribute the code. This model stands in contrast to proprietary software, where the source code is kept secret and only the original creators have the right to make changes. Open projects are developed collaboratively, often by communities of developers from around the world, and are typically distributed under licenses that promote sharing and openness.", + }, + { + type: "paragraph", + content: + "One of the primary benefits of open source is the promotion of digital autonomy. By providing access to the source code, these programs empower users to control their own technology, customize software to fit their needs, and avoid vendor lock-in. This level of transparency also allows for greater security, as anyone can inspect the code for vulnerabilities or malicious elements. As a result, users are not solely dependent on a single company for updates, bug fixes, or continued support.", + }, + { + type: "paragraph", + content: + "Additionally, open development fosters innovation and collaboration. Developers can build upon existing projects, share improvements, and learn from each other, accelerating the pace of technological advancement. The open nature of these projects often leads to higher quality software, as bugs are identified and fixed more quickly by a diverse group of contributors. Furthermore, using open source can reduce costs for individuals, businesses, and governments, as it is often available for free and can be tailored to specific requirements without expensive licensing fees.", + }, + ], + }); + + useEffect(() => { + ctx.setEditor(editor); + }, [ctx, editor]); + + const streamTools = useMemo(() => { + return aiDocumentFormats.html + .getStreamToolsProvider() + .getStreamTools(editor); // TODO: not document dependent? + }, [editor]); + + const BlockNoteTool = makeAssistantTool({ + // type: "frontend", + toolName: "applyDocumentOperations", + ...tool({ + parameters: createStreamToolsArraySchema(streamTools), + execute: async (args) => { + // debugger; + console.log(args); + }, + }), + render: (props) => , + }); + + debugger; + const DocumentStateTool = makeAssistantToolUI({ + toolName: "document-state", + render: (props) => null, + }); + // Renders the editor instance using a React component. + return ( +
+ + + + {/* Add the AI Command menu to the editor */} + + + {/* We disabled the default formatting toolbar with `formattingToolbar=false` + and replace it for one with an "AI button" (defined below). + (See "Formatting Toolbar" in docs) + */} + + + {/* We disabled the default SlashMenu with `slashMenu=false` + and replace it for one with an AI option (defined below). + (See "Suggestion Menus" in docs) + */} + + +
+ ); +} + +// Formatting toolbar with the `AIToolbarButton` added +function FormattingToolbarWithAI() { + return ( + ( + + {...getFormattingToolbarItems()} + {/* Add the AI button */} + + + )} + /> + ); +} + +// Slash menu with the AI option added +function SuggestionMenuWithAI(props: { + editor: BlockNoteEditor; +}) { + return ( + + filterSuggestionItems( + [ + ...getDefaultReactSlashMenuItems(props.editor), + // add the default AI slash menu items, or define your own + ...getAISlashMenuItems(props.editor), + ], + query, + ) + } + /> + ); +} diff --git a/agent-demo/app/favicon.ico b/agent-demo/app/favicon.ico new file mode 100644 index 0000000000..718d6fea48 Binary files /dev/null and b/agent-demo/app/favicon.ico differ diff --git a/agent-demo/app/globals.css b/agent-demo/app/globals.css new file mode 100644 index 0000000000..a6986b058a --- /dev/null +++ b/agent-demo/app/globals.css @@ -0,0 +1,156 @@ +@import "tailwindcss"; +@import "tw-animate-css"; + +@custom-variant dark (&:is(.dark *)); + +@theme inline { + --color-background: var(--background); + --color-foreground: var(--foreground); + --font-sans: var(--font-geist-sans); + --font-mono: var(--font-geist-mono); + --color-sidebar-ring: var(--sidebar-ring); + --color-sidebar-border: var(--sidebar-border); + --color-sidebar-accent-foreground: var(--sidebar-accent-foreground); + --color-sidebar-accent: var(--sidebar-accent); + --color-sidebar-primary-foreground: var(--sidebar-primary-foreground); + --color-sidebar-primary: var(--sidebar-primary); + --color-sidebar-foreground: var(--sidebar-foreground); + --color-sidebar: var(--sidebar); + --color-chart-5: var(--chart-5); + --color-chart-4: var(--chart-4); + --color-chart-3: var(--chart-3); + --color-chart-2: var(--chart-2); + --color-chart-1: var(--chart-1); + --color-ring: var(--ring); + --color-input: var(--input); + --color-border: var(--border); + --color-destructive: var(--destructive); + --color-accent-foreground: var(--accent-foreground); + --color-accent: var(--accent); + --color-muted-foreground: var(--muted-foreground); + --color-muted: var(--muted); + --color-secondary-foreground: var(--secondary-foreground); + --color-secondary: var(--secondary); + --color-primary-foreground: var(--primary-foreground); + --color-primary: var(--primary); + --color-popover-foreground: var(--popover-foreground); + --color-popover: var(--popover); + --color-card-foreground: var(--card-foreground); + --color-card: var(--card); + --radius-sm: calc(var(--radius) - 4px); + --radius-md: calc(var(--radius) - 2px); + --radius-lg: var(--radius); + --radius-xl: calc(var(--radius) + 4px); + --animate-shimmer: shimmer-sweep var(--shimmer-duration, 1000ms) linear infinite both; + @keyframes shimmer-sweep { + from { + background-position: 150% 0; + } + to { + background-position: -100% 0; + } + } + @keyframes shimmer-sweep { + from { + background-position: 150% 0; + } + to { + background-position: -100% 0; + } + } + @keyframes shimmer-sweep { + from { + background-position: 150% 0; + } + to { + background-position: -100% 0; + } + } +} + +:root { + --radius: 0.625rem; + --background: oklch(1 0 0); + --foreground: oklch(0.141 0.005 285.823); + --card: oklch(1 0 0); + --card-foreground: oklch(0.141 0.005 285.823); + --popover: oklch(1 0 0); + --popover-foreground: oklch(0.141 0.005 285.823); + --primary: oklch(0.21 0.006 285.885); + --primary-foreground: oklch(0.985 0 0); + --secondary: oklch(0.967 0.001 286.375); + --secondary-foreground: oklch(0.21 0.006 285.885); + --muted: oklch(0.967 0.001 286.375); + --muted-foreground: oklch(0.552 0.016 285.938); + --accent: oklch(0.967 0.001 286.375); + --accent-foreground: oklch(0.21 0.006 285.885); + --destructive: oklch(0.577 0.245 27.325); + --border: oklch(0.92 0.004 286.32); + --input: oklch(0.92 0.004 286.32); + --ring: oklch(0.705 0.015 286.067); + --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); + --sidebar: oklch(0.985 0 0); + --sidebar-foreground: oklch(0.141 0.005 285.823); + --sidebar-primary: oklch(0.21 0.006 285.885); + --sidebar-primary-foreground: oklch(0.985 0 0); + --sidebar-accent: oklch(0.967 0.001 286.375); + --sidebar-accent-foreground: oklch(0.21 0.006 285.885); + --sidebar-border: oklch(0.92 0.004 286.32); + --sidebar-ring: oklch(0.705 0.015 286.067); +} + +.dark { + --background: oklch(0.141 0.005 285.823); + --foreground: oklch(0.985 0 0); + --card: oklch(0.21 0.006 285.885); + --card-foreground: oklch(0.985 0 0); + --popover: oklch(0.21 0.006 285.885); + --popover-foreground: oklch(0.985 0 0); + --primary: oklch(0.92 0.004 286.32); + --primary-foreground: oklch(0.21 0.006 285.885); + --secondary: oklch(0.274 0.006 286.033); + --secondary-foreground: oklch(0.985 0 0); + --muted: oklch(0.274 0.006 286.033); + --muted-foreground: oklch(0.705 0.015 286.067); + --accent: oklch(0.274 0.006 286.033); + --accent-foreground: oklch(0.985 0 0); + --destructive: oklch(0.704 0.191 22.216); + --border: oklch(1 0 0 / 10%); + --input: oklch(1 0 0 / 15%); + --ring: oklch(0.552 0.016 285.938); + --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.21 0.006 285.885); + --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.274 0.006 286.033); + --sidebar-accent-foreground: oklch(0.985 0 0); + --sidebar-border: oklch(1 0 0 / 10%); + --sidebar-ring: oklch(0.552 0.016 285.938); +} + +@layer base { + * { + @apply border-border outline-ring/50; + } + + :root { + color-scheme: light; + } + + :root.dark { + color-scheme: dark; + } + + body { + @apply bg-background text-foreground; + } +} \ No newline at end of file diff --git a/agent-demo/app/layout.tsx b/agent-demo/app/layout.tsx new file mode 100644 index 0000000000..3e747110b2 --- /dev/null +++ b/agent-demo/app/layout.tsx @@ -0,0 +1,34 @@ +import type { Metadata } from "next"; +import { Geist, Geist_Mono } from "next/font/google"; +import "./globals.css"; + +const geistSans = Geist({ + variable: "--font-geist-sans", + subsets: ["latin"], +}); + +const geistMono = Geist_Mono({ + variable: "--font-geist-mono", + subsets: ["latin"], +}); + +export const metadata: Metadata = { + title: "assistant-ui Starter App", + description: "Generated by create-assistant-ui", +}; + +export default function RootLayout({ + children, +}: Readonly<{ + children: React.ReactNode; +}>) { + return ( + + + {children} + + + ); +} diff --git a/agent-demo/app/page.tsx b/agent-demo/app/page.tsx new file mode 100644 index 0000000000..1d63a99332 --- /dev/null +++ b/agent-demo/app/page.tsx @@ -0,0 +1,30 @@ +"use client"; +import { Assistant } from "./assistant"; + +import { useRef } from "react"; + +import { UIMessage } from "ai"; + +import { BlockNoteAISDKChat } from "@/components/BlockNoteChat"; +import { ChatContextProvider } from "@/components/ChatContext"; +import { AssistantChatTransport } from "@assistant-ui/react-ai-sdk"; + +const transport = new AssistantChatTransport({ + api: "/api/chat", +}); + +export default function Home() { + const chatRef = useRef>(undefined); + if (!chatRef.current) { + // TODO: would be better to get rid of this, but currently we need the raw chat object (on context) + chatRef.current = new BlockNoteAISDKChat({ + // api: "/api/chat", + transport, + }); + } + return ( + + + + ); +} diff --git a/agent-demo/assistant-ui-feedback.md b/agent-demo/assistant-ui-feedback.md new file mode 100644 index 0000000000..1048e06668 --- /dev/null +++ b/agent-demo/assistant-ui-feedback.md @@ -0,0 +1,4 @@ +- need to call `transport.setRuntime` with manual ai-sdk setup is not document +- npx create creates an API handler that doesn't read tools (route.ts) +- some components don't have examples / screenshots (reasoning) +- sendmessage metadata diff --git a/agent-demo/components.json b/agent-demo/components.json new file mode 100644 index 0000000000..a64445d7a5 --- /dev/null +++ b/agent-demo/components.json @@ -0,0 +1,21 @@ +{ + "$schema": "https://ui.shadcn.com/schema.json", + "style": "new-york", + "rsc": true, + "tsx": true, + "tailwind": { + "config": "", + "css": "app/globals.css", + "baseColor": "zinc", + "cssVariables": true, + "prefix": "" + }, + "aliases": { + "components": "@/components", + "utils": "@/lib/utils", + "ui": "@/components/ui", + "lib": "@/lib", + "hooks": "@/hooks" + }, + "iconLibrary": "lucide" +} diff --git a/agent-demo/components/BlockNoteChat.tsx b/agent-demo/components/BlockNoteChat.tsx new file mode 100644 index 0000000000..60c8adc993 --- /dev/null +++ b/agent-demo/components/BlockNoteChat.tsx @@ -0,0 +1,74 @@ +import { Chat } from "@ai-sdk/react"; +import { BlockNoteEditor } from "@blocknote/core"; +import { + aiDocumentFormats, + AIRequest, + buildAIRequest, + sendMessageWithAIRequest, +} from "@blocknote/xl-ai"; +import { ChatInit, UIMessage } from "ai"; + +// maybe better to use a custom transport for this? +class BlockNoteAISDKChatRaw< + UI_MESSAGE extends UIMessage, +> extends Chat { + public readonly sendMessageOriginal: Chat["sendMessage"]; + constructor( + init: ChatInit, + public editor?: BlockNoteEditor, + ) { + super(init); + this.sendMessageOriginal = this.sendMessage; + } + + public sendMessageAlt = async ( + message: Parameters["sendMessage"]>[0], + options?: Parameters["sendMessage"]>[1], + aiRequest?: AIRequest, + ) => { + if ((options?.metadata as any)?.source === "blocknote-ai") { + // prevent loops + return this.sendMessageOriginal(message, options); + } + + aiRequest = + aiRequest ?? + (await buildAIRequest({ + editor: this.editor!, + useSelection: false, + deleteEmptyCursorBlock: true, + streamToolsProvider: aiDocumentFormats.html.getStreamToolsProvider(), + onBlockUpdated: () => {}, + documentStateBuilder: + aiDocumentFormats.html.defaultDocumentStateBuilder, + })); + + await sendMessageWithAIRequest( + this as Chat, + aiRequest, + message, + options, + ); + }; +} + +/** + * A properly typed version of BlockNoteAISDKChatRaw that correctly types + * the sendMessage method with extended options (streamTools, onStart). + * + * This class can be instantiated and provides proper TypeScript typing for + * the sendMessage method, ensuring that the extended options are recognized + * by the type system. + */ +export class BlockNoteAISDKChat< + UI_MESSAGE extends UIMessage, +> extends BlockNoteAISDKChatRaw { + /** + * Override sendMessage with the correct type signature that includes + * streamTools and onStart options. + */ + public override sendMessage: (typeof BlockNoteAISDKChatRaw)["prototype"]["sendMessageAlt"] = + async (message, options, aiRequest) => { + return this.sendMessageAlt(message as any, options, aiRequest); + }; +} diff --git a/agent-demo/components/ChatContext.tsx b/agent-demo/components/ChatContext.tsx new file mode 100644 index 0000000000..fecde61cd0 --- /dev/null +++ b/agent-demo/components/ChatContext.tsx @@ -0,0 +1,237 @@ +"use client"; + +import { Chat } from "@ai-sdk/react"; +import { BlockNoteEditor } from "@blocknote/core"; +import { _getApplySuggestionsTr, AIExtension } from "@blocknote/xl-ai"; +import { + ChatTransport, + isToolOrDynamicToolUIPart, + isToolUIPart, + UIMessage, +} from "ai"; +import { + createContext, + useContext, + useEffect, + useMemo, + useRef, + useState, +} from "react"; + +import { BlockNoteAISDKChat } from "@/components/BlockNoteChat"; +import { trackToolCallCheckpoints } from "@/components/checkpoints"; +import { Node, Slice } from "prosemirror-model"; + +export const ChatContextProvider = ({ + children, + chat, + transport, +}: { + children: React.ReactNode; + chat: BlockNoteAISDKChat; + transport: ChatTransport; +}) => { + // Note: lots of functionality on context, maybe some can be moved to components? + const [editor, setEditor] = useState< + BlockNoteEditor | undefined + >(undefined); + + useEffect(() => { + chat.editor = editor; + }, [editor, chat]); + + // TODO: dispose is never called + const { checkpoints } = useMemo( + () => + editor + ? trackToolCallCheckpoints(chat, editor) + : { checkpoints: new Map(), dispose: () => {} }, + [chat, editor], + ); + + const [hasSuggestions, setHasSuggestions] = useState(false); + + const acceptChanges = () => { + if (!editor) { + throw new Error("Editor not found"); + } + editor.getExtension(AIExtension)!.acceptChanges(); + chat.messages = chat.messages.map((message) => { + if ( + message.role === "assistant" && + message.parts.some( + (part) => + isToolUIPart(part) && part.type === "tool-applyDocumentOperations", + ) + ) { + if ((message.metadata as any)?.applied === undefined) { + message.metadata = { ...(message.metadata || {}), applied: true }; + } + } + return message; + }); + }; + const rejectChanges = () => { + if (!editor) { + throw new Error("Editor not found"); + } + editor.getExtension(AIExtension)!.rejectChanges(); + chat.messages = chat.messages.map((message) => { + if ( + message.role === "assistant" && + message.parts.some( + (part) => + isToolUIPart(part) && part.type === "tool-applyDocumentOperations", + ) + ) { + if ((message.metadata as any)?.applied === undefined) { + message.metadata = { ...(message.metadata || {}), applied: false }; + } + } + return message; + }); + }; + + const [previewDocument, setPreviewDocument] = useState(); + + useEffect(() => { + if (!editor) { + return; + } + const unsubscribe = editor.onChange((editor) => { + const tr = _getApplySuggestionsTr(editor); + if (tr.docChanged) { + setHasSuggestions(true); + } else { + setHasSuggestions(false); + } + }); + return unsubscribe; + }, [editor]); + + const previousDocumentState = useRef(undefined); + + useEffect(() => { + console.log("previewDocument", previewDocument); + if (!editor) { + return; + } + if (!previewDocument) { + console.log("resetting document 1"); + if (previousDocumentState.current) { + // reset + console.log("resetting document"); + + editor.transact((tr) => { + tr.replace( + 0, + tr.doc.content.size, + new Slice(previousDocumentState.current!.content, 0, 0), + ); + }); + previousDocumentState.current = undefined; + } + return; + } + + // set preview document + const documentState = checkpoints.get(previewDocument)!; + + if (!documentState || !documentState.pm) { + throw new Error("Document state not found"); + } + + const newNode = Node.fromJSON(editor.pmSchema, documentState.pm); + + if (editor.prosemirrorState.doc.eq(newNode)) { + // duplicate strict mode render? + // smelly, probably wrong architecture between effect / ref / etc. + return; + } + + console.log("documentState", documentState); + previousDocumentState.current = editor._tiptapEditor.state.doc; + editor.transact((tr) => { + tr.replace(0, tr.doc.content.size, new Slice(newNode.content, 0, 0)); + }); + }, [chat, editor, previewDocument, previousDocumentState]); + + return ( + + {children} + + ); +}; + +export function getMessageWithToolCallId( + chat: Chat, + toolCallId: string, +) { + return chat.messages.find( + (m) => + m.role === "assistant" && + m.parts.some( + (p) => isToolOrDynamicToolUIPart(p) && p.toolCallId === toolCallId, + ), + ); +} +// function getDocumentStateBeforeToolCall( +// chat: Chat, +// toolCallId: string, +// ) { +// let index = chat.messages.findIndex( +// (m) => +// m.role === "assistant" && +// m.parts.some( +// (p) => isToolOrDynamicToolUIPart(p) && p.toolCallId === toolCallId, +// ), +// ); +// if (index === -1) { +// throw new Error("Tool call not found"); +// } +// return chat.messages[index].metadata?.documentState; +// // while (index >= 0) { +// // const message = chat.messages[index]; +// // // find the last document-state message before the tool call and get the metadata +// // if ( +// // message.role === "assistant" && +// // message.id.startsWith("document-state-") +// // ) { +// // const state = (message.metadata as any)?.documentState; +// // return state; +// // } + +// // index--; +// // } +// throw new Error("Document state not found"); +// } + +type ChatContextType = { + chat: BlockNoteAISDKChat; + hasSuggestions: boolean; + setEditor: (editor: BlockNoteEditor) => void; + acceptChanges: () => void; + rejectChanges: () => void; + setPreviewDocument: (toolCallId: string | undefined) => void; + transport: ChatTransport; +}; + +const ChatContext = createContext(undefined); + +export const useChatContext = () => { + const chat = useContext(ChatContext); + if (!chat) { + throw new Error("ChatContext not found"); + } + return chat; +}; diff --git a/agent-demo/components/DocumentList.tsx b/agent-demo/components/DocumentList.tsx new file mode 100644 index 0000000000..556a0cd3ad --- /dev/null +++ b/agent-demo/components/DocumentList.tsx @@ -0,0 +1,104 @@ +import { + ThreadListItemPrimitive, + ThreadListPrimitive, + useAssistantState, +} from "@assistant-ui/react"; +import { ArchiveIcon, PlusIcon } from "lucide-react"; +import type { FC } from "react"; + +import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { Button } from "@/components/ui/button"; +import { Skeleton } from "@/components/ui/skeleton"; + +export const DocumentList: FC = () => { + return ( + + + + + ); +}; + +const DocumentListNew: FC = () => { + return ( + { + window.alert("Not implemented"); + }} + > + + + ); +}; + +const DocumentListItems: FC = () => { + const isLoading = useAssistantState(({ threads }) => threads.isLoading); + + if (isLoading) { + return ; + } + + return ( + + ); +}; + +const DocumentListSkeleton: FC = () => { + return ( + <> + {Array.from({ length: 5 }, (_, i) => ( +
+ +
+ ))} + + ); +}; + +const DocumentListItem: FC = () => { + return ( + + + + + + + ); +}; + +const DocumentListItemTitle: FC = () => { + return ( + + + + ); +}; + +const DocumentListItemArchive: FC = () => { + return ( + + + + + + ); +}; diff --git a/agent-demo/components/assistant-ui/assistant-modal.tsx b/agent-demo/components/assistant-ui/assistant-modal.tsx new file mode 100644 index 0000000000..08f2a22769 --- /dev/null +++ b/agent-demo/components/assistant-ui/assistant-modal.tsx @@ -0,0 +1,65 @@ +"use client"; + +import { BotIcon, ChevronDownIcon } from "lucide-react"; + +import { AssistantModalPrimitive } from "@assistant-ui/react"; +import { type FC, forwardRef } from "react"; + +import { Thread } from "@/components/assistant-ui/thread"; +import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; + +export const AssistantModal: FC = () => { + return ( + + + + + + + + + + + ); +}; + +type AssistantModalButtonProps = { "data-state"?: "open" | "closed" }; + +const AssistantModalButton = forwardRef< + HTMLButtonElement, + AssistantModalButtonProps +>(({ "data-state": state, ...rest }, ref) => { + const tooltip = state === "open" ? "Close Assistant" : "Open Assistant"; + + return ( + + + + + {tooltip} + + ); +}); + +AssistantModalButton.displayName = "AssistantModalButton"; diff --git a/agent-demo/components/assistant-ui/attachment.tsx b/agent-demo/components/assistant-ui/attachment.tsx new file mode 100644 index 0000000000..4f513e43ff --- /dev/null +++ b/agent-demo/components/assistant-ui/attachment.tsx @@ -0,0 +1,235 @@ +"use client"; + +import { PropsWithChildren, useEffect, useState, type FC } from "react"; +import Image from "next/image"; +import { XIcon, PlusIcon, FileText } from "lucide-react"; +import { + AttachmentPrimitive, + ComposerPrimitive, + MessagePrimitive, + useAssistantState, + useAssistantApi, +} from "@assistant-ui/react"; +import { useShallow } from "zustand/shallow"; +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from "@/components/ui/tooltip"; +import { + Dialog, + DialogTitle, + DialogContent, + DialogTrigger, +} from "@/components/ui/dialog"; +import { Avatar, AvatarImage, AvatarFallback } from "@/components/ui/avatar"; +import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { cn } from "@/lib/utils"; + +const useFileSrc = (file: File | undefined) => { + const [src, setSrc] = useState(undefined); + + useEffect(() => { + if (!file) { + setSrc(undefined); + return; + } + + const objectUrl = URL.createObjectURL(file); + setSrc(objectUrl); + + return () => { + URL.revokeObjectURL(objectUrl); + }; + }, [file]); + + return src; +}; + +const useAttachmentSrc = () => { + const { file, src } = useAssistantState( + useShallow(({ attachment }): { file?: File; src?: string } => { + if (attachment.type !== "image") return {}; + if (attachment.file) return { file: attachment.file }; + const src = attachment.content?.filter((c) => c.type === "image")[0] + ?.image; + if (!src) return {}; + return { src }; + }), + ); + + return useFileSrc(file) ?? src; +}; + +type AttachmentPreviewProps = { + src: string; +}; + +const AttachmentPreview: FC = ({ src }) => { + const [isLoaded, setIsLoaded] = useState(false); + return ( + Image Preview setIsLoaded(true)} + priority={false} + /> + ); +}; + +const AttachmentPreviewDialog: FC = ({ children }) => { + const src = useAttachmentSrc(); + + if (!src) return children; + + return ( + + + {children} + + + + Image Attachment Preview + +
+ +
+ + + ); +}; + +const AttachmentThumb: FC = () => { + const isImage = useAssistantState( + ({ attachment }) => attachment.type === "image", + ); + const src = useAttachmentSrc(); + + return ( + + + + + + + ); +}; + +const AttachmentUI: FC = () => { + const api = useAssistantApi(); + const isComposer = api.attachment.source === "composer"; + + const isImage = useAssistantState( + ({ attachment }) => attachment.type === "image", + ); + const typeLabel = useAssistantState(({ attachment }) => { + const type = attachment.type; + switch (type) { + case "image": + return "Image"; + case "document": + return "Document"; + case "file": + return "File"; + default: + const _exhaustiveCheck: never = type; + throw new Error(`Unknown attachment type: ${_exhaustiveCheck}`); + } + }); + + return ( + + #attachment-tile]:size-24", + )} + > + + +
+ +
+ + + {isComposer && } + + + + + + ); +}; + +const AttachmentRemove: FC = () => { + return ( + + + + + + ); +}; + +export const UserMessageAttachments: FC = () => { + return ( +
+ +
+ ); +}; + +export const ComposerAttachments: FC = () => { + return ( +
+ +
+ ); +}; + +export const ComposerAddAttachment: FC = () => { + return ( + + + + + + ); +}; diff --git a/agent-demo/components/assistant-ui/markdown-text.tsx b/agent-demo/components/assistant-ui/markdown-text.tsx new file mode 100644 index 0000000000..5f3ee56982 --- /dev/null +++ b/agent-demo/components/assistant-ui/markdown-text.tsx @@ -0,0 +1,228 @@ +"use client"; + +import "@assistant-ui/react-markdown/styles/dot.css"; + +import { + type CodeHeaderProps, + MarkdownTextPrimitive, + unstable_memoizeMarkdownComponents as memoizeMarkdownComponents, + useIsMarkdownCodeBlock, +} from "@assistant-ui/react-markdown"; +import remarkGfm from "remark-gfm"; +import { type FC, memo, useState } from "react"; +import { CheckIcon, CopyIcon } from "lucide-react"; + +import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; +import { cn } from "@/lib/utils"; + +const MarkdownTextImpl = () => { + return ( + + ); +}; + +export const MarkdownText = memo(MarkdownTextImpl); + +const CodeHeader: FC = ({ language, code }) => { + const { isCopied, copyToClipboard } = useCopyToClipboard(); + const onCopy = () => { + if (!code || isCopied) return; + copyToClipboard(code); + }; + + return ( +
+ + {language} + + + {!isCopied && } + {isCopied && } + +
+ ); +}; + +const useCopyToClipboard = ({ + copiedDuration = 3000, +}: { + copiedDuration?: number; +} = {}) => { + const [isCopied, setIsCopied] = useState(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 }) => ( +

+ ), + h2: ({ className, ...props }) => ( +

+ ), + h3: ({ className, ...props }) => ( +

+ ), + h4: ({ className, ...props }) => ( +

+ ), + h5: ({ className, ...props }) => ( +

+ ), + h6: ({ className, ...props }) => ( +
+ ), + p: ({ className, ...props }) => ( +

+ ), + a: ({ className, ...props }) => ( + + ), + blockquote: ({ className, ...props }) => ( +

+ ), + ul: ({ className, ...props }) => ( +
    li]:mt-2", className)} + {...props} + /> + ), + ol: ({ className, ...props }) => ( +
      li]:mt-2", className)} + {...props} + /> + ), + hr: ({ className, ...props }) => ( +
      + ), + table: ({ className, ...props }) => ( + + ), + th: ({ className, ...props }) => ( + td:first-child]:rounded-bl-lg [&:last-child>td:last-child]:rounded-br-lg", + className, + )} + {...props} + /> + ), + sup: ({ className, ...props }) => ( + a]:text-xs [&>a]:no-underline", className)} + {...props} + /> + ), + pre: ({ className, ...props }) => ( + 
      +  ),
+  code: function Code({ className, ...props }) {
+    const isCodeBlock = useIsMarkdownCodeBlock();
+    return (
+      
+    );
+  },
+  CodeHeader,
+});
diff --git a/agent-demo/components/assistant-ui/thread-list.tsx b/agent-demo/components/assistant-ui/thread-list.tsx
new file mode 100644
index 0000000000..e6ea558479
--- /dev/null
+++ b/agent-demo/components/assistant-ui/thread-list.tsx
@@ -0,0 +1,95 @@
+import {
+  ThreadListItemPrimitive,
+  ThreadListPrimitive,
+  useAssistantState,
+} from "@assistant-ui/react";
+import { ArchiveIcon, PlusIcon } from "lucide-react";
+import type { FC } from "react";
+
+import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+import { Button } from "@/components/ui/button";
+import { Skeleton } from "@/components/ui/skeleton";
+
+export const ThreadList: FC = () => {
+  return (
+    
+      
+      
+    
+  );
+};
+
+const ThreadListNew: FC = () => {
+  return (
+    
+      
+    
+  );
+};
+
+const ThreadListItems: FC = () => {
+  const isLoading = useAssistantState(({ threads }) => threads.isLoading);
+
+  if (isLoading) {
+    return ;
+  }
+
+  return ;
+};
+
+const ThreadListSkeleton: FC = () => {
+  return (
+    <>
+      {Array.from({ length: 5 }, (_, i) => (
+        
      
+          
+        
      
+      ))}
+    
+  );
+};
+
+const ThreadListItem: FC = () => {
+  return (
+    
+      
+        
+      
+      
+    
+  );
+};
+
+const ThreadListItemTitle: FC = () => {
+  return (
+    
+      
+    
+  );
+};
+
+const ThreadListItemArchive: FC = () => {
+  return (
+    
+      
+        
+      
+    
+  );
+};
diff --git a/agent-demo/components/assistant-ui/thread.tsx b/agent-demo/components/assistant-ui/thread.tsx
new file mode 100644
index 0000000000..38afe8748e
--- /dev/null
+++ b/agent-demo/components/assistant-ui/thread.tsx
@@ -0,0 +1,446 @@
+import {
+  ArrowDownIcon,
+  ArrowUpIcon,
+  CheckIcon,
+  ChevronLeftIcon,
+  ChevronRightIcon,
+  CopyIcon,
+  PencilIcon,
+  Square,
+} from "lucide-react";
+
+import {
+  ActionBarPrimitive,
+  BranchPickerPrimitive,
+  ComposerPrimitive,
+  ErrorPrimitive,
+  MessagePrimitive,
+  ThreadPrimitive,
+  useAssistantState,
+} from "@assistant-ui/react";
+
+import { LazyMotion, MotionConfig, domAnimation } from "motion/react";
+import * as m from "motion/react-m";
+import type { FC } from "react";
+
+import {
+  ComposerAddAttachment,
+  ComposerAttachments,
+  UserMessageAttachments,
+} from "@/components/assistant-ui/attachment";
+import { MarkdownText } from "@/components/assistant-ui/markdown-text";
+import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
+import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
+import { useChatContext } from "@/components/ChatContext";
+import { Button } from "@/components/ui/button";
+import { cn } from "@/lib/utils";
+
+export const Thread: FC = () => {
+  return (
+    
+      
+        
+          
+            
+              
+            
+
+            
+
+            
+              
      
+            
+            
+            
+          
+        
+      
+    
+  );
+};
+
+const ThreadScrollToBottom: FC = () => {
+  return (
+    
+      
+        
+      
+    
+  );
+};
+
+const ThreadWelcome: FC = () => {
+  return (
+    
      
+      
      
+        
      
+          
+            Hello there!
+          
+          
+            How can I help you today?
+          
+        
      
+      
      
+      
+    
      
+  );
+};
+
+const ThreadSuggestions: FC = () => {
+  return (
+    
      
+      {[
+        {
+          title: "Summarize",
+          label: "the document",
+          action: "Create a short summary of the document",
+        },
+        {
+          title: "Add action items",
+          label: "to the document",
+          action: "Add action items to the document",
+        },
+      ].map((suggestedAction, index) => (
+        
+          
+            
+          
+        
+      ))}
+    
      
+  );
+};
+
+// Added
+const ThreadAcceptReject: FC = () => {
+  const ctx = useChatContext();
+  if (!ctx.hasSuggestions) {
+    return null;
+  }
+  return (
+    
      
+      {[
+        {
+          title: "Accept",
+          label: "suggested changes",
+          action: () => {
+            ctx.acceptChanges();
+          },
+        },
+        {
+          title: "Reject",
+          label: "revert the document",
+          action: () => {
+            ctx.rejectChanges();
+          },
+        },
+      ].map((suggestedAction, index) => (
+        
+          {/*  */}
+          
+          {/*  */}
+        
+      ))}
+    
      
+  );
+};
+
+const Composer: FC = () => {
+  return (
+    
      
+      
+      
+        
+        
+        
+      
+    
      
+  );
+};
+
+const ComposerAction: FC = () => {
+  return (
+    
      
+      
+
+      
+        
+          
+            
+          
+        
+      
+
+      
+        
+          
+        
+      
+    
      
+  );
+};
+
+const MessageError: FC = () => {
+  return (
+    
+      
+        
+      
+    
+  );
+};
+
+const AssistantMessage: FC = () => {
+  // added to hide specific messages
+  const isDocumentStateMessage = useAssistantState(
+    ({ message }) =>
+      message.parts.length > 0 &&
+      message.parts[0].type === "tool-call" &&
+      message.parts[0].toolName === "document-state",
+  );
+
+  if (isDocumentStateMessage) {
+    return null;
+  }
+
+  return (
+    
+      
      
+        
      
+          
+          
+        
      
+
+        {/* 
      
+          
+          
+        
       */}
+      
      
+          
+  );
+};
+
+const AssistantActionBar: FC = () => {
+  return (
+    
+      
+        
+          
+            
+          
+          
+            
+          
+        
+      
+      {/* 
+        
+          
+        
+       */}
+    
+  );
+};
+
+const UserMessage: FC = () => {
+  return (
+    
+      
      
+        
+
+        
      
+          
      
+            
+          
      
+          {/* 
      
+            
+          
       */}
+        
      
+
+        
+      
      
+          
+  );
+};
+
+const UserActionBar: FC = () => {
+  return (
+    
+      
+        
+          
+        
+      
+    
+  );
+};
+
+const EditComposer: FC = () => {
+  return (
+    
      
+      
+        
+
+        
      
+          
+            
+          
+          
+            
+          
+        
      
+            
+    
      
+  );
+};
+
+const BranchPicker: FC = ({
+  className,
+  ...rest
+}) => {
+  return (
+    
+      
+        
+          
+        
+      
+      
+         / 
+      
+      
+        
+          
+        
+      
+    
+  );
+};
diff --git a/agent-demo/components/assistant-ui/threadlist-sidebar.tsx b/agent-demo/components/assistant-ui/threadlist-sidebar.tsx
new file mode 100644
index 0000000000..656ae4280b
--- /dev/null
+++ b/agent-demo/components/assistant-ui/threadlist-sidebar.tsx
@@ -0,0 +1,73 @@
+import { DocumentList } from "@/components/DocumentList";
+import {
+  Sidebar,
+  SidebarContent,
+  SidebarHeader,
+  SidebarMenu,
+  SidebarMenuButton,
+  SidebarMenuItem,
+  SidebarRail,
+} from "@/components/ui/sidebar";
+import { MessagesSquare } from "lucide-react";
+import Link from "next/link";
+import * as React from "react";
+
+export function ThreadListSidebar({
+  ...props
+}: React.ComponentProps) {
+  return (
+    
+      
+        
      
+          
+            
+              
+                
+                  
      
+                    
+                  
      
+                  
      
+                    
+                      BlockNote AI
+                    
+                  
      
+                
+                    
+                  
+                
+        
      
+            
+      
+        {/*  */}
+        
+      
+      
+      {/* 
+        
+          
+            
+              
+                
      
+                  
+                
      
+                
      
+                  
+                    GitHub
+                  
+                  View Source
+                
      
+              
+                  
+                
+              
+             */}
+          
+  );
+}
diff --git a/agent-demo/components/assistant-ui/tool-fallback.tsx b/agent-demo/components/assistant-ui/tool-fallback.tsx
new file mode 100644
index 0000000000..aca40305b5
--- /dev/null
+++ b/agent-demo/components/assistant-ui/tool-fallback.tsx
@@ -0,0 +1,46 @@
+import type { ToolCallMessagePartComponent } from "@assistant-ui/react";
+import { CheckIcon, ChevronDownIcon, ChevronUpIcon } from "lucide-react";
+import { useState } from "react";
+import { Button } from "@/components/ui/button";
+
+export const ToolFallback: ToolCallMessagePartComponent = ({
+  toolName,
+  argsText,
+  result,
+}) => {
+  const [isCollapsed, setIsCollapsed] = useState(true);
+  return (
+    
      
+      
      
+        
+        
      
+          Used tool: {toolName}
+        
      
+        
+      
      
+      {!isCollapsed && (
+        
      
+          
      
+            
      +              {argsText}
+            
      
+          
      
+          {result !== undefined && (
+            
      
+              
      
+                Result:
+              
      
+              
      +                {typeof result === "string"
+                  ? result
+                  : JSON.stringify(result, null, 2)}
+              
      
+            
      
+          )}
+        
      
+      )}
+    
      
+  );
+};
diff --git a/agent-demo/components/assistant-ui/tooltip-icon-button.tsx b/agent-demo/components/assistant-ui/tooltip-icon-button.tsx
new file mode 100644
index 0000000000..54b5fa27a7
--- /dev/null
+++ b/agent-demo/components/assistant-ui/tooltip-icon-button.tsx
@@ -0,0 +1,42 @@
+"use client";
+
+import { ComponentPropsWithRef, forwardRef } from "react";
+import { Slottable } from "@radix-ui/react-slot";
+
+import {
+  Tooltip,
+  TooltipContent,
+  TooltipTrigger,
+} from "@/components/ui/tooltip";
+import { Button } from "@/components/ui/button";
+import { cn } from "@/lib/utils";
+
+export type TooltipIconButtonProps = ComponentPropsWithRef & {
+  tooltip: string;
+  side?: "top" | "bottom" | "left" | "right";
+};
+
+export const TooltipIconButton = forwardRef<
+  HTMLButtonElement,
+  TooltipIconButtonProps
+>(({ children, tooltip, side = "bottom", className, ...rest }, ref) => {
+  return (
+    
+      
+        
+      
+      {tooltip}
+    
+  );
+});
+
+TooltipIconButton.displayName = "TooltipIconButton";
diff --git a/agent-demo/components/checkpoints.ts b/agent-demo/components/checkpoints.ts
new file mode 100644
index 0000000000..969e579074
--- /dev/null
+++ b/agent-demo/components/checkpoints.ts
@@ -0,0 +1,53 @@
+import { Chat } from "@ai-sdk/react";
+import { BlockNoteEditor } from "@blocknote/core";
+import { isToolUIPart } from "ai";
+
+export function trackToolCallCheckpoints(
+  chat: Chat,
+  editor: BlockNoteEditor,
+) {
+  const checkpoints = new Map<
+    string,
+    {
+      bn: any;
+      pm: any;
+    }
+  >();
+
+  /**
+   * Note: this approach has a flaw that "output-available" (in chatHandlers)
+   * is only set once all tool calls have been completed.
+   *
+   * If the LLM uses parallel tool calling, all checkpoints will be the same
+   *
+   * It would be nicer to set "output-available" at the correct time,
+   * and maybe also integrate this checkpoint logic there.
+   *
+   * However, this would require diving deep on streams, for now we consider this edge-case ok
+   */
+  const dispose = chat["~registerMessagesCallback"](() => {
+    for (const message of chat.messages) {
+      for (const part of message.parts) {
+        if (
+          isToolUIPart(part) &&
+          part.type === "tool-applyDocumentOperations"
+        ) {
+          if (
+            part.state === "output-available" &&
+            part.output !== "" && // this  is an assistant-ui thing
+            !checkpoints.has(part.toolCallId)
+          ) {
+            checkpoints.set(part.toolCallId, {
+              bn: editor.document,
+              pm: editor.prosemirrorState.doc.toJSON(),
+            });
+          }
+        }
+      }
+    }
+  });
+  return {
+    checkpoints,
+    dispose,
+  };
+}
diff --git a/agent-demo/components/ui/avatar.tsx b/agent-demo/components/ui/avatar.tsx
new file mode 100644
index 0000000000..f7923d49a9
--- /dev/null
+++ b/agent-demo/components/ui/avatar.tsx
@@ -0,0 +1,53 @@
+"use client";
+
+import * as React from "react";
+import * as AvatarPrimitive from "@radix-ui/react-avatar";
+
+import { cn } from "@/lib/utils";
+
+function Avatar({
+  className,
+  ...props
+}: React.ComponentProps) {
+  return (
+    
+  );
+}
+
+function AvatarImage({
+  className,
+  ...props
+}: React.ComponentProps) {
+  return (
+    
+  );
+}
+
+function AvatarFallback({
+  className,
+  ...props
+}: React.ComponentProps) {
+  return (
+    
+  );
+}
+
+export { Avatar, AvatarImage, AvatarFallback };
diff --git a/agent-demo/components/ui/breadcrumb.tsx b/agent-demo/components/ui/breadcrumb.tsx
new file mode 100644
index 0000000000..a0df2b3669
--- /dev/null
+++ b/agent-demo/components/ui/breadcrumb.tsx
@@ -0,0 +1,109 @@
+import * as React from "react";
+import { Slot } from "@radix-ui/react-slot";
+import { ChevronRight, MoreHorizontal } from "lucide-react";
+
+import { cn } from "@/lib/utils";
+
+function Breadcrumb({ ...props }: React.ComponentProps<"nav">) {
+  return
      + ), + td: ({ className, ...props }) => ( + + ), + tr: ({ className, ...props }) => ( +