add artifact tool component

Author: leehuwuj

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

@@ -1,9 +1,9 @@
 import { ChatMessage } from "@llamaindex/chat-ui";
 import { DeepResearchCard } from "./custom/deep-research-card";
+import { ArtifactToolComponent } from "./tools/artifact";
 import { ToolAnnotations } from "./tools/chat-tools";
 import { ChatSourcesComponent, RetrieverComponent } from "./tools/query-index";
 import { WeatherToolComponent } from "./tools/weather-card";
-
 export function ChatMessageContent() {
   return (
     <ChatMessage.Content>
@@ -13,6 +13,7 @@ export function ChatMessageContent() {
       <WeatherToolComponent />
       <DeepResearchCard />
       <ToolAnnotations />
+      <ArtifactToolComponent />
       <ChatMessage.Content.Image />
       <ChatMessage.Content.Markdown />
       <ChatMessage.Content.DocumentFile />

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

@@ -1,7 +1,13 @@
 "use client";
 
+import {
+  getCustomAnnotation,
+  useChatMessage,
+  useChatUI,
+} from "@llamaindex/chat-ui";
 import { Check, ChevronDown, Code, Copy, Loader2 } from "lucide-react";
-import { useEffect, useRef, useState } from "react";
+import { useEffect, useMemo, useRef, useState } from "react";
+import { z } from "zod";
 import { Button, buttonVariants } from "../../button";
 import {
   Collapsible,
@@ -386,3 +392,59 @@ function closePanel() {
     panel.classList.add("hidden");
   });
 }
+
+const ArtifactToolSchema = z.object({
+  tool_name: z.literal("artifact"),
+  tool_kwargs: z.object({
+    query: z.string(),
+  }),
+  tool_id: z.string(),
+  tool_output: z.object({
+    content: z.string(),
+    tool_name: z.string(),
+    raw_input: z.object({
+      args: z.array(z.unknown()),
+      kwargs: z.object({
+        query: z.string(),
+      }),
+    }),
+    raw_output: z.custom<CodeArtifact>(),
+    is_error: z.boolean(),
+  }),
+  return_direct: z.boolean().optional(),
+});
+
+type ArtifactTool = z.infer<typeof ArtifactToolSchema>;
+
+export function ArtifactToolComponent() {
+  const { message } = useChatMessage();
+  const { messages } = useChatUI();
+
+  const artifactOutputEvent = getCustomAnnotation<ArtifactTool>(
+    message.annotations,
+    (annotation: unknown) => {
+      const result = ArtifactToolSchema.safeParse(annotation);
+      return result.success;
+    },
+  ).at(0);
+
+  const artifactVersion = useMemo(() => {
+    const artifactToolCalls = messages.filter((m) =>
+      m.annotations?.some(
+        (a: unknown) => (a as ArtifactTool).tool_name === "artifact",
+      ),
+    );
+    return artifactToolCalls.length;
+  }, [messages]);
+
+  return (
+    <div className="flex flex-col gap-4">
+      {artifactOutputEvent && (
+        <Artifact
+          artifact={artifactOutputEvent.tool_output.raw_output}
+          version={artifactVersion}
+        />
+      )}
+    </div>
+  );
+}