Skip to content
Snippets Groups Projects
Commit 7e23d779 authored by leehuwuj's avatar leehuwuj
Browse files

add new query index and weather card for agent workflows

parent 5a230bec
No related branches found
No related tags found
No related merge requests found
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 />
......
"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>
);
}
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>
);
}
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment