Skip to content

Commit b58913e

Browse files
leehuwujleehuwuj
committed
add new query index and weather card for agent workflows
1 parent 5a230be commit b58913e

File tree

2 files changed

+377
-0
lines changed

2 files changed

+377
-0
lines changed

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

+81
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,81 @@
1+
"use client";
2+
3+
import { getCustomAnnotation, useChatMessage } from "@llamaindex/chat-ui";
4+
import { ChatEvents } from "@llamaindex/chat-ui/widgets";
5+
import { useMemo } from "react";
6+
import { z } from "zod";
7+
8+
const QueryIndexSchema = z.object({
9+
tool_name: z.literal("query_index"),
10+
tool_kwargs: z.object({
11+
input: z.string(),
12+
}),
13+
tool_id: z.string(),
14+
tool_output: z.optional(
15+
z
16+
.object({
17+
content: z.string(),
18+
tool_name: z.string(),
19+
raw_input: z.record(z.unknown()),
20+
raw_output: z.record(z.unknown()),
21+
is_error: z.boolean().optional(),
22+
})
23+
.optional(),
24+
),
25+
return_direct: z.boolean().optional(),
26+
});
27+
type QueryIndex = z.infer<typeof QueryIndexSchema>;
28+
29+
type GroupedIndexQuery = {
30+
initial: QueryIndex;
31+
output?: QueryIndex;
32+
};
33+
34+
export function RetrieverComponent() {
35+
const { message } = useChatMessage();
36+
37+
const queryIndexEvents = getCustomAnnotation<QueryIndex>(
38+
message.annotations,
39+
(annotation) => {
40+
const result = QueryIndexSchema.safeParse(annotation);
41+
return result.success;
42+
},
43+
);
44+
45+
// Group events by tool_id and render them in a single ChatEvents component
46+
const groupedIndexQueries = useMemo(() => {
47+
const groups = new Map<string, GroupedIndexQuery>();
48+
49+
queryIndexEvents?.forEach((event) => {
50+
groups.set(event.tool_id, { initial: event });
51+
});
52+
53+
return Array.from(groups.values());
54+
}, [queryIndexEvents]);
55+
56+
return (
57+
<div className="space-y-4">
58+
{groupedIndexQueries.map(({ initial }) => {
59+
const eventData = [
60+
{
61+
title: `Searching index with query: ${initial.tool_kwargs.input}`,
62+
},
63+
];
64+
65+
if (initial.tool_output) {
66+
eventData.push({
67+
title: `Got ${JSON.stringify((initial.tool_output?.raw_output as any).source_nodes?.length ?? 0)} sources for query: ${initial.tool_kwargs.input}`,
68+
});
69+
}
70+
71+
return (
72+
<ChatEvents
73+
key={initial.tool_id}
74+
data={eventData}
75+
showLoading={!initial.tool_output}
76+
/>
77+
);
78+
})}
79+
</div>
80+
);
81+
}

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

+296
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,296 @@
1+
import { getCustomAnnotation, useChatMessage } from "@llamaindex/chat-ui";
2+
import { ChatEvents } from "@llamaindex/chat-ui/widgets";
3+
import { useMemo } from "react";
4+
import { z } from "zod";
5+
6+
export interface WeatherData {
7+
latitude: number;
8+
longitude: number;
9+
generationtime_ms: number;
10+
utc_offset_seconds: number;
11+
timezone: string;
12+
timezone_abbreviation: string;
13+
elevation: number;
14+
current_units: {
15+
time: string;
16+
interval: string;
17+
temperature_2m: string;
18+
weather_code: string;
19+
};
20+
current: {
21+
time: string;
22+
interval: number;
23+
temperature_2m: number;
24+
weather_code: number;
25+
};
26+
hourly_units: {
27+
time: string;
28+
temperature_2m: string;
29+
weather_code: string;
30+
};
31+
hourly: {
32+
time: string[];
33+
temperature_2m: number[];
34+
weather_code: number[];
35+
};
36+
daily_units: {
37+
time: string;
38+
weather_code: string;
39+
};
40+
daily: {
41+
time: string[];
42+
weather_code: number[];
43+
};
44+
}
45+
46+
// Follow WMO Weather interpretation codes (WW)
47+
const weatherCodeDisplayMap: Record<
48+
string,
49+
{
50+
icon: React.ReactNode;
51+
status: string;
52+
}
53+
> = {
54+
"0": {
55+
icon: <span>☀️</span>,
56+
status: "Clear sky",
57+
},
58+
"1": {
59+
icon: <span>🌤️</span>,
60+
status: "Mainly clear",
61+
},
62+
"2": {
63+
icon: <span>☁️</span>,
64+
status: "Partly cloudy",
65+
},
66+
"3": {
67+
icon: <span>☁️</span>,
68+
status: "Overcast",
69+
},
70+
"45": {
71+
icon: <span>🌫️</span>,
72+
status: "Fog",
73+
},
74+
"48": {
75+
icon: <span>🌫️</span>,
76+
status: "Depositing rime fog",
77+
},
78+
"51": {
79+
icon: <span>🌧️</span>,
80+
status: "Drizzle",
81+
},
82+
"53": {
83+
icon: <span>🌧️</span>,
84+
status: "Drizzle",
85+
},
86+
"55": {
87+
icon: <span>🌧️</span>,
88+
status: "Drizzle",
89+
},
90+
"56": {
91+
icon: <span>🌧️</span>,
92+
status: "Freezing Drizzle",
93+
},
94+
"57": {
95+
icon: <span>🌧️</span>,
96+
status: "Freezing Drizzle",
97+
},
98+
"61": {
99+
icon: <span>🌧️</span>,
100+
status: "Rain",
101+
},
102+
"63": {
103+
icon: <span>🌧️</span>,
104+
status: "Rain",
105+
},
106+
"65": {
107+
icon: <span>🌧️</span>,
108+
status: "Rain",
109+
},
110+
"66": {
111+
icon: <span>🌧️</span>,
112+
status: "Freezing Rain",
113+
},
114+
"67": {
115+
icon: <span>🌧️</span>,
116+
status: "Freezing Rain",
117+
},
118+
"71": {
119+
icon: <span>❄️</span>,
120+
status: "Snow fall",
121+
},
122+
"73": {
123+
icon: <span>❄️</span>,
124+
status: "Snow fall",
125+
},
126+
"75": {
127+
icon: <span>❄️</span>,
128+
status: "Snow fall",
129+
},
130+
"77": {
131+
icon: <span>❄️</span>,
132+
status: "Snow grains",
133+
},
134+
"80": {
135+
icon: <span>🌧️</span>,
136+
status: "Rain showers",
137+
},
138+
"81": {
139+
icon: <span>🌧️</span>,
140+
status: "Rain showers",
141+
},
142+
"82": {
143+
icon: <span>🌧️</span>,
144+
status: "Rain showers",
145+
},
146+
"85": {
147+
icon: <span>❄️</span>,
148+
status: "Snow showers",
149+
},
150+
"86": {
151+
icon: <span>❄️</span>,
152+
status: "Snow showers",
153+
},
154+
"95": {
155+
icon: <span>⛈️</span>,
156+
status: "Thunderstorm",
157+
},
158+
"96": {
159+
icon: <span>⛈️</span>,
160+
status: "Thunderstorm",
161+
},
162+
"99": {
163+
icon: <span>⛈️</span>,
164+
status: "Thunderstorm",
165+
},
166+
};
167+
168+
const displayDay = (time: string) => {
169+
return new Date(time).toLocaleDateString("en-US", {
170+
weekday: "long",
171+
});
172+
};
173+
174+
export function WeatherCard({ data }: { data: WeatherData }) {
175+
const currentDayString = new Date(data.current.time).toLocaleDateString(
176+
"en-US",
177+
{
178+
weekday: "long",
179+
month: "long",
180+
day: "numeric",
181+
},
182+
);
183+
184+
return (
185+
<div className="bg-[#61B9F2] rounded-2xl shadow-xl p-5 space-y-4 text-white w-fit">
186+
<div className="flex justify-between">
187+
<div className="space-y-2">
188+
<div className="text-xl">{currentDayString}</div>
189+
<div className="text-5xl font-semibold flex gap-4">
190+
<span>
191+
{data.current.temperature_2m} {data.current_units.temperature_2m}
192+
</span>
193+
{weatherCodeDisplayMap[data.current.weather_code].icon}
194+
</div>
195+
</div>
196+
<span className="text-xl">
197+
{weatherCodeDisplayMap[data.current.weather_code].status}
198+
</span>
199+
</div>
200+
<div className="gap-2 grid grid-cols-6">
201+
{data.daily.time.map((time, index) => {
202+
if (index === 0) return null; // skip the current day
203+
return (
204+
<div key={time} className="flex flex-col items-center gap-4">
205+
<span>{displayDay(time)}</span>
206+
<div className="text-4xl">
207+
{weatherCodeDisplayMap[data.daily.weather_code[index]].icon}
208+
</div>
209+
<span className="text-sm">
210+
{weatherCodeDisplayMap[data.daily.weather_code[index]].status}
211+
</span>
212+
</div>
213+
);
214+
})}
215+
</div>
216+
</div>
217+
);
218+
}
219+
220+
// A new component for the weather tool which uses the WeatherCard component with the new data schema from agent workflow events
221+
const WeatherToolSchema = z.object({
222+
tool_name: z.literal("get_weather_information"),
223+
tool_kwargs: z.object({
224+
location: z.string(),
225+
}),
226+
tool_id: z.string(),
227+
tool_output: z.optional(
228+
z
229+
.object({
230+
content: z.string(),
231+
tool_name: z.string(),
232+
raw_input: z.record(z.unknown()),
233+
raw_output: z.custom<WeatherData>(),
234+
is_error: z.boolean().optional(),
235+
})
236+
.optional(),
237+
),
238+
return_direct: z.boolean().optional(),
239+
});
240+
241+
type WeatherTool = z.infer<typeof WeatherToolSchema>;
242+
243+
type GroupedWeatherQuery = {
244+
initial: WeatherTool;
245+
output?: WeatherTool;
246+
};
247+
248+
export function WeatherToolComponent() {
249+
const { message } = useChatMessage();
250+
251+
const weatherEvents = getCustomAnnotation<WeatherTool>(
252+
message.annotations,
253+
(annotation: unknown) => {
254+
const result = WeatherToolSchema.safeParse(annotation);
255+
return result.success;
256+
},
257+
);
258+
259+
// Group events by tool_id
260+
const groupedWeatherQueries = useMemo(() => {
261+
const groups = new Map<string, GroupedWeatherQuery>();
262+
263+
weatherEvents?.forEach((event: WeatherTool) => {
264+
groups.set(event.tool_id, { initial: event });
265+
});
266+
267+
return Array.from(groups.values());
268+
}, [weatherEvents]);
269+
270+
return (
271+
<div className="space-y-4">
272+
{groupedWeatherQueries.map(({ initial }) => {
273+
if (!initial.tool_output?.raw_output) {
274+
return (
275+
<ChatEvents
276+
key={initial.tool_id}
277+
data={[
278+
{
279+
title: `Loading weather information for ${initial.tool_kwargs.location}...`,
280+
},
281+
]}
282+
showLoading={true}
283+
/>
284+
);
285+
}
286+
287+
return (
288+
<WeatherCard
289+
key={initial.tool_id}
290+
data={initial.tool_output.raw_output as WeatherData}
291+
/>
292+
);
293+
})}
294+
</div>
295+
);
296+
}

0 commit comments

Comments
 (0)