diff --git a/.changeset/nice-garlics-repeat.md b/.changeset/nice-garlics-repeat.md
new file mode 100644
index 000000000..2ce66a1b5
--- /dev/null
+++ b/.changeset/nice-garlics-repeat.md
@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+feat: use llamaindex chat-ui for nextjs frontend
diff --git a/templates/types/streaming/nextjs/app/components/chat-section.tsx b/templates/types/streaming/nextjs/app/components/chat-section.tsx
index e7e489ba1..483ca7bb3 100644
--- a/templates/types/streaming/nextjs/app/components/chat-section.tsx
+++ b/templates/types/streaming/nextjs/app/components/chat-section.tsx
@@ -1,57 +1,26 @@
 "use client";
 
+import { ChatSection as ChatSectionUI } from "@llamaindex/chat-ui";
+import "@llamaindex/chat-ui/styles/code.css";
+import "@llamaindex/chat-ui/styles/katex.css";
 import { useChat } from "ai/react";
-import { useState } from "react";
-import { ChatInput, ChatMessages } from "./ui/chat";
+import CustomChatInput from "./ui/chat/chat-input";
+import CustomChatMessages from "./ui/chat/chat-messages";
 import { useClientConfig } from "./ui/chat/hooks/use-config";
 
 export default function ChatSection() {
   const { backend } = useClientConfig();
-  const [requestData, setRequestData] = useState<any>();
-  const {
-    messages,
-    input,
-    isLoading,
-    handleSubmit,
-    handleInputChange,
-    reload,
-    stop,
-    append,
-    setInput,
-  } = useChat({
-    body: { data: requestData },
+  const handler = useChat({
     api: `${backend}/api/chat`,
-    headers: {
-      "Content-Type": "application/json", // using JSON because of vercel/ai 2.2.26
-    },
     onError: (error: unknown) => {
       if (!(error instanceof Error)) throw error;
-      const message = JSON.parse(error.message);
-      alert(message.detail);
+      alert(JSON.parse(error.message).detail);
     },
-    sendExtraMessageFields: true,
   });
-
   return (
-    <div className="space-y-4 w-full h-full flex flex-col">
-      <ChatMessages
-        messages={messages}
-        isLoading={isLoading}
-        reload={reload}
-        stop={stop}
-        append={append}
-      />
-      <ChatInput
-        input={input}
-        handleSubmit={handleSubmit}
-        handleInputChange={handleInputChange}
-        isLoading={isLoading}
-        messages={messages}
-        append={append}
-        setInput={setInput}
-        requestParams={{ params: requestData }}
-        setRequestData={setRequestData}
-      />
-    </div>
+    <ChatSectionUI handler={handler} className="w-full h-full">
+      <CustomChatMessages />
+      <CustomChatInput />
+    </ChatSectionUI>
   );
 }
diff --git a/templates/types/streaming/nextjs/app/components/ui/README-template.md b/templates/types/streaming/nextjs/app/components/ui/README-template.md
deleted file mode 100644
index ebfcf48c9..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/README-template.md
+++ /dev/null
@@ -1 +0,0 @@
-Using the chat component from https://github.com/marcusschiesser/ui (based on https://ui.shadcn.com/)
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-actions.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-actions.tsx
deleted file mode 100644
index 151ef61a9..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-actions.tsx
+++ /dev/null
@@ -1,28 +0,0 @@
-import { PauseCircle, RefreshCw } from "lucide-react";
-
-import { Button } from "../button";
-import { ChatHandler } from "./chat.interface";
-
-export default function ChatActions(
-  props: Pick<ChatHandler, "stop" | "reload"> & {
-    showReload?: boolean;
-    showStop?: boolean;
-  },
-) {
-  return (
-    <div className="space-x-4">
-      {props.showStop && (
-        <Button variant="outline" size="sm" onClick={props.stop}>
-          <PauseCircle className="mr-2 h-4 w-4" />
-          Stop generating
-        </Button>
-      )}
-      {props.showReload && (
-        <Button variant="outline" size="sm" onClick={props.reload}>
-          <RefreshCw className="mr-2 h-4 w-4" />
-          Regenerate
-        </Button>
-      )}
-    </div>
-  );
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-avatar.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-avatar.tsx
similarity index 78%
rename from templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-avatar.tsx
rename to templates/types/streaming/nextjs/app/components/ui/chat/chat-avatar.tsx
index ce04e306a..cfa307cbf 100644
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-avatar.tsx
+++ b/templates/types/streaming/nextjs/app/components/ui/chat/chat-avatar.tsx
@@ -1,8 +1,10 @@
+import { useChatMessage } from "@llamaindex/chat-ui";
 import { User2 } from "lucide-react";
 import Image from "next/image";
 
-export default function ChatAvatar({ role }: { role: string }) {
-  if (role === "user") {
+export function ChatMessageAvatar() {
+  const { message } = useChatMessage();
+  if (message.role === "user") {
     return (
       <div className="flex h-8 w-8 shrink-0 select-none items-center justify-center rounded-md border bg-background shadow">
         <User2 className="h-4 w-4" />
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-input.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-input.tsx
index 0e5c318b3..8800bfef7 100644
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-input.tsx
+++ b/templates/types/streaming/nextjs/app/components/ui/chat/chat-input.tsx
@@ -1,34 +1,13 @@
-import { JSONValue } from "ai";
-import React from "react";
-import { DocumentFile } from ".";
-import { Button } from "../button";
-import { DocumentPreview } from "../document-preview";
-import FileUploader from "../file-uploader";
-import { Textarea } from "../textarea";
-import UploadImagePreview from "../upload-image-preview";
-import { ChatHandler } from "./chat.interface";
-import { useFile } from "./hooks/use-file";
-import { LlamaCloudSelector } from "./widgets/LlamaCloudSelector";
+"use client";
 
-const ALLOWED_EXTENSIONS = ["png", "jpg", "jpeg", "csv", "pdf", "txt", "docx"];
+import { ChatInput, useChatUI, useFile } from "@llamaindex/chat-ui";
+import { DocumentPreview, ImagePreview } from "@llamaindex/chat-ui/widgets";
+import { LlamaCloudSelector } from "./custom/llama-cloud-selector";
+import { useClientConfig } from "./hooks/use-config";
 
-export default function ChatInput(
-  props: Pick<
-    ChatHandler,
-    | "isLoading"
-    | "input"
-    | "onFileUpload"
-    | "onFileError"
-    | "handleSubmit"
-    | "handleInputChange"
-    | "messages"
-    | "setInput"
-    | "append"
-  > & {
-    requestParams?: any;
-    setRequestData?: React.Dispatch<any>;
-  },
-) {
+export default function CustomChatInput() {
+  const { requestData, isLoading, input } = useChatUI();
+  const { backend } = useClientConfig();
   const {
     imageUrl,
     setImageUrl,
@@ -37,107 +16,65 @@ export default function ChatInput(
     removeDoc,
     reset,
     getAnnotations,
-  } = useFile();
-
-  // default submit function does not handle including annotations in the message
-  // so we need to use append function to submit new message with annotations
-  const handleSubmitWithAnnotations = (
-    e: React.FormEvent<HTMLFormElement>,
-    annotations: JSONValue[] | undefined,
-  ) => {
-    e.preventDefault();
-    props.append!({
-      content: props.input,
-      role: "user",
-      createdAt: new Date(),
-      annotations,
-    });
-    props.setInput!("");
-  };
-
-  const onSubmit = (e: React.FormEvent<HTMLFormElement>) => {
-    e.preventDefault();
-    const annotations = getAnnotations();
-    if (annotations.length) {
-      handleSubmitWithAnnotations(e, annotations);
-      return reset();
-    }
-    props.handleSubmit(e);
-  };
+  } = useFile({ uploadAPI: `${backend}/api/chat/upload` });
 
+  /**
+   * Handles file uploads. Overwrite to hook into the file upload behavior.
+   * @param file The file to upload
+   */
   const handleUploadFile = async (file: File) => {
+    // There's already an image uploaded, only allow one image at a time
     if (imageUrl) {
       alert("You can only upload one image at a time.");
       return;
     }
+
     try {
-      await uploadFile(file, props.requestParams);
-      props.onFileUpload?.(file);
+      // Upload the file and send with it the current request data
+      await uploadFile(file, requestData);
     } catch (error: any) {
-      const onFileUploadError = props.onFileError || window.alert;
-      onFileUploadError(error.message);
+      // Show error message if upload fails
+      alert(error.message);
     }
   };
 
-  const handleKeyDown = (e: React.KeyboardEvent<HTMLTextAreaElement>) => {
-    if (e.key === "Enter" && !e.shiftKey) {
-      e.preventDefault();
-      onSubmit(e as unknown as React.FormEvent<HTMLFormElement>);
-    }
-  };
+  // Get references to the upload files in message annotations format, see https://github.com/run-llama/chat-ui/blob/main/packages/chat-ui/src/hook/use-file.tsx#L56
+  const annotations = getAnnotations();
 
   return (
-    <form
-      onSubmit={onSubmit}
-      className="rounded-xl bg-white p-4 shadow-xl space-y-4 shrink-0"
+    <ChatInput
+      className="shadow-xl rounded-xl"
+      resetUploadedFiles={reset}
+      annotations={annotations}
     >
-      {imageUrl && (
-        <UploadImagePreview url={imageUrl} onRemove={() => setImageUrl(null)} />
-      )}
-      {files.length > 0 && (
-        <div className="flex gap-4 w-full overflow-auto py-2">
-          {files.map((file: DocumentFile) => (
-            <DocumentPreview
-              key={file.id}
-              file={file}
-              onRemove={() => removeDoc(file)}
-            />
-          ))}
-        </div>
-      )}
-      <div className="flex w-full items-start justify-between gap-4 ">
-        <Textarea
-          id="chat-input"
-          autoFocus
-          name="message"
-          placeholder="Type a message"
-          className="flex-1 min-h-0 h-[40px]"
-          value={props.input}
-          onChange={props.handleInputChange}
-          onKeyDown={handleKeyDown}
-        />
-        <FileUploader
-          onFileUpload={handleUploadFile}
-          onFileError={props.onFileError}
-          config={{
-            allowedExtensions: ALLOWED_EXTENSIONS,
-            disabled: props.isLoading,
-            multiple: true,
-          }}
-        />
-        {process.env.NEXT_PUBLIC_USE_LLAMACLOUD === "true" &&
-          props.setRequestData && (
-            <LlamaCloudSelector setRequestData={props.setRequestData} />
-          )}
-        <Button
-          type="submit"
+      <div>
+        {/* Image preview section */}
+        {imageUrl && (
+          <ImagePreview url={imageUrl} onRemove={() => setImageUrl(null)} />
+        )}
+        {/* Document previews section */}
+        {files.length > 0 && (
+          <div className="flex gap-4 w-full overflow-auto py-2">
+            {files.map((file) => (
+              <DocumentPreview
+                key={file.id}
+                file={file}
+                onRemove={() => removeDoc(file)}
+              />
+            ))}
+          </div>
+        )}
+      </div>
+      <ChatInput.Form>
+        <ChatInput.Field />
+        <ChatInput.Upload onUpload={handleUploadFile} />
+        <LlamaCloudSelector />
+        <ChatInput.Submit
           disabled={
-            props.isLoading || (!props.input.trim() && files.length === 0)
+            isLoading || (!input.trim() && files.length === 0 && !imageUrl)
           }
-        >
-          Send message
-        </Button>
-      </div>
-    </form>
+        />
+      </ChatInput.Form>
+    </ChatInput>
   );
 }
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message-content.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message-content.tsx
new file mode 100644
index 000000000..d4ddfbf47
--- /dev/null
+++ b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message-content.tsx
@@ -0,0 +1,30 @@
+import {
+  ChatMessage,
+  ContentPosition,
+  getSourceAnnotationData,
+  useChatMessage,
+} from "@llamaindex/chat-ui";
+import { Markdown } from "./custom/markdown";
+import { ToolAnnotations } from "./tools/chat-tools";
+
+export function ChatMessageContent() {
+  const { message } = useChatMessage();
+  const customContent = [
+    {
+      // override the default markdown component
+      position: ContentPosition.MARKDOWN,
+      component: (
+        <Markdown
+          content={message.content}
+          sources={getSourceAnnotationData(message.annotations)?.[0]}
+        />
+      ),
+    },
+    {
+      // add the tool annotations after events
+      position: ContentPosition.AFTER_EVENTS,
+      component: <ToolAnnotations message={message} />,
+    },
+  ];
+  return <ChatMessage.Content content={customContent} />;
+}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-agent-events.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-agent-events.tsx
deleted file mode 100644
index a385754bc..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-agent-events.tsx
+++ /dev/null
@@ -1,222 +0,0 @@
-import { icons, LucideIcon } from "lucide-react";
-import { useMemo } from "react";
-import { Button } from "../../button";
-import {
-  Drawer,
-  DrawerClose,
-  DrawerContent,
-  DrawerHeader,
-  DrawerTitle,
-  DrawerTrigger,
-} from "../../drawer";
-import { Progress } from "../../progress";
-import { AgentEventData, ProgressData } from "../index";
-import Markdown from "./markdown";
-
-const AgentIcons: Record<string, LucideIcon> = {
-  bot: icons.Bot,
-  researcher: icons.ScanSearch,
-  writer: icons.PenLine,
-  reviewer: icons.MessageCircle,
-  publisher: icons.BookCheck,
-};
-
-type StepText = {
-  text: string;
-};
-
-type StepProgress = {
-  text: string;
-  progress: ProgressData;
-};
-
-type MergedEvent = {
-  agent: string;
-  icon: LucideIcon;
-  steps: Array<StepText | StepProgress>;
-};
-
-export function ChatAgentEvents({
-  data,
-  isFinished,
-}: {
-  data: AgentEventData[];
-  isFinished: boolean;
-}) {
-  const events = useMemo(() => mergeAdjacentEvents(data), [data]);
-  return (
-    <div className="pl-2">
-      <div className="text-sm space-y-4">
-        {events.map((eventItem, index) => (
-          <AgentEventContent
-            key={index}
-            event={eventItem}
-            isLast={index === events.length - 1}
-            isFinished={isFinished}
-          />
-        ))}
-      </div>
-    </div>
-  );
-}
-
-const MAX_TEXT_LENGTH = 150;
-
-function TextContent({ agent, step }: { agent: string; step: StepText }) {
-  const { displayText, showMore } = useMemo(
-    () => ({
-      displayText: step.text.slice(0, MAX_TEXT_LENGTH),
-      showMore: step.text.length > MAX_TEXT_LENGTH,
-    }),
-    [step.text],
-  );
-
-  return (
-    <>
-      <div className="whitespace-break-spaces">
-        {!showMore && <span>{step.text}</span>}
-        {showMore && (
-          <div>
-            <span>{displayText}...</span>
-            <AgentEventDialog content={step.text} title={`Agent "${agent}"`}>
-              <span className="font-semibold underline cursor-pointer ml-2">
-                Show more
-              </span>
-            </AgentEventDialog>
-          </div>
-        )}
-      </div>
-    </>
-  );
-}
-
-function ProgressContent({ step }: { step: StepProgress }) {
-  const progressValue =
-    step.progress.total !== 0
-      ? Math.round(((step.progress.current + 1) / step.progress.total) * 100)
-      : 0;
-
-  return (
-    <div className="space-y-2 mt-2">
-      {step.text && (
-        <p className="text-sm text-muted-foreground">{step.text}</p>
-      )}
-      <Progress value={progressValue} className="w-full h-2" />
-      <p className="text-sm text-muted-foreground">
-        Processing {step.progress.current + 1} of {step.progress.total} steps...
-      </p>
-    </div>
-  );
-}
-
-function AgentEventContent({
-  event,
-  isLast,
-  isFinished,
-}: {
-  event: MergedEvent;
-  isLast: boolean;
-  isFinished: boolean;
-}) {
-  const { agent, steps } = event;
-  const AgentIcon = event.icon;
-  const textSteps = steps.filter((step) => !("progress" in step));
-  const progressSteps = steps.filter(
-    (step) => "progress" in step,
-  ) as StepProgress[];
-  // We only show progress at the last step
-  // TODO: once we support steps that work in parallel, we need to update this
-  const lastProgressStep =
-    progressSteps.length > 0
-      ? progressSteps[progressSteps.length - 1]
-      : undefined;
-
-  return (
-    <div className="flex gap-4 border-b pb-4 items-center fadein-agent">
-      <div className="w-[100px] flex flex-col items-center gap-2">
-        <div className="relative">
-          {isLast && !isFinished && (
-            <div className="absolute -top-0 -right-4">
-              <span className="relative flex h-3 w-3">
-                <span className="animate-ping absolute inline-flex h-full w-full rounded-full bg-sky-400 opacity-75"></span>
-                <span className="relative inline-flex rounded-full h-3 w-3 bg-sky-500"></span>
-              </span>
-            </div>
-          )}
-          <AgentIcon />
-        </div>
-        <span className="font-bold">{agent}</span>
-      </div>
-      {textSteps.length > 0 && (
-        <div className="flex-1">
-          <ul className="list-decimal space-y-2">
-            {textSteps.map((step, index) => (
-              <li key={index}>
-                <TextContent agent={agent} step={step} />
-              </li>
-            ))}
-          </ul>
-          {lastProgressStep && !isFinished && (
-            <ProgressContent step={lastProgressStep} />
-          )}
-        </div>
-      )}
-    </div>
-  );
-}
-
-type AgentEventDialogProps = {
-  title: string;
-  content: string;
-  children: React.ReactNode;
-};
-
-function AgentEventDialog(props: AgentEventDialogProps) {
-  return (
-    <Drawer direction="left">
-      <DrawerTrigger asChild>{props.children}</DrawerTrigger>
-      <DrawerContent className="w-3/5 mt-24 h-full max-h-[96%] ">
-        <DrawerHeader className="flex justify-between">
-          <div className="space-y-2">
-            <DrawerTitle>{props.title}</DrawerTitle>
-          </div>
-          <DrawerClose asChild>
-            <Button variant="outline">Close</Button>
-          </DrawerClose>
-        </DrawerHeader>
-        <div className="m-4 overflow-auto">
-          <Markdown content={props.content} />
-        </div>
-      </DrawerContent>
-    </Drawer>
-  );
-}
-
-function mergeAdjacentEvents(events: AgentEventData[]): MergedEvent[] {
-  const mergedEvents: MergedEvent[] = [];
-
-  for (const event of events) {
-    const lastMergedEvent = mergedEvents[mergedEvents.length - 1];
-
-    const eventStep: StepText | StepProgress = event.data
-      ? ({
-          text: event.text,
-          progress: event.data,
-        } as StepProgress)
-      : ({
-          text: event.text,
-        } as StepText);
-
-    if (lastMergedEvent && lastMergedEvent.agent === event.agent) {
-      lastMergedEvent.steps.push(eventStep);
-    } else {
-      mergedEvents.push({
-        agent: event.agent,
-        steps: [eventStep],
-        icon: AgentIcons[event.agent.toLowerCase()] ?? icons.Bot,
-      });
-    }
-  }
-
-  return mergedEvents;
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-events.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-events.tsx
deleted file mode 100644
index 3dfad75d2..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-events.tsx
+++ /dev/null
@@ -1,50 +0,0 @@
-import { ChevronDown, ChevronRight, Loader2 } from "lucide-react";
-import { useState } from "react";
-import { Button } from "../../button";
-import {
-  Collapsible,
-  CollapsibleContent,
-  CollapsibleTrigger,
-} from "../../collapsible";
-import { EventData } from "../index";
-
-export function ChatEvents({
-  data,
-  isLoading,
-}: {
-  data: EventData[];
-  isLoading: boolean;
-}) {
-  const [isOpen, setIsOpen] = useState(false);
-
-  const buttonLabel = isOpen ? "Hide events" : "Show events";
-
-  const EventIcon = isOpen ? (
-    <ChevronDown className="h-4 w-4" />
-  ) : (
-    <ChevronRight className="h-4 w-4" />
-  );
-
-  return (
-    <div className="border-l-2 border-indigo-400 pl-2">
-      <Collapsible open={isOpen} onOpenChange={setIsOpen}>
-        <CollapsibleTrigger asChild>
-          <Button variant="secondary" className="space-x-2">
-            {isLoading ? <Loader2 className="h-4 w-4 animate-spin" /> : null}
-            <span>{buttonLabel}</span>
-            {EventIcon}
-          </Button>
-        </CollapsibleTrigger>
-        <CollapsibleContent asChild>
-          <div className="mt-4 text-sm space-y-2">
-            {data.map((eventItem, index) => (
-              <div className="whitespace-break-spaces" key={index}>
-                {eventItem.title}
-              </div>
-            ))}
-          </div>
-        </CollapsibleContent>
-      </Collapsible>
-    </div>
-  );
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-files.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-files.tsx
deleted file mode 100644
index 0a5859a96..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-files.tsx
+++ /dev/null
@@ -1,13 +0,0 @@
-import { DocumentPreview } from "../../document-preview";
-import { DocumentFileData } from "../index";
-
-export function ChatFiles({ data }: { data: DocumentFileData }) {
-  if (!data.files.length) return null;
-  return (
-    <div className="flex gap-2 items-center">
-      {data.files.map((file, index) => (
-        <DocumentPreview key={file.id} file={file} />
-      ))}
-    </div>
-  );
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-image.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-image.tsx
deleted file mode 100644
index 2de28c3d3..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-image.tsx
+++ /dev/null
@@ -1,17 +0,0 @@
-import Image from "next/image";
-import { type ImageData } from "../index";
-
-export function ChatImage({ data }: { data: ImageData }) {
-  return (
-    <div className="rounded-md max-w-[200px] shadow-md">
-      <Image
-        src={data.url}
-        width={0}
-        height={0}
-        sizes="100vw"
-        style={{ width: "100%", height: "auto" }}
-        alt=""
-      />
-    </div>
-  );
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-sources.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-sources.tsx
deleted file mode 100644
index c0da0031a..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-sources.tsx
+++ /dev/null
@@ -1,173 +0,0 @@
-import { Check, Copy } from "lucide-react";
-import { useMemo } from "react";
-import { Button } from "../../button";
-import { PreviewCard } from "../../document-preview";
-import {
-  HoverCard,
-  HoverCardContent,
-  HoverCardTrigger,
-} from "../../hover-card";
-import { cn } from "../../lib/utils";
-import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
-import { DocumentFileType, SourceData, SourceNode } from "../index";
-import PdfDialog from "../widgets/PdfDialog";
-
-type Document = {
-  url: string;
-  sources: SourceNode[];
-};
-
-export function ChatSources({ data }: { data: SourceData }) {
-  const documents: Document[] = useMemo(() => {
-    // group nodes by document (a document must have a URL)
-    const nodesByUrl: Record<string, SourceNode[]> = {};
-    data.nodes.forEach((node) => {
-      const key = node.url;
-      nodesByUrl[key] ??= [];
-      nodesByUrl[key].push(node);
-    });
-
-    // convert to array of documents
-    return Object.entries(nodesByUrl).map(([url, sources]) => ({
-      url,
-      sources,
-    }));
-  }, [data.nodes]);
-
-  if (documents.length === 0) return null;
-
-  return (
-    <div className="space-y-2 text-sm">
-      <div className="font-semibold text-lg">Sources:</div>
-      <div className="flex gap-3 flex-wrap">
-        {documents.map((document) => {
-          return <DocumentInfo key={document.url} document={document} />;
-        })}
-      </div>
-    </div>
-  );
-}
-
-function SourceInfo({ node, index }: { node?: SourceNode; index: number }) {
-  if (!node) return <SourceNumberButton index={index} />;
-  return (
-    <HoverCard>
-      <HoverCardTrigger
-        className="cursor-default"
-        onClick={(e) => {
-          e.preventDefault();
-          e.stopPropagation();
-        }}
-      >
-        <SourceNumberButton
-          index={index}
-          className="hover:text-white hover:bg-primary"
-        />
-      </HoverCardTrigger>
-      <HoverCardContent className="w-[400px]">
-        <NodeInfo nodeInfo={node} />
-      </HoverCardContent>
-    </HoverCard>
-  );
-}
-
-export function SourceNumberButton({
-  index,
-  className,
-}: {
-  index: number;
-  className?: string;
-}) {
-  return (
-    <span
-      className={cn(
-        "text-xs w-5 h-5 rounded-full bg-gray-100 inline-flex items-center justify-center",
-        className,
-      )}
-    >
-      {index + 1}
-    </span>
-  );
-}
-
-export function DocumentInfo({
-  document,
-  className,
-}: {
-  document: Document;
-  className?: string;
-}) {
-  const { url, sources } = document;
-  // Extract filename from URL
-  const urlParts = url.split("/");
-  const fileName = urlParts.length > 0 ? urlParts[urlParts.length - 1] : url;
-  const fileExt = fileName?.split(".").pop() as DocumentFileType | undefined;
-
-  const previewFile = {
-    name: fileName,
-    type: fileExt as DocumentFileType,
-  };
-
-  const DocumentDetail = (
-    <div className={`relative ${className}`}>
-      <PreviewCard className={"cursor-pointer"} file={previewFile} />
-      <div className="absolute bottom-2 right-2 space-x-2 flex">
-        {sources.map((node: SourceNode, index: number) => (
-          <div key={node.id}>
-            <SourceInfo node={node} index={index} />
-          </div>
-        ))}
-      </div>
-    </div>
-  );
-
-  if (url.endsWith(".pdf")) {
-    // open internal pdf dialog for pdf files when click document card
-    return <PdfDialog documentId={url} url={url} trigger={DocumentDetail} />;
-  }
-  // open external link when click document card for other file types
-  return <div onClick={() => window.open(url, "_blank")}>{DocumentDetail}</div>;
-}
-
-function NodeInfo({ nodeInfo }: { nodeInfo: SourceNode }) {
-  const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 1000 });
-
-  const pageNumber =
-    // XXX: page_label is used in Python, but page_number is used by Typescript
-    (nodeInfo.metadata?.page_number as number) ??
-    (nodeInfo.metadata?.page_label as number) ??
-    null;
-
-  return (
-    <div className="space-y-4">
-      <div className="flex justify-between items-center">
-        <span className="font-semibold">
-          {pageNumber ? `On page ${pageNumber}:` : "Node content:"}
-        </span>
-        {nodeInfo.text && (
-          <Button
-            onClick={(e) => {
-              e.stopPropagation();
-              copyToClipboard(nodeInfo.text);
-            }}
-            size="icon"
-            variant="ghost"
-            className="h-12 w-12 shrink-0"
-          >
-            {isCopied ? (
-              <Check className="h-4 w-4" />
-            ) : (
-              <Copy className="h-4 w-4" />
-            )}
-          </Button>
-        )}
-      </div>
-
-      {nodeInfo.text && (
-        <pre className="max-h-[200px] overflow-auto whitespace-pre-line">
-          &ldquo;{nodeInfo.text}&rdquo;
-        </pre>
-      )}
-    </div>
-  );
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-suggestedQuestions.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-suggestedQuestions.tsx
deleted file mode 100644
index 56353856a..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-suggestedQuestions.tsx
+++ /dev/null
@@ -1,31 +0,0 @@
-import { ChatHandler, SuggestedQuestionsData } from "..";
-
-export function SuggestedQuestions({
-  questions,
-  append,
-  isLastMessage,
-}: {
-  questions: SuggestedQuestionsData;
-  append: Pick<ChatHandler, "append">["append"];
-  isLastMessage: boolean;
-}) {
-  const showQuestions = isLastMessage && questions.length > 0;
-  return (
-    showQuestions &&
-    append !== undefined && (
-      <div className="flex flex-col space-y-2">
-        {questions.map((question, index) => (
-          <a
-            key={index}
-            onClick={() => {
-              append({ role: "user", content: question });
-            }}
-            className="text-sm italic hover:underline cursor-pointer"
-          >
-            {"->"} {question}
-          </a>
-        ))}
-      </div>
-    )
-  );
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-tools.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-tools.tsx
deleted file mode 100644
index 9f867470d..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/chat-tools.tsx
+++ /dev/null
@@ -1,40 +0,0 @@
-import { ToolData } from "../index";
-import { Artifact, CodeArtifact } from "../widgets/Artifact";
-import { WeatherCard, WeatherData } from "../widgets/WeatherCard";
-
-// TODO: If needed, add displaying more tool outputs here
-export default function ChatTools({
-  data,
-  artifactVersion,
-}: {
-  data: ToolData;
-  artifactVersion?: number;
-}) {
-  if (!data) return null;
-  const { toolCall, toolOutput } = data;
-
-  if (toolOutput.isError) {
-    return (
-      <div className="border-l-2 border-red-400 pl-2">
-        There was an error when calling the tool {toolCall.name} with input:{" "}
-        <br />
-        {JSON.stringify(toolCall.input)}
-      </div>
-    );
-  }
-
-  switch (toolCall.name) {
-    case "get_weather_information":
-      const weatherData = toolOutput.output as unknown as WeatherData;
-      return <WeatherCard data={weatherData} />;
-    case "artifact":
-      return (
-        <Artifact
-          artifact={toolOutput.output as CodeArtifact}
-          version={artifactVersion}
-        />
-      );
-    default:
-      return null;
-  }
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/codeblock.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/codeblock.tsx
deleted file mode 100644
index e61762ae2..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/codeblock.tsx
+++ /dev/null
@@ -1,131 +0,0 @@
-"use client";
-
-import hljs from "highlight.js";
-// instead of atom-one-dark theme, there are a lot of others: https://highlightjs.org/demo
-import "highlight.js/styles/atom-one-dark-reasonable.css";
-import { Check, Copy, Download } from "lucide-react";
-import { FC, memo, useEffect, useRef } from "react";
-import { Button } from "../../button";
-import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
-
-interface Props {
-  language: string;
-  value: string;
-  className?: string;
-}
-
-interface languageMap {
-  [key: string]: string | undefined;
-}
-
-export const programmingLanguages: languageMap = {
-  javascript: ".js",
-  python: ".py",
-  java: ".java",
-  c: ".c",
-  cpp: ".cpp",
-  "c++": ".cpp",
-  "c#": ".cs",
-  ruby: ".rb",
-  php: ".php",
-  swift: ".swift",
-  "objective-c": ".m",
-  kotlin: ".kt",
-  typescript: ".ts",
-  go: ".go",
-  perl: ".pl",
-  rust: ".rs",
-  scala: ".scala",
-  haskell: ".hs",
-  lua: ".lua",
-  shell: ".sh",
-  sql: ".sql",
-  html: ".html",
-  css: ".css",
-  // add more file extensions here, make sure the key is same as language prop in CodeBlock.tsx component
-};
-
-export const generateRandomString = (length: number, lowercase = false) => {
-  const chars = "ABCDEFGHJKLMNPQRSTUVWXY3456789"; // excluding similar looking characters like Z, 2, I, 1, O, 0
-  let result = "";
-  for (let i = 0; i < length; i++) {
-    result += chars.charAt(Math.floor(Math.random() * chars.length));
-  }
-  return lowercase ? result.toLowerCase() : result;
-};
-
-const CodeBlock: FC<Props> = memo(({ language, value, className }) => {
-  const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 2000 });
-  const codeRef = useRef<HTMLElement>(null);
-
-  useEffect(() => {
-    if (codeRef.current && codeRef.current.dataset.highlighted !== "yes") {
-      hljs.highlightElement(codeRef.current);
-    }
-  }, [language, value]);
-
-  const downloadAsFile = () => {
-    if (typeof window === "undefined") {
-      return;
-    }
-    const fileExtension = programmingLanguages[language] || ".file";
-    const suggestedFileName = `file-${generateRandomString(
-      3,
-      true,
-    )}${fileExtension}`;
-    const fileName = window.prompt("Enter file name", suggestedFileName);
-
-    if (!fileName) {
-      // User pressed cancel on prompt.
-      return;
-    }
-
-    const blob = new Blob([value], { type: "text/plain" });
-    const url = URL.createObjectURL(blob);
-    const link = document.createElement("a");
-    link.download = fileName;
-    link.href = url;
-    link.style.display = "none";
-    document.body.appendChild(link);
-    link.click();
-    document.body.removeChild(link);
-    URL.revokeObjectURL(url);
-  };
-
-  const onCopy = () => {
-    if (isCopied) return;
-    copyToClipboard(value);
-  };
-
-  return (
-    <div
-      className={`codeblock relative w-full bg-zinc-950 font-sans ${className}`}
-    >
-      <div className="flex w-full items-center justify-between bg-zinc-800 px-6 py-2 pr-4 text-zinc-100">
-        <span className="text-xs lowercase">{language}</span>
-        <div className="flex items-center space-x-1">
-          <Button variant="ghost" onClick={downloadAsFile} size="icon">
-            <Download />
-            <span className="sr-only">Download</span>
-          </Button>
-          <Button variant="ghost" size="icon" onClick={onCopy}>
-            {isCopied ? (
-              <Check className="h-4 w-4" />
-            ) : (
-              <Copy className="h-4 w-4" />
-            )}
-            <span className="sr-only">Copy code</span>
-          </Button>
-        </div>
-      </div>
-      <pre className="border border-zinc-700">
-        <code ref={codeRef} className={`language-${language} font-mono`}>
-          {value}
-        </code>
-      </pre>
-    </div>
-  );
-});
-CodeBlock.displayName = "CodeBlock";
-
-export { CodeBlock };
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/index.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/index.tsx
deleted file mode 100644
index 47ec2ba81..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/index.tsx
+++ /dev/null
@@ -1,184 +0,0 @@
-import { Check, Copy } from "lucide-react";
-
-import { Message } from "ai";
-import { Fragment } from "react";
-import { Button } from "../../button";
-import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
-import {
-  AgentEventData,
-  ChatHandler,
-  DocumentFileData,
-  EventData,
-  ImageData,
-  MessageAnnotation,
-  MessageAnnotationType,
-  SuggestedQuestionsData,
-  ToolData,
-  getAnnotationData,
-  getSourceAnnotationData,
-} from "../index";
-import { ChatAgentEvents } from "./chat-agent-events";
-import ChatAvatar from "./chat-avatar";
-import { ChatEvents } from "./chat-events";
-import { ChatFiles } from "./chat-files";
-import { ChatImage } from "./chat-image";
-import { ChatSources } from "./chat-sources";
-import { SuggestedQuestions } from "./chat-suggestedQuestions";
-import ChatTools from "./chat-tools";
-import Markdown from "./markdown";
-
-type ContentDisplayConfig = {
-  order: number;
-  component: JSX.Element | null;
-};
-
-function ChatMessageContent({
-  message,
-  isLoading,
-  append,
-  isLastMessage,
-  artifactVersion,
-}: {
-  message: Message;
-  isLoading: boolean;
-  append: Pick<ChatHandler, "append">["append"];
-  isLastMessage: boolean;
-  artifactVersion: number | undefined;
-}) {
-  const annotations = message.annotations as MessageAnnotation[] | undefined;
-  if (!annotations?.length) return <Markdown content={message.content} />;
-
-  const imageData = getAnnotationData<ImageData>(
-    annotations,
-    MessageAnnotationType.IMAGE,
-  );
-  const contentFileData = getAnnotationData<DocumentFileData>(
-    annotations,
-    MessageAnnotationType.DOCUMENT_FILE,
-  );
-  const eventData = getAnnotationData<EventData>(
-    annotations,
-    MessageAnnotationType.EVENTS,
-  );
-  const agentEventData = getAnnotationData<AgentEventData>(
-    annotations,
-    MessageAnnotationType.AGENT_EVENTS,
-  );
-
-  const sourceData = getSourceAnnotationData(annotations);
-
-  const toolData = getAnnotationData<ToolData>(
-    annotations,
-    MessageAnnotationType.TOOLS,
-  );
-  const suggestedQuestionsData = getAnnotationData<SuggestedQuestionsData>(
-    annotations,
-    MessageAnnotationType.SUGGESTED_QUESTIONS,
-  );
-
-  const contents: ContentDisplayConfig[] = [
-    {
-      order: 1,
-      component: imageData[0] ? <ChatImage data={imageData[0]} /> : null,
-    },
-    {
-      order: -3,
-      component:
-        eventData.length > 0 ? (
-          <ChatEvents isLoading={isLoading} data={eventData} />
-        ) : null,
-    },
-    {
-      order: -2,
-      component:
-        agentEventData.length > 0 ? (
-          <ChatAgentEvents
-            data={agentEventData}
-            isFinished={!!message.content}
-          />
-        ) : null,
-    },
-    {
-      order: 2,
-      component: contentFileData[0] ? (
-        <ChatFiles data={contentFileData[0]} />
-      ) : null,
-    },
-    {
-      order: -1,
-      component: toolData[0] ? (
-        <ChatTools data={toolData[0]} artifactVersion={artifactVersion} />
-      ) : null,
-    },
-    {
-      order: 0,
-      component: <Markdown content={message.content} sources={sourceData[0]} />,
-    },
-    {
-      order: 3,
-      component: sourceData[0] ? <ChatSources data={sourceData[0]} /> : null,
-    },
-    {
-      order: 4,
-      component: suggestedQuestionsData[0] ? (
-        <SuggestedQuestions
-          questions={suggestedQuestionsData[0]}
-          append={append}
-          isLastMessage={isLastMessage}
-        />
-      ) : null,
-    },
-  ];
-
-  return (
-    <div className="flex-1 gap-4 flex flex-col">
-      {contents
-        .sort((a, b) => a.order - b.order)
-        .map((content, index) => (
-          <Fragment key={index}>{content.component}</Fragment>
-        ))}
-    </div>
-  );
-}
-
-export default function ChatMessage({
-  chatMessage,
-  isLoading,
-  append,
-  isLastMessage,
-  artifactVersion,
-}: {
-  chatMessage: Message;
-  isLoading: boolean;
-  append: Pick<ChatHandler, "append">["append"];
-  isLastMessage: boolean;
-  artifactVersion: number | undefined;
-}) {
-  const { isCopied, copyToClipboard } = useCopyToClipboard({ timeout: 2000 });
-  return (
-    <div className="flex items-start gap-4 pr-5 pt-5">
-      <ChatAvatar role={chatMessage.role} />
-      <div className="group flex flex-1 justify-between gap-2">
-        <ChatMessageContent
-          message={chatMessage}
-          isLoading={isLoading}
-          append={append}
-          isLastMessage={isLastMessage}
-          artifactVersion={artifactVersion}
-        />
-        <Button
-          onClick={() => copyToClipboard(chatMessage.content)}
-          size="icon"
-          variant="ghost"
-          className="h-8 w-8 opacity-0 group-hover:opacity-100"
-        >
-          {isCopied ? (
-            <Check className="h-4 w-4" />
-          ) : (
-            <Copy className="h-4 w-4" />
-          )}
-        </Button>
-      </div>
-    </div>
-  );
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/markdown.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/markdown.tsx
deleted file mode 100644
index 9074e208d..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-message/markdown.tsx
+++ /dev/null
@@ -1,172 +0,0 @@
-import "katex/dist/katex.min.css";
-import { FC, memo } from "react";
-import ReactMarkdown, { Options } from "react-markdown";
-import rehypeKatex from "rehype-katex";
-import remarkGfm from "remark-gfm";
-import remarkMath from "remark-math";
-
-import { DOCUMENT_FILE_TYPES, DocumentFileType, SourceData } from "..";
-import { useClientConfig } from "../hooks/use-config";
-import { DocumentInfo, SourceNumberButton } from "./chat-sources";
-import { CodeBlock } from "./codeblock";
-
-const MemoizedReactMarkdown: FC<Options> = memo(
-  ReactMarkdown,
-  (prevProps, nextProps) =>
-    prevProps.children === nextProps.children &&
-    prevProps.className === nextProps.className,
-);
-
-const preprocessLaTeX = (content: string) => {
-  // Replace block-level LaTeX delimiters \[ \] with $$ $$
-  const blockProcessedContent = content.replace(
-    /\\\[([\s\S]*?)\\\]/g,
-    (_, equation) => `$$${equation}$$`,
-  );
-  // Replace inline LaTeX delimiters \( \) with $ $
-  const inlineProcessedContent = blockProcessedContent.replace(
-    /\\\[([\s\S]*?)\\\]/g,
-    (_, equation) => `$${equation}$`,
-  );
-  return inlineProcessedContent;
-};
-
-const preprocessMedia = (content: string) => {
-  // Remove `sandbox:` from the beginning of the URL
-  // to fix OpenAI's models issue appending `sandbox:` to the relative URL
-  return content.replace(/(sandbox|attachment|snt):/g, "");
-};
-
-/**
- * Update the citation flag [citation:id]() to the new format [citation:index](url)
- */
-const preprocessCitations = (content: string, sources?: SourceData) => {
-  if (sources) {
-    const citationRegex = /\[citation:(.+?)\]\(\)/g;
-    let match;
-    // Find all the citation references in the content
-    while ((match = citationRegex.exec(content)) !== null) {
-      const citationId = match[1];
-      // Find the source node with the id equal to the citation-id, also get the index of the source node
-      const sourceNode = sources.nodes.find((node) => node.id === citationId);
-      // If the source node is found, replace the citation reference with the new format
-      if (sourceNode !== undefined) {
-        content = content.replace(
-          match[0],
-          `[citation:${sources.nodes.indexOf(sourceNode)}]()`,
-        );
-      } else {
-        // If the source node is not found, remove the citation reference
-        content = content.replace(match[0], "");
-      }
-    }
-  }
-  return content;
-};
-
-const preprocessContent = (content: string, sources?: SourceData) => {
-  return preprocessCitations(
-    preprocessMedia(preprocessLaTeX(content)),
-    sources,
-  );
-};
-
-export default function Markdown({
-  content,
-  sources,
-}: {
-  content: string;
-  sources?: SourceData;
-}) {
-  const processedContent = preprocessContent(content, sources);
-  const { backend } = useClientConfig();
-
-  return (
-    <MemoizedReactMarkdown
-      className="prose dark:prose-invert prose-p:leading-relaxed prose-pre:p-0 break-words custom-markdown"
-      remarkPlugins={[remarkGfm, remarkMath]}
-      rehypePlugins={[rehypeKatex as any]}
-      components={{
-        p({ children }) {
-          return <div className="mb-2 last:mb-0">{children}</div>;
-        },
-        code({ node, inline, className, children, ...props }) {
-          if (children.length) {
-            if (children[0] == "▍") {
-              return (
-                <span className="mt-1 animate-pulse cursor-default">▍</span>
-              );
-            }
-
-            children[0] = (children[0] as string).replace("`▍`", "▍");
-          }
-
-          const match = /language-(\w+)/.exec(className || "");
-
-          if (inline) {
-            return (
-              <code className={className} {...props}>
-                {children}
-              </code>
-            );
-          }
-
-          return (
-            <CodeBlock
-              key={Math.random()}
-              language={(match && match[1]) || ""}
-              value={String(children).replace(/\n$/, "")}
-              className="mb-2"
-              {...props}
-            />
-          );
-        },
-        a({ href, children }) {
-          // If href starts with `{backend}/api/files`, then it's a local document and we use DocumenInfo for rendering
-          if (href?.startsWith(backend + "/api/files")) {
-            // Check if the file is document file type
-            const fileExtension = href.split(".").pop()?.toLowerCase();
-
-            if (
-              fileExtension &&
-              DOCUMENT_FILE_TYPES.includes(fileExtension as DocumentFileType)
-            ) {
-              return (
-                <DocumentInfo
-                  document={{
-                    url: backend
-                      ? new URL(decodeURIComponent(href)).href
-                      : href,
-                    sources: [],
-                  }}
-                  className="mb-2 mt-2"
-                />
-              );
-            }
-          }
-          // If a text link starts with 'citation:', then render it as a citation reference
-          if (
-            Array.isArray(children) &&
-            typeof children[0] === "string" &&
-            children[0].startsWith("citation:")
-          ) {
-            const index = Number(children[0].replace("citation:", ""));
-            if (!isNaN(index)) {
-              return <SourceNumberButton index={index} />;
-            } else {
-              // citation is not looked up yet, don't render anything
-              return <></>;
-            }
-          }
-          return (
-            <a href={href} target="_blank">
-              {children}
-            </a>
-          );
-        },
-      }}
-    >
-      {processedContent}
-    </MemoizedReactMarkdown>
-  );
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-messages.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-messages.tsx
index 2f29def29..17c4e021a 100644
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat-messages.tsx
+++ b/templates/types/streaming/nextjs/app/components/ui/chat/chat-messages.tsx
@@ -1,136 +1,30 @@
-import { Loader2 } from "lucide-react";
-import { useEffect, useMemo, useRef, useState } from "react";
+"use client";
 
-import { ToolData } from ".";
-import { Button } from "../button";
-import ChatActions from "./chat-actions";
-import ChatMessage from "./chat-message";
-import { ChatHandler } from "./chat.interface";
-import { useClientConfig } from "./hooks/use-config";
-
-export default function ChatMessages(
-  props: Pick<
-    ChatHandler,
-    "messages" | "isLoading" | "reload" | "stop" | "append"
-  >,
-) {
-  const { backend } = useClientConfig();
-  const [starterQuestions, setStarterQuestions] = useState<string[]>();
-
-  const scrollableChatContainerRef = useRef<HTMLDivElement>(null);
-  const messageLength = props.messages.length;
-  const lastMessage = props.messages[messageLength - 1];
-
-  const scrollToBottom = () => {
-    if (scrollableChatContainerRef.current) {
-      scrollableChatContainerRef.current.scrollTop =
-        scrollableChatContainerRef.current.scrollHeight;
-    }
-  };
-
-  const isLastMessageFromAssistant =
-    messageLength > 0 && lastMessage?.role !== "user";
-  const showReload =
-    props.reload && !props.isLoading && isLastMessageFromAssistant;
-  const showStop = props.stop && props.isLoading;
-
-  // `isPending` indicate
-  // that stream response is not yet received from the server,
-  // so we show a loading indicator to give a better UX.
-  const isPending = props.isLoading && !isLastMessageFromAssistant;
-
-  useEffect(() => {
-    scrollToBottom();
-  }, [messageLength, lastMessage]);
-
-  useEffect(() => {
-    if (!starterQuestions) {
-      fetch(`${backend}/api/chat/config`)
-        .then((response) => response.json())
-        .then((data) => {
-          if (data?.starterQuestions) {
-            setStarterQuestions(data.starterQuestions);
-          }
-        })
-        .catch((error) => console.error("Error fetching config", error));
-    }
-  }, [starterQuestions, backend]);
-
-  // build a map of message id to artifact version
-  const artifactVersionMap = useMemo(() => {
-    const map = new Map<string, number | undefined>();
-    let versionIndex = 1;
-    props.messages.forEach((m) => {
-      m.annotations?.forEach((annotation) => {
-        if (
-          typeof annotation === "object" &&
-          annotation != null &&
-          "type" in annotation &&
-          annotation.type === "tools"
-        ) {
-          const data = annotation.data as ToolData;
-          if (data?.toolCall?.name === "artifact") {
-            map.set(m.id, versionIndex);
-            versionIndex++;
-          }
-        }
-      });
-    });
-    return map;
-  }, [props.messages]);
+import { ChatMessage, ChatMessages, useChatUI } from "@llamaindex/chat-ui";
+import { ChatMessageAvatar } from "./chat-avatar";
+import { ChatMessageContent } from "./chat-message-content";
+import { ChatStarter } from "./chat-starter";
 
+export default function CustomChatMessages() {
+  const { messages } = useChatUI();
   return (
-    <div
-      className="flex-1 w-full rounded-xl bg-white p-4 shadow-xl relative overflow-y-auto"
-      ref={scrollableChatContainerRef}
-    >
-      <div className="flex flex-col gap-5 divide-y">
-        {props.messages.map((m, i) => {
-          const isLoadingMessage = i === messageLength - 1 && props.isLoading;
-          return (
-            <ChatMessage
-              key={m.id}
-              chatMessage={m}
-              isLoading={isLoadingMessage}
-              append={props.append!}
-              isLastMessage={i === messageLength - 1}
-              artifactVersion={artifactVersionMap.get(m.id)}
-            />
-          );
-        })}
-        {isPending && (
-          <div className="flex justify-center items-center pt-10">
-            <Loader2 className="h-4 w-4 animate-spin" />
-          </div>
-        )}
-      </div>
-      {(showReload || showStop) && (
-        <div className="flex justify-end py-4">
-          <ChatActions
-            reload={props.reload}
-            stop={props.stop}
-            showReload={showReload}
-            showStop={showStop}
-          />
-        </div>
-      )}
-      {!messageLength && starterQuestions?.length && props.append && (
-        <div className="absolute bottom-6 left-0 w-full">
-          <div className="grid grid-cols-2 gap-2 mx-20">
-            {starterQuestions.map((question, i) => (
-              <Button
-                variant="outline"
-                key={i}
-                onClick={() =>
-                  props.append!({ role: "user", content: question })
-                }
-              >
-                {question}
-              </Button>
-            ))}
-          </div>
-        </div>
-      )}
-    </div>
+    <ChatMessages className="shadow-xl rounded-xl">
+      <ChatMessages.List>
+        {messages.map((message, index) => (
+          <ChatMessage
+            key={message.id}
+            message={message}
+            isLast={index === messages.length - 1}
+          >
+            <ChatMessageAvatar />
+            <ChatMessageContent />
+            <ChatMessage.Actions />
+          </ChatMessage>
+        ))}
+        <ChatMessages.Loading />
+      </ChatMessages.List>
+      <ChatMessages.Actions />
+      <ChatStarter />
+    </ChatMessages>
   );
 }
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat-starter.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/chat-starter.tsx
new file mode 100644
index 000000000..0f455316d
--- /dev/null
+++ b/templates/types/streaming/nextjs/app/components/ui/chat/chat-starter.tsx
@@ -0,0 +1,26 @@
+import { useChatUI } from "@llamaindex/chat-ui";
+import { StarterQuestions } from "@llamaindex/chat-ui/widgets";
+import { useEffect, useState } from "react";
+import { useClientConfig } from "./hooks/use-config";
+
+export function ChatStarter() {
+  const { append } = useChatUI();
+  const { backend } = useClientConfig();
+  const [starterQuestions, setStarterQuestions] = useState<string[]>();
+
+  useEffect(() => {
+    if (!starterQuestions) {
+      fetch(`${backend}/api/chat/config`)
+        .then((response) => response.json())
+        .then((data) => {
+          if (data?.starterQuestions) {
+            setStarterQuestions(data.starterQuestions);
+          }
+        })
+        .catch((error) => console.error("Error fetching config", error));
+    }
+  }, [starterQuestions, backend]);
+
+  if (!starterQuestions?.length) return null;
+  return <StarterQuestions append={append} questions={starterQuestions} />;
+}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/chat.interface.ts b/templates/types/streaming/nextjs/app/components/ui/chat/chat.interface.ts
deleted file mode 100644
index 5483abde0..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/chat.interface.ts
+++ /dev/null
@@ -1,25 +0,0 @@
-import { Message } from "ai";
-
-export interface ChatHandler {
-  messages: Message[];
-  input: string;
-  isLoading: boolean;
-  handleSubmit: (
-    e: React.FormEvent<HTMLFormElement>,
-    ops?: {
-      data?: any;
-    },
-  ) => void;
-  handleInputChange: (e: React.ChangeEvent<HTMLTextAreaElement>) => void;
-  reload?: () => void;
-  stop?: () => void;
-  onFileUpload?: (file: File) => Promise<void>;
-  onFileError?: (errMsg: string) => void;
-  setInput?: (input: string) => void;
-  append?: (
-    message: Message | Omit<Message, "id">,
-    ops?: {
-      data: any;
-    },
-  ) => Promise<string | null | undefined>;
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/widgets/LlamaCloudSelector.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/custom/llama-cloud-selector.tsx
similarity index 96%
rename from templates/types/streaming/nextjs/app/components/ui/chat/widgets/LlamaCloudSelector.tsx
rename to templates/types/streaming/nextjs/app/components/ui/chat/custom/llama-cloud-selector.tsx
index 6d67081c9..f40a33a72 100644
--- a/templates/types/streaming/nextjs/app/components/ui/chat/widgets/LlamaCloudSelector.tsx
+++ b/templates/types/streaming/nextjs/app/components/ui/chat/custom/llama-cloud-selector.tsx
@@ -1,3 +1,4 @@
+import { useChatUI } from "@llamaindex/chat-ui";
 import { Loader2 } from "lucide-react";
 import { useCallback, useEffect, useState } from "react";
 import {
@@ -35,19 +36,18 @@ type LlamaCloudConfig = {
 };
 
 export interface LlamaCloudSelectorProps {
-  setRequestData?: React.Dispatch<any>;
   onSelect?: (pipeline: PipelineConfig | undefined) => void;
   defaultPipeline?: PipelineConfig;
   shouldCheckValid?: boolean;
 }
 
 export function LlamaCloudSelector({
-  setRequestData,
   onSelect,
   defaultPipeline,
   shouldCheckValid = false,
 }: LlamaCloudSelectorProps) {
   const { backend } = useClientConfig();
+  const { setRequestData } = useChatUI();
   const [config, setConfig] = useState<LlamaCloudConfig>();
 
   const updateRequestParams = useCallback(
@@ -97,6 +97,10 @@ export function LlamaCloudSelector({
     setPipeline(JSON.parse(value) as PipelineConfig);
   };
 
+  if (process.env.NEXT_PUBLIC_USE_LLAMACLOUD !== "true") {
+    return null;
+  }
+
   if (!config) {
     return (
       <div className="flex justify-center items-center p-3">
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/custom/markdown.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/custom/markdown.tsx
new file mode 100644
index 000000000..88925e8bc
--- /dev/null
+++ b/templates/types/streaming/nextjs/app/components/ui/chat/custom/markdown.tsx
@@ -0,0 +1,27 @@
+import { SourceData } from "@llamaindex/chat-ui";
+import { Markdown as MarkdownUI } from "@llamaindex/chat-ui/widgets";
+import { useClientConfig } from "../hooks/use-config";
+
+const preprocessMedia = (content: string) => {
+  // Remove `sandbox:` from the beginning of the URL before rendering markdown
+  // OpenAI models sometimes prepend `sandbox:` to relative URLs - this fixes it
+  return content.replace(/(sandbox|attachment|snt):/g, "");
+};
+
+export function Markdown({
+  content,
+  sources,
+}: {
+  content: string;
+  sources?: SourceData;
+}) {
+  const { backend } = useClientConfig();
+  const processedContent = preprocessMedia(content);
+  return (
+    <MarkdownUI
+      content={processedContent}
+      backend={backend}
+      sources={sources}
+    />
+  );
+}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/hooks/use-file.ts b/templates/types/streaming/nextjs/app/components/ui/chat/hooks/use-file.ts
deleted file mode 100644
index 049db6ba7..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/hooks/use-file.ts
+++ /dev/null
@@ -1,121 +0,0 @@
-"use client";
-
-import { JSONValue } from "llamaindex";
-import { useState } from "react";
-import {
-  DocumentFile,
-  DocumentFileType,
-  MessageAnnotation,
-  MessageAnnotationType,
-} from "..";
-import { useClientConfig } from "./use-config";
-
-const docMineTypeMap: Record<string, DocumentFileType> = {
-  "text/csv": "csv",
-  "application/pdf": "pdf",
-  "text/plain": "txt",
-  "application/vnd.openxmlformats-officedocument.wordprocessingml.document":
-    "docx",
-};
-
-export function useFile() {
-  const { backend } = useClientConfig();
-  const [imageUrl, setImageUrl] = useState<string | null>(null);
-  const [files, setFiles] = useState<DocumentFile[]>([]);
-
-  const addDoc = (file: DocumentFile) => {
-    const existedFile = files.find((f) => f.id === file.id);
-    if (!existedFile) {
-      setFiles((prev) => [...prev, file]);
-      return true;
-    }
-    return false;
-  };
-
-  const removeDoc = (file: DocumentFile) => {
-    setFiles((prev) => prev.filter((f) => f.id !== file.id));
-  };
-
-  const reset = () => {
-    imageUrl && setImageUrl(null);
-    files.length && setFiles([]);
-  };
-
-  const uploadContent = async (
-    file: File,
-    requestParams: any = {},
-  ): Promise<DocumentFile> => {
-    const base64 = await readContent({ file, asUrl: true });
-    const uploadAPI = `${backend}/api/chat/upload`;
-    const response = await fetch(uploadAPI, {
-      method: "POST",
-      headers: {
-        "Content-Type": "application/json",
-      },
-      body: JSON.stringify({
-        ...requestParams,
-        base64,
-        name: file.name,
-      }),
-    });
-    if (!response.ok) throw new Error("Failed to upload document.");
-    return (await response.json()) as DocumentFile;
-  };
-
-  const getAnnotations = () => {
-    const annotations: MessageAnnotation[] = [];
-    if (imageUrl) {
-      annotations.push({
-        type: MessageAnnotationType.IMAGE,
-        data: { url: imageUrl },
-      });
-    }
-    if (files.length > 0) {
-      annotations.push({
-        type: MessageAnnotationType.DOCUMENT_FILE,
-        data: { files },
-      });
-    }
-    return annotations as JSONValue[];
-  };
-
-  const readContent = async (input: {
-    file: File;
-    asUrl?: boolean;
-  }): Promise<string> => {
-    const { file, asUrl } = input;
-    const content = await new Promise<string>((resolve, reject) => {
-      const reader = new FileReader();
-      if (asUrl) {
-        reader.readAsDataURL(file);
-      } else {
-        reader.readAsText(file);
-      }
-      reader.onload = () => resolve(reader.result as string);
-      reader.onerror = (error) => reject(error);
-    });
-    return content;
-  };
-
-  const uploadFile = async (file: File, requestParams: any = {}) => {
-    if (file.type.startsWith("image/")) {
-      const base64 = await readContent({ file, asUrl: true });
-      return setImageUrl(base64);
-    }
-
-    const filetype = docMineTypeMap[file.type];
-    if (!filetype) throw new Error("Unsupported document type.");
-    const newDoc = await uploadContent(file, requestParams);
-    return addDoc(newDoc);
-  };
-
-  return {
-    imageUrl,
-    setImageUrl,
-    files,
-    removeDoc,
-    reset,
-    getAnnotations,
-    uploadFile,
-  };
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/index.ts b/templates/types/streaming/nextjs/app/components/ui/chat/index.ts
deleted file mode 100644
index e78a8c277..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/index.ts
+++ /dev/null
@@ -1,139 +0,0 @@
-import { JSONValue } from "ai";
-import ChatInput from "./chat-input";
-import ChatMessages from "./chat-messages";
-
-export { type ChatHandler } from "./chat.interface";
-export { ChatInput, ChatMessages };
-
-export enum MessageAnnotationType {
-  IMAGE = "image",
-  DOCUMENT_FILE = "document_file",
-  SOURCES = "sources",
-  EVENTS = "events",
-  TOOLS = "tools",
-  SUGGESTED_QUESTIONS = "suggested_questions",
-  AGENT_EVENTS = "agent",
-}
-
-export type ImageData = {
-  url: string;
-};
-
-export type DocumentFileType = "csv" | "pdf" | "txt" | "docx";
-export const DOCUMENT_FILE_TYPES: DocumentFileType[] = [
-  "csv",
-  "pdf",
-  "txt",
-  "docx",
-];
-
-export type DocumentFile = {
-  id: string;
-  name: string; // The uploaded file name in the backend
-  size: number; // The file size in bytes
-  type: DocumentFileType;
-  url: string; // The URL of the uploaded file in the backend
-  refs?: string[]; // DocumentIDs of the uploaded file in the vector index
-};
-
-export type DocumentFileData = {
-  files: DocumentFile[];
-};
-
-export type SourceNode = {
-  id: string;
-  metadata: Record<string, unknown>;
-  score?: number;
-  text: string;
-  url: string;
-};
-
-export type SourceData = {
-  nodes: SourceNode[];
-};
-
-export type EventData = {
-  title: string;
-};
-
-export type ProgressData = {
-  id: string;
-  total: number;
-  current: number;
-};
-
-export type AgentEventData = {
-  agent: string;
-  text: string;
-  type: "text" | "progress";
-  data?: ProgressData;
-};
-
-export type ToolData = {
-  toolCall: {
-    id: string;
-    name: string;
-    input: {
-      [key: string]: JSONValue;
-    };
-  };
-  toolOutput: {
-    output: JSONValue;
-    isError: boolean;
-  };
-};
-
-export type SuggestedQuestionsData = string[];
-
-export type AnnotationData =
-  | ImageData
-  | DocumentFileData
-  | SourceData
-  | EventData
-  | AgentEventData
-  | ToolData
-  | SuggestedQuestionsData;
-
-export type MessageAnnotation = {
-  type: MessageAnnotationType;
-  data: AnnotationData;
-};
-
-const NODE_SCORE_THRESHOLD = 0.25;
-
-export function getAnnotationData<T extends AnnotationData>(
-  annotations: MessageAnnotation[],
-  type: MessageAnnotationType,
-): T[] {
-  return annotations.filter((a) => a.type === type).map((a) => a.data as T);
-}
-
-export function getSourceAnnotationData(
-  annotations: MessageAnnotation[],
-): SourceData[] {
-  const data = getAnnotationData<SourceData>(
-    annotations,
-    MessageAnnotationType.SOURCES,
-  );
-  if (data.length > 0) {
-    const sourceData = data[0] as SourceData;
-    if (sourceData.nodes) {
-      sourceData.nodes = preprocessSourceNodes(sourceData.nodes);
-    }
-  }
-  return data;
-}
-
-function preprocessSourceNodes(nodes: SourceNode[]): SourceNode[] {
-  // Filter source nodes has lower score
-  nodes = nodes
-    .filter((node) => (node.score ?? 1) > NODE_SCORE_THRESHOLD)
-    .filter((node) => node.url && node.url.trim() !== "")
-    .sort((a, b) => (b.score ?? 1) - (a.score ?? 1))
-    .map((node) => {
-      // remove trailing slash for node url if exists
-      node.url = node.url.replace(/\/$/, "");
-      return node;
-    });
-  return nodes;
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/widgets/Artifact.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/tools/artifact.tsx
similarity index 98%
rename from templates/types/streaming/nextjs/app/components/ui/chat/widgets/Artifact.tsx
rename to templates/types/streaming/nextjs/app/components/ui/chat/tools/artifact.tsx
index fa2a60592..fe6e81998 100644
--- a/templates/types/streaming/nextjs/app/components/ui/chat/widgets/Artifact.tsx
+++ b/templates/types/streaming/nextjs/app/components/ui/chat/tools/artifact.tsx
@@ -10,7 +10,7 @@ import {
 } from "../../collapsible";
 import { cn } from "../../lib/utils";
 import { Tabs, TabsContent, TabsList, TabsTrigger } from "../../tabs";
-import Markdown from "../chat-message/markdown";
+import { Markdown } from "../custom/markdown";
 import { useClientConfig } from "../hooks/use-config";
 import { useCopyToClipboard } from "../hooks/use-copy-to-clipboard";
 
@@ -29,12 +29,17 @@ export type CodeArtifact = {
   files?: string[];
 };
 
+type OutputUrl = {
+  url: string;
+  filename: string;
+};
+
 type ArtifactResult = {
   template: string;
   stdout: string[];
   stderr: string[];
   runtimeError?: { name: string; value: string; tracebackRaw: string[] };
-  outputUrls: Array<{ url: string; filename: string }>;
+  outputUrls: OutputUrl[];
   url: string;
 };
 
@@ -272,11 +277,7 @@ function CodeSandboxPreview({ url }: { url: string }) {
   );
 }
 
-function InterpreterOutput({
-  outputUrls,
-}: {
-  outputUrls: Array<{ url: string; filename: string }>;
-}) {
+function InterpreterOutput({ outputUrls }: { outputUrls: OutputUrl[] }) {
   return (
     <ul className="flex flex-col gap-2 mt-4">
       {outputUrls.map((url) => (
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/tools/chat-tools.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/tools/chat-tools.tsx
new file mode 100644
index 000000000..c57a6d5b5
--- /dev/null
+++ b/templates/types/streaming/nextjs/app/components/ui/chat/tools/chat-tools.tsx
@@ -0,0 +1,89 @@
+import {
+  getAnnotationData,
+  MessageAnnotation,
+  useChatMessage,
+  useChatUI,
+} from "@llamaindex/chat-ui";
+import { JSONValue, Message } from "ai";
+import { useMemo } from "react";
+import { Artifact, CodeArtifact } from "./artifact";
+import { WeatherCard, WeatherData } from "./weather-card";
+
+export function ToolAnnotations({ message }: { message: Message }) {
+  const annotations = message.annotations as MessageAnnotation[] | undefined;
+  const toolData = annotations
+    ? (getAnnotationData(annotations, "tools") as unknown as ToolData[])
+    : null;
+  return toolData?.[0] ? <ChatTools data={toolData[0]} /> : null;
+}
+
+// TODO: Used to render outputs of tools. If needed, add more renderers here.
+function ChatTools({ data }: { data: ToolData }) {
+  const { messages } = useChatUI();
+  const { message } = useChatMessage();
+
+  // build a map of message id to artifact version
+  const artifactVersionMap = useMemo(() => {
+    const map = new Map<string, number | undefined>();
+    let versionIndex = 1;
+    messages.forEach((m) => {
+      m.annotations?.forEach((annotation: any) => {
+        if (
+          typeof annotation === "object" &&
+          annotation != null &&
+          "type" in annotation &&
+          annotation.type === "tools"
+        ) {
+          const data = annotation.data as ToolData;
+          if (data?.toolCall?.name === "artifact") {
+            map.set(m.id, versionIndex);
+            versionIndex++;
+          }
+        }
+      });
+    });
+    return map;
+  }, [messages]);
+
+  if (!data) return null;
+  const { toolCall, toolOutput } = data;
+
+  if (toolOutput.isError) {
+    return (
+      <div className="border-l-2 border-red-400 pl-2">
+        There was an error when calling the tool {toolCall.name} with input:{" "}
+        <br />
+        {JSON.stringify(toolCall.input)}
+      </div>
+    );
+  }
+
+  switch (toolCall.name) {
+    case "get_weather_information":
+      const weatherData = toolOutput.output as unknown as WeatherData;
+      return <WeatherCard data={weatherData} />;
+    case "artifact":
+      return (
+        <Artifact
+          artifact={toolOutput.output as CodeArtifact}
+          version={artifactVersionMap.get(message.id)}
+        />
+      );
+    default:
+      return null;
+  }
+}
+
+type ToolData = {
+  toolCall: {
+    id: string;
+    name: string;
+    input: {
+      [key: string]: JSONValue;
+    };
+  };
+  toolOutput: {
+    output: JSONValue;
+    isError: boolean;
+  };
+};
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/widgets/WeatherCard.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/tools/weather-card.tsx
similarity index 100%
rename from templates/types/streaming/nextjs/app/components/ui/chat/widgets/WeatherCard.tsx
rename to templates/types/streaming/nextjs/app/components/ui/chat/tools/weather-card.tsx
diff --git a/templates/types/streaming/nextjs/app/components/ui/chat/widgets/PdfDialog.tsx b/templates/types/streaming/nextjs/app/components/ui/chat/widgets/PdfDialog.tsx
deleted file mode 100644
index 36abc5cf8..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/chat/widgets/PdfDialog.tsx
+++ /dev/null
@@ -1,67 +0,0 @@
-import dynamic from "next/dynamic";
-import { Button } from "../../button";
-import {
-  Drawer,
-  DrawerClose,
-  DrawerContent,
-  DrawerDescription,
-  DrawerHeader,
-  DrawerTitle,
-  DrawerTrigger,
-} from "../../drawer";
-
-export interface PdfDialogProps {
-  documentId: string;
-  url: string;
-  trigger: React.ReactNode;
-}
-
-// Dynamic imports for client-side rendering only
-const PDFViewer = dynamic(
-  () => import("@llamaindex/pdf-viewer").then((module) => module.PDFViewer),
-  { ssr: false },
-);
-
-const PdfFocusProvider = dynamic(
-  () =>
-    import("@llamaindex/pdf-viewer").then((module) => module.PdfFocusProvider),
-  { ssr: false },
-);
-
-export default function PdfDialog(props: PdfDialogProps) {
-  return (
-    <Drawer direction="left">
-      <DrawerTrigger asChild>{props.trigger}</DrawerTrigger>
-      <DrawerContent className="w-3/5 mt-24 h-full max-h-[96%] ">
-        <DrawerHeader className="flex justify-between">
-          <div className="space-y-2">
-            <DrawerTitle>PDF Content</DrawerTitle>
-            <DrawerDescription>
-              File URL:{" "}
-              <a
-                className="hover:text-blue-900"
-                href={props.url}
-                target="_blank"
-              >
-                {props.url}
-              </a>
-            </DrawerDescription>
-          </div>
-          <DrawerClose asChild>
-            <Button variant="outline">Close</Button>
-          </DrawerClose>
-        </DrawerHeader>
-        <div className="m-4">
-          <PdfFocusProvider>
-            <PDFViewer
-              file={{
-                id: props.documentId,
-                url: props.url,
-              }}
-            />
-          </PdfFocusProvider>
-        </div>
-      </DrawerContent>
-    </Drawer>
-  );
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/document-preview.tsx b/templates/types/streaming/nextjs/app/components/ui/document-preview.tsx
deleted file mode 100644
index ee3059a05..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/document-preview.tsx
+++ /dev/null
@@ -1,129 +0,0 @@
-import { XCircleIcon } from "lucide-react";
-import Image from "next/image";
-import DocxIcon from "../ui/icons/docx.svg";
-import PdfIcon from "../ui/icons/pdf.svg";
-import SheetIcon from "../ui/icons/sheet.svg";
-import TxtIcon from "../ui/icons/txt.svg";
-import { Button } from "./button";
-import { DocumentFile, DocumentFileType } from "./chat";
-import {
-  Drawer,
-  DrawerClose,
-  DrawerContent,
-  DrawerDescription,
-  DrawerHeader,
-  DrawerTitle,
-  DrawerTrigger,
-} from "./drawer";
-import { cn } from "./lib/utils";
-
-export interface DocumentPreviewProps {
-  file: DocumentFile;
-  onRemove?: () => void;
-}
-
-export function DocumentPreview(props: DocumentPreviewProps) {
-  const { name, size, type, refs } = props.file;
-
-  if (refs?.length) {
-    return (
-      <div title={`Document IDs: ${refs.join(", ")}`}>
-        <PreviewCard {...props} />
-      </div>
-    );
-  }
-
-  return (
-    <Drawer direction="left">
-      <DrawerTrigger asChild>
-        <div>
-          <PreviewCard className="cursor-pointer" {...props} />
-        </div>
-      </DrawerTrigger>
-      <DrawerContent className="w-3/5 mt-24 h-full max-h-[96%] ">
-        <DrawerHeader className="flex justify-between">
-          <div className="space-y-2">
-            <DrawerTitle>{type.toUpperCase()} Raw Content</DrawerTitle>
-            <DrawerDescription>
-              {name} ({inKB(size)} KB)
-            </DrawerDescription>
-          </div>
-          <DrawerClose asChild>
-            <Button variant="outline">Close</Button>
-          </DrawerClose>
-        </DrawerHeader>
-        <div className="m-4 max-h-[80%] overflow-auto">
-          {refs?.length && (
-            <pre className="bg-secondary rounded-md p-4 block text-sm">
-              {refs.join(", ")}
-            </pre>
-          )}
-        </div>
-      </DrawerContent>
-    </Drawer>
-  );
-}
-
-export const FileIcon: Record<DocumentFileType, string> = {
-  csv: SheetIcon,
-  pdf: PdfIcon,
-  docx: DocxIcon,
-  txt: TxtIcon,
-};
-
-export function PreviewCard(props: {
-  file: {
-    name: string;
-    size?: number;
-    type: DocumentFileType;
-  };
-  onRemove?: () => void;
-  className?: string;
-}) {
-  const { onRemove, file, className } = props;
-  return (
-    <div
-      className={cn(
-        "p-2 w-60 max-w-60 bg-secondary rounded-lg text-sm relative",
-        className,
-      )}
-    >
-      <div className="flex flex-row items-center gap-2">
-        <div className="relative h-8 w-8 shrink-0 overflow-hidden rounded-md flex items-center justify-center">
-          <Image
-            className="h-full w-auto object-contain"
-            priority
-            src={FileIcon[file.type]}
-            alt="Icon"
-          />
-        </div>
-        <div className="overflow-hidden">
-          <div className="truncate font-semibold">
-            {file.name} {file.size ? `(${inKB(file.size)} KB)` : ""}
-          </div>
-          {file.type && (
-            <div className="truncate text-token-text-tertiary flex items-center gap-2">
-              <span>{file.type.toUpperCase()} File</span>
-            </div>
-          )}
-        </div>
-      </div>
-      {onRemove && (
-        <div
-          className={cn(
-            "absolute -top-2 -right-2 w-6 h-6 z-10 bg-gray-500 text-white rounded-full",
-          )}
-        >
-          <XCircleIcon
-            className="w-6 h-6 bg-gray-500 text-white rounded-full"
-            onClick={onRemove}
-          />
-        </div>
-      )}
-    </div>
-  );
-}
-
-function inKB(size: number) {
-  return Math.round((size / 1024) * 10) / 10;
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/file-uploader.tsx b/templates/types/streaming/nextjs/app/components/ui/file-uploader.tsx
deleted file mode 100644
index 15f9e4035..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/file-uploader.tsx
+++ /dev/null
@@ -1,136 +0,0 @@
-"use client";
-
-import { Loader2, Paperclip } from "lucide-react";
-import { ChangeEvent, useState } from "react";
-import { buttonVariants } from "./button";
-import { cn } from "./lib/utils";
-
-export interface FileUploaderProps {
-  config?: {
-    inputId?: string;
-    fileSizeLimit?: number;
-    allowedExtensions?: string[];
-    checkExtension?: (extension: string) => string | null;
-    disabled: boolean;
-    multiple?: boolean;
-  };
-  onFileUpload: (file: File) => Promise<void>;
-  onFileError?: (errMsg: string) => void;
-}
-
-const DEFAULT_INPUT_ID = "fileInput";
-const DEFAULT_FILE_SIZE_LIMIT = 1024 * 1024 * 50; // 50 MB
-
-export default function FileUploader({
-  config,
-  onFileUpload,
-  onFileError,
-}: FileUploaderProps) {
-  const [uploading, setUploading] = useState(false);
-  const [remainingFiles, setRemainingFiles] = useState<number>(0);
-
-  const inputId = config?.inputId || DEFAULT_INPUT_ID;
-  const fileSizeLimit = config?.fileSizeLimit || DEFAULT_FILE_SIZE_LIMIT;
-  const allowedExtensions = config?.allowedExtensions;
-  const defaultCheckExtension = (extension: string) => {
-    if (allowedExtensions && !allowedExtensions.includes(extension)) {
-      return `Invalid file type. Please select a file with one of these formats: ${allowedExtensions!.join(
-        ",",
-      )}`;
-    }
-    return null;
-  };
-  const checkExtension = config?.checkExtension ?? defaultCheckExtension;
-
-  const isFileSizeExceeded = (file: File) => {
-    return file.size > fileSizeLimit;
-  };
-
-  const resetInput = () => {
-    const fileInput = document.getElementById(inputId) as HTMLInputElement;
-    fileInput.value = "";
-  };
-
-  const onFileChange = async (e: ChangeEvent<HTMLInputElement>) => {
-    const files = Array.from(e.target.files || []);
-    if (!files.length) return;
-
-    setUploading(true);
-
-    await handleUpload(files);
-
-    resetInput();
-    setUploading(false);
-  };
-
-  const handleUpload = async (files: File[]) => {
-    const onFileUploadError = onFileError || window.alert;
-    // Validate files
-    // If multiple files with image or multiple images
-    if (
-      files.length > 1 &&
-      files.some((file) => file.type.startsWith("image/"))
-    ) {
-      onFileUploadError("Multiple files with image are not supported");
-      return;
-    }
-
-    for (const file of files) {
-      const fileExtension = file.name.split(".").pop() || "";
-      const extensionFileError = checkExtension(fileExtension);
-      if (extensionFileError) {
-        onFileUploadError(extensionFileError);
-        return;
-      }
-
-      if (isFileSizeExceeded(file)) {
-        onFileUploadError(
-          `File size exceeded. Limit is ${fileSizeLimit / 1024 / 1024} MB`,
-        );
-        return;
-      }
-    }
-
-    setRemainingFiles(files.length);
-    for (const file of files) {
-      await onFileUpload(file);
-      setRemainingFiles((prev) => prev - 1);
-    }
-    setRemainingFiles(0);
-  };
-
-  return (
-    <div className="self-stretch">
-      <input
-        type="file"
-        id={inputId}
-        style={{ display: "none" }}
-        onChange={onFileChange}
-        accept={allowedExtensions?.join(",")}
-        disabled={config?.disabled || uploading}
-        multiple={config?.multiple}
-      />
-      <label
-        htmlFor={inputId}
-        className={cn(
-          buttonVariants({ variant: "secondary", size: "icon" }),
-          "cursor-pointer relative",
-          uploading && "opacity-50",
-        )}
-      >
-        {uploading ? (
-          <div className="relative flex items-center justify-center h-full w-full">
-            <Loader2 className="h-6 w-6 animate-spin absolute" />
-            {remainingFiles > 0 && (
-              <span className="text-xs absolute inset-0 flex items-center justify-center">
-                {remainingFiles}
-              </span>
-            )}
-          </div>
-        ) : (
-          <Paperclip className="-rotate-45 w-4 h-4" />
-        )}
-      </label>
-    </div>
-  );
-}
diff --git a/templates/types/streaming/nextjs/app/components/ui/upload-image-preview.tsx b/templates/types/streaming/nextjs/app/components/ui/upload-image-preview.tsx
deleted file mode 100644
index 55ef6e9c2..000000000
--- a/templates/types/streaming/nextjs/app/components/ui/upload-image-preview.tsx
+++ /dev/null
@@ -1,32 +0,0 @@
-import { XCircleIcon } from "lucide-react";
-import Image from "next/image";
-import { cn } from "./lib/utils";
-
-export default function UploadImagePreview({
-  url,
-  onRemove,
-}: {
-  url: string;
-  onRemove: () => void;
-}) {
-  return (
-    <div className="relative w-20 h-20 group">
-      <Image
-        src={url}
-        alt="Uploaded image"
-        fill
-        className="object-cover w-full h-full rounded-xl hover:brightness-75"
-      />
-      <div
-        className={cn(
-          "absolute -top-2 -right-2 w-6 h-6 z-10 bg-gray-500 text-white rounded-full hidden group-hover:block",
-        )}
-      >
-        <XCircleIcon
-          className="w-6 h-6 bg-gray-500 text-white rounded-full"
-          onClick={onRemove}
-        />
-      </div>
-    </div>
-  );
-}
diff --git a/templates/types/streaming/nextjs/app/observability/index.ts b/templates/types/streaming/nextjs/app/observability/index.ts
index 2e4ce2b18..206eff4d8 100644
--- a/templates/types/streaming/nextjs/app/observability/index.ts
+++ b/templates/types/streaming/nextjs/app/observability/index.ts
@@ -1 +1,2 @@
+// TODO: You can add observability here. For templates re-start `create-llama` with `--pro` flag to generate a new project with observability.
 export const initObservability = () => {};
diff --git a/templates/types/streaming/nextjs/package.json b/templates/types/streaming/nextjs/package.json
index d54224829..3a9de9ffd 100644
--- a/templates/types/streaming/nextjs/package.json
+++ b/templates/types/streaming/nextjs/package.json
@@ -12,7 +12,6 @@
   "dependencies": {
     "@apidevtools/swagger-parser": "^10.1.0",
     "@e2b/code-interpreter": "0.0.9-beta.3",
-    "@llamaindex/pdf-viewer": "^1.1.3",
     "@radix-ui/react-collapsible": "^1.0.3",
     "@radix-ui/react-hover-card": "^1.0.7",
     "@radix-ui/react-progress": "^1.1.0",
@@ -32,19 +31,13 @@
     "next": "^14.2.4",
     "react": "^18.2.0",
     "react-dom": "^18.2.0",
-    "react-markdown": "^8.0.7",
-    "rehype-katex": "^7.0.0",
-    "remark": "^14.0.3",
-    "remark-code-import": "^1.2.0",
-    "remark-gfm": "^3.0.1",
-    "remark-math": "^5.1.1",
     "supports-color": "^8.1.1",
     "tailwind-merge": "^2.1.0",
     "tiktoken": "^1.0.15",
     "uuid": "^9.0.1",
     "vaul": "^0.9.1",
     "marked": "^14.1.2",
-    "highlight.js": "^11.10.0"
+    "@llamaindex/chat-ui": "0.0.4"
   },
   "devDependencies": {
     "@types/node": "^20.10.3",
diff --git a/templates/types/streaming/nextjs/tailwind.config.ts b/templates/types/streaming/nextjs/tailwind.config.ts
index aa5580aff..d441c0508 100644
--- a/templates/types/streaming/nextjs/tailwind.config.ts
+++ b/templates/types/streaming/nextjs/tailwind.config.ts
@@ -3,7 +3,11 @@ import { fontFamily } from "tailwindcss/defaultTheme";
 
 const config: Config = {
   darkMode: ["class"],
-  content: ["app/**/*.{ts,tsx}", "components/**/*.{ts,tsx}"],
+  content: [
+    "app/**/*.{ts,tsx}",
+    "components/**/*.{ts,tsx}",
+    "node_modules/@llamaindex/chat-ui/**/*.{ts,tsx}",
+  ],
   theme: {
     container: {
       center: true,