add new query index and weather card for agent workflows

Author: leehuwuj

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

@@ -1,11 +1,15 @@
 import { ChatMessage } from "@llamaindex/chat-ui";
 import { DeepResearchCard } from "./custom/deep-research-card";
 import { ToolAnnotations } from "./tools/chat-tools";
+import { RetrieverComponent } from "./tools/query-index";
+import { WeatherToolComponent } from "./tools/weather-card";
 
 export function ChatMessageContent() {
   return (
     <ChatMessage.Content>
       <ChatMessage.Content.Event />
+      <RetrieverComponent />
+      <WeatherToolComponent />
       <ChatMessage.Content.AgentEvent />
       <DeepResearchCard />
       <ToolAnnotations />

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

@@ -0,0 +1,81 @@
+"use client";
+
+import { getCustomAnnotation, useChatMessage } from "@llamaindex/chat-ui";
+import { ChatEvents } from "@llamaindex/chat-ui/widgets";
+import { useMemo } from "react";
+import { z } from "zod";
+
+const QueryIndexSchema = z.object({
+  tool_name: z.literal("query_index"),
+  tool_kwargs: z.object({
+    input: z.string(),
+  }),
+  tool_id: z.string(),
+  tool_output: z.optional(
+    z
+      .object({
+        content: z.string(),
+        tool_name: z.string(),
+        raw_input: z.record(z.unknown()),
+        raw_output: z.record(z.unknown()),
+        is_error: z.boolean().optional(),
+      })
+      .optional(),
+  ),
+  return_direct: z.boolean().optional(),
+});
+type QueryIndex = z.infer<typeof QueryIndexSchema>;
+
+type GroupedIndexQuery = {
+  initial: QueryIndex;
+  output?: QueryIndex;
+};
+
+export function RetrieverComponent() {
+  const { message } = useChatMessage();
+
+  const queryIndexEvents = getCustomAnnotation<QueryIndex>(
+    message.annotations,
+    (annotation) => {
+      const result = QueryIndexSchema.safeParse(annotation);
+      return result.success;
+    },
+  );
+
+  // Group events by tool_id and render them in a single ChatEvents component
+  const groupedIndexQueries = useMemo(() => {
+    const groups = new Map<string, GroupedIndexQuery>();
+
+    queryIndexEvents?.forEach((event) => {
+      groups.set(event.tool_id, { initial: event });
+    });
+
+    return Array.from(groups.values());
+  }, [queryIndexEvents]);
+
+  return (
+    <div className="space-y-4">
+      {groupedIndexQueries.map(({ initial }) => {
+        const eventData = [
+          {
+            title: `Searching index with query: ${initial.tool_kwargs.input}`,
+          },
+        ];
+
+        if (initial.tool_output) {
+          eventData.push({
+            title: `Got ${JSON.stringify((initial.tool_output?.raw_output as any).source_nodes?.length ?? 0)} sources for query: ${initial.tool_kwargs.input}`,
+          });
+        }
+
+        return (
+          <ChatEvents
+            key={initial.tool_id}
+            data={eventData}
+            showLoading={!initial.tool_output}
+          />
+        );
+      })}
+    </div>
+  );
+}

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

@@ -1,3 +1,8 @@
+import { getCustomAnnotation, useChatMessage } from "@llamaindex/chat-ui";
+import { ChatEvents } from "@llamaindex/chat-ui/widgets";
+import { useMemo } from "react";
+import { z } from "zod";
+
 export interface WeatherData {
   latitude: number;
   longitude: number;
@@ -211,3 +216,81 @@ export function WeatherCard({ data }: { data: WeatherData }) {
     </div>
   );
 }
+
+// A new component for the weather tool which uses the WeatherCard component with the new data schema from agent workflow events
+const WeatherToolSchema = z.object({
+  tool_name: z.literal("get_weather_information"),
+  tool_kwargs: z.object({
+    location: z.string(),
+  }),
+  tool_id: z.string(),
+  tool_output: z.optional(
+    z
+      .object({
+        content: z.string(),
+        tool_name: z.string(),
+        raw_input: z.record(z.unknown()),
+        raw_output: z.custom<WeatherData>(),
+        is_error: z.boolean().optional(),
+      })
+      .optional(),
+  ),
+  return_direct: z.boolean().optional(),
+});
+
+type WeatherTool = z.infer<typeof WeatherToolSchema>;
+
+type GroupedWeatherQuery = {
+  initial: WeatherTool;
+  output?: WeatherTool;
+};
+
+export function WeatherToolComponent() {
+  const { message } = useChatMessage();
+
+  const weatherEvents = getCustomAnnotation<WeatherTool>(
+    message.annotations,
+    (annotation: unknown) => {
+      const result = WeatherToolSchema.safeParse(annotation);
+      return result.success;
+    },
+  );
+
+  // Group events by tool_id
+  const groupedWeatherQueries = useMemo(() => {
+    const groups = new Map<string, GroupedWeatherQuery>();
+
+    weatherEvents?.forEach((event: WeatherTool) => {
+      groups.set(event.tool_id, { initial: event });
+    });
+
+    return Array.from(groups.values());
+  }, [weatherEvents]);
+
+  return (
+    <div className="space-y-4">
+      {groupedWeatherQueries.map(({ initial }) => {
+        if (!initial.tool_output?.raw_output) {
+          return (
+            <ChatEvents
+              key={initial.tool_id}
+              data={[
+                {
+                  title: `Loading weather information for ${initial.tool_kwargs.location}...`,
+                },
+              ]}
+              showLoading={true}
+            />
+          );
+        }
+
+        return (
+          <WeatherCard
+            key={initial.tool_id}
+            data={initial.tool_output.raw_output as WeatherData}
+          />
+        );
+      })}
+    </div>
+  );
+}