feat: use latest chat UI (#418)

---------

Co-authored-by: Marcus Schiesser <[email protected]>

Author: thucpn

.changeset/wet-pears-matter.md

@@ -0,0 +1,5 @@
+---
+"create-llama": patch
+---
+
+feat: use latest chat UI

templates/types/streaming/nextjs/app/components/ui/chat/chat-input.tsx

@@ -1,7 +1,7 @@
 "use client";
 
 import { ChatInput, useChatUI, useFile } from "@llamaindex/chat-ui";
-import { DocumentPreview, ImagePreview } from "@llamaindex/chat-ui/widgets";
+import { DocumentInfo, ImagePreview } from "@llamaindex/chat-ui/widgets";
 import { LlamaCloudSelector } from "./custom/llama-cloud-selector";
 import { useClientConfig } from "./hooks/use-config";
 
@@ -56,9 +56,10 @@ export default function CustomChatInput() {
         {files.length > 0 && (
           <div className="flex gap-4 w-full overflow-auto py-2">
             {files.map((file) => (
-              <DocumentPreview
+              <DocumentInfo
                 key={file.id}
-                file={file}
+                document={{ url: file.url, sources: [] }}
+                className="mb-2 mt-2"
                 onRemove={() => removeDoc(file)}
               />
             ))}

templates/types/streaming/nextjs/app/components/ui/chat/chat-message-content.tsx

@@ -3,11 +3,13 @@ import {
   ContentPosition,
   getSourceAnnotationData,
   useChatMessage,
+  useChatUI,
 } from "@llamaindex/chat-ui";
 import { Markdown } from "./custom/markdown";
 import { ToolAnnotations } from "./tools/chat-tools";
 
 export function ChatMessageContent() {
+  const { isLoading, append } = useChatUI();
   const { message } = useChatMessage();
   const customContent = [
     {
@@ -26,5 +28,11 @@ export function ChatMessageContent() {
       component: <ToolAnnotations message={message} />,
     },
   ];
-  return <ChatMessage.Content content={customContent} />;
+  return (
+    <ChatMessage.Content
+      content={customContent}
+      isLoading={isLoading}
+      append={append}
+    />
+  );
 }

templates/types/streaming/nextjs/app/components/ui/chat/chat-messages.tsx

@@ -12,7 +12,7 @@ export default function CustomChatMessages() {
       <ChatMessages.List>
         {messages.map((message, index) => (
           <ChatMessage
-            key={message.id}
+            key={index}
             message={message}
             isLast={index === messages.length - 1}
           >

templates/types/streaming/nextjs/app/components/ui/chat/tools/chat-tools.tsx

@@ -1,50 +1,40 @@
 import {
-  getAnnotationData,
+  Message,
   MessageAnnotation,
-  useChatMessage,
+  getAnnotationData,
   useChatUI,
 } from "@llamaindex/chat-ui";
-import { JSONValue, Message } from "ai";
+import { JSONValue } from "ai";
 import { useMemo } from "react";
 import { Artifact, CodeArtifact } from "./artifact";
 import { WeatherCard, WeatherData } from "./weather-card";
 
 export function ToolAnnotations({ message }: { message: Message }) {
+  // TODO: This is a bit of a hack to get the artifact version. better to generate the version in the tool call and
+  // store it in CodeArtifact
+  const { messages } = useChatUI();
+  const artifactVersion = useMemo(
+    () => getArtifactVersion(messages, message),
+    [messages, message],
+  );
+  // Get the tool data from the message annotations
   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;
+  return toolData?.[0] ? (
+    <ChatTools data={toolData[0]} artifactVersion={artifactVersion} />
+  ) : 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]);
-
+function ChatTools({
+  data,
+  artifactVersion,
+}: {
+  data: ToolData;
+  artifactVersion: number | undefined;
+}) {
   if (!data) return null;
   const { toolCall, toolOutput } = data;
 
@@ -66,7 +56,7 @@ function ChatTools({ data }: { data: ToolData }) {
       return (
         <Artifact
           artifact={toolOutput.output as CodeArtifact}
-          version={artifactVersionMap.get(message.id)}
+          version={artifactVersion}
         />
       );
     default:
@@ -87,3 +77,25 @@ type ToolData = {
     isError: boolean;
   };
 };
+
+function getArtifactVersion(
+  messages: Message[],
+  message: Message,
+): number | undefined {
+  const messageId = "id" in message ? message.id : undefined;
+  if (!messageId) return undefined;
+  let versionIndex = 1;
+  for (const m of messages) {
+    const toolData = m.annotations
+      ? (getAnnotationData(m.annotations, "tools") as unknown as ToolData[])
+      : null;
+
+    if (toolData?.some((t) => t.toolCall.name === "artifact")) {
+      if ("id" in m && m.id === messageId) {
+        return versionIndex;
+      }
+      versionIndex++;
+    }
+  }
+  return undefined;
+}

templates/types/streaming/nextjs/app/layout.tsx

@@ -1,7 +1,6 @@
 import type { Metadata } from "next";
 import { Inter } from "next/font/google";
 import "./globals.css";
-import "./markdown.css";
 
 const inter = Inter({ subsets: ["latin"] });
 

templates/types/streaming/nextjs/app/markdown.css

@@ -18,7 +18,7 @@
     "@radix-ui/react-select": "^2.1.1",
     "@radix-ui/react-slot": "^1.0.2",
     "@radix-ui/react-tabs": "^1.1.0",
-    "@llamaindex/chat-ui": "0.0.5",
+    "@llamaindex/chat-ui": "0.0.7",
     "ai": "3.3.42",
     "ajv": "^8.12.0",
     "class-variance-authority": "^0.7.0",