Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 901e4cc

Browse files
jeremyphilemonjeremyphilemon
authoredMay 15, 2024
Use core messages spec (vercel#337)
1 parent 095550d commit 901e4cc

File tree

7 files changed

+231
-107
lines changed

7 files changed

+231
-107
lines changed
 

‎components/chat.tsx

Lines changed: 2 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,9 +7,8 @@ import { EmptyScreen } from '@/components/empty-screen'
77
import { useLocalStorage } from '@/lib/hooks/use-local-storage'
88
import { useEffect, useState } from 'react'
99
import { useUIState, useAIState } from 'ai/rsc'
10-
import { Session } from '@/lib/types'
10+
import { Message, Session } from '@/lib/types'
1111
import { usePathname, useRouter } from 'next/navigation'
12-
import { Message } from '@/lib/chat/actions'
1312
import { useScrollAnchor } from '@/lib/hooks/use-scroll-anchor'
1413
import { toast } from 'sonner'
1514

@@ -71,7 +70,7 @@ export function Chat({ id, className, session, missingKeys }: ChatProps) {
7170
) : (
7271
<EmptyScreen />
7372
)}
74-
<div className="h-px w-full" ref={visibilityRef} />
73+
<div className="w-full h-px" ref={visibilityRef} />
7574
</div>
7675
<ChatPanel
7776
id={id}

‎components/stocks/message.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import { CodeBlock } from '../ui/codeblock'
77
import { MemoizedReactMarkdown } from '../markdown'
88
import remarkGfm from 'remark-gfm'
99
import remarkMath from 'remark-math'
10-
import { StreamableValue } from 'ai/rsc'
10+
import { StreamableValue, useStreamableValue } from 'ai/rsc'
1111
import { useStreamableText } from '@/lib/hooks/use-streamable-text'
1212

1313
// Different types of message bubbles.

‎components/stocks/stock-purchase.tsx

Lines changed: 14 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ interface Purchase {
1414
}
1515

1616
export function Purchase({
17-
props: { numberOfShares, symbol, price, status = 'requires_action' }
17+
props: { numberOfShares, symbol, price, status = 'expired' }
1818
}: {
1919
props: Purchase
2020
}) {
@@ -60,8 +60,8 @@ export function Purchase({
6060
}
6161

6262
return (
63-
<div className="rounded-xl border bg-zinc-950 p-4 text-green-400">
64-
<div className="float-right inline-block rounded-full bg-white/10 px-2 py-1 text-xs">
63+
<div className="p-4 text-green-400 border rounded-xl bg-zinc-950">
64+
<div className="inline-block float-right px-2 py-1 text-xs rounded-full bg-white/10">
6565
+1.23% ↑
6666
</div>
6767
<div className="text-lg text-zinc-300">{symbol}</div>
@@ -70,7 +70,7 @@ export function Purchase({
7070
<div className="mt-4 text-zinc-200">{purchasingUI}</div>
7171
) : status === 'requires_action' ? (
7272
<>
73-
<div className="relative mt-6 pb-6">
73+
<div className="relative pb-6 mt-6">
7474
<p>Shares to purchase</p>
7575
<input
7676
id="labels-range-input"
@@ -79,46 +79,46 @@ export function Purchase({
7979
onChange={onSliderChange}
8080
min="10"
8181
max="1000"
82-
className="h-1 w-full cursor-pointer appearance-none rounded-lg bg-zinc-600 accent-green-500 dark:bg-zinc-700"
82+
className="w-full h-1 rounded-lg appearance-none cursor-pointer bg-zinc-600 accent-green-500 dark:bg-zinc-700"
8383
/>
84-
<span className="absolute bottom-1 start-0 text-xs text-zinc-400">
84+
<span className="absolute text-xs bottom-1 start-0 text-zinc-400">
8585
10
8686
</span>
87-
<span className="absolute bottom-1 start-1/3 -translate-x-1/2 text-xs text-zinc-400 rtl:translate-x-1/2">
87+
<span className="absolute text-xs -translate-x-1/2 bottom-1 start-1/3 text-zinc-400 rtl:translate-x-1/2">
8888
100
8989
</span>
90-
<span className="absolute bottom-1 start-2/3 -translate-x-1/2 text-xs text-zinc-400 rtl:translate-x-1/2">
90+
<span className="absolute text-xs -translate-x-1/2 bottom-1 start-2/3 text-zinc-400 rtl:translate-x-1/2">
9191
500
9292
</span>
93-
<span className="absolute bottom-1 end-0 text-xs text-zinc-400">
93+
<span className="absolute text-xs bottom-1 end-0 text-zinc-400">
9494
1000
9595
</span>
9696
</div>
9797

9898
<div className="mt-6">
9999
<p>Total cost</p>
100100
<div className="flex flex-wrap items-center text-xl font-bold sm:items-end sm:gap-2 sm:text-3xl">
101-
<div className="flex basis-1/3 flex-col tabular-nums sm:basis-auto sm:flex-row sm:items-center sm:gap-2">
101+
<div className="flex flex-col basis-1/3 tabular-nums sm:basis-auto sm:flex-row sm:items-center sm:gap-2">
102102
{value}
103103
<span className="mb-1 text-sm font-normal text-zinc-600 sm:mb-0 dark:text-zinc-400">
104104
shares
105105
</span>
106106
</div>
107-
<div className="basis-1/3 text-center sm:basis-auto">×</div>
108-
<span className="flex basis-1/3 flex-col tabular-nums sm:basis-auto sm:flex-row sm:items-center sm:gap-2">
107+
<div className="text-center basis-1/3 sm:basis-auto">×</div>
108+
<span className="flex flex-col basis-1/3 tabular-nums sm:basis-auto sm:flex-row sm:items-center sm:gap-2">
109109
${price}
110110
<span className="mb-1 ml-1 text-sm font-normal text-zinc-600 sm:mb-0 dark:text-zinc-400">
111111
per share
112112
</span>
113113
</span>
114-
<div className="mt-2 basis-full border-t border-t-zinc-700 pt-2 text-center sm:mt-0 sm:basis-auto sm:border-0 sm:pt-0 sm:text-left">
114+
<div className="pt-2 mt-2 text-center border-t basis-full border-t-zinc-700 sm:mt-0 sm:basis-auto sm:border-0 sm:pt-0 sm:text-left">
115115
= <span>{formatNumber(value * price)}</span>
116116
</div>
117117
</div>
118118
</div>
119119

120120
<button
121-
className="mt-6 w-full rounded-lg bg-green-400 px-4 py-2 font-bold text-zinc-900 hover:bg-green-500"
121+
className="w-full px-4 py-2 mt-6 font-bold bg-green-400 rounded-lg text-zinc-900 hover:bg-green-500"
122122
onClick={async () => {
123123
const response = await confirmPurchase(symbol, price, value)
124124
setPurchasingUI(response.purchasingUI)

‎lib/chat/actions.tsx

Lines changed: 183 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ import {
3333
} from '@/lib/utils'
3434
import { saveChat } from '@/app/actions'
3535
import { SpinnerMessage, UserMessage } from '@/components/stocks/message'
36-
import { Chat } from '@/lib/types'
36+
import { Chat, Message } from '@/lib/types'
3737
import { auth } from '@/auth'
3838

3939
async function confirmPurchase(symbol: string, price: number, amount: number) {
@@ -85,18 +85,7 @@ async function confirmPurchase(symbol: string, price: number, amount: number) {
8585
aiState.done({
8686
...aiState.get(),
8787
messages: [
88-
...aiState.get().messages.slice(0, -1),
89-
{
90-
id: nanoid(),
91-
role: 'function',
92-
name: 'showStockPurchase',
93-
content: JSON.stringify({
94-
symbol,
95-
price,
96-
defaultAmount: amount,
97-
status: 'completed'
98-
})
99-
},
88+
...aiState.get().messages,
10089
{
10190
id: nanoid(),
10291
role: 'system',
@@ -208,15 +197,35 @@ async function submitUserMessage(content: string) {
208197

209198
await sleep(1000)
210199

200+
const toolCallId = nanoid()
201+
211202
aiState.done({
212203
...aiState.get(),
213204
messages: [
214205
...aiState.get().messages,
215206
{
216207
id: nanoid(),
217-
role: 'function',
218-
name: 'listStocks',
219-
content: JSON.stringify(stocks)
208+
role: 'assistant',
209+
content: [
210+
{
211+
type: 'tool-call',
212+
toolName: 'listStocks',
213+
toolCallId,
214+
args: { stocks }
215+
}
216+
]
217+
},
218+
{
219+
id: nanoid(),
220+
role: 'tool',
221+
content: [
222+
{
223+
type: 'tool-result',
224+
toolName: 'listStocks',
225+
toolCallId,
226+
result: stocks
227+
}
228+
]
220229
}
221230
]
222231
})
@@ -249,15 +258,35 @@ async function submitUserMessage(content: string) {
249258

250259
await sleep(1000)
251260

261+
const toolCallId = nanoid()
262+
252263
aiState.done({
253264
...aiState.get(),
254265
messages: [
255266
...aiState.get().messages,
256267
{
257268
id: nanoid(),
258-
role: 'function',
259-
name: 'showStockPrice',
260-
content: JSON.stringify({ symbol, price, delta })
269+
role: 'assistant',
270+
content: [
271+
{
272+
type: 'tool-call',
273+
toolName: 'showStockPrice',
274+
toolCallId,
275+
args: { symbol, price, delta }
276+
}
277+
]
278+
},
279+
{
280+
id: nanoid(),
281+
role: 'tool',
282+
content: [
283+
{
284+
type: 'tool-result',
285+
toolName: 'showStockPrice',
286+
toolCallId,
287+
result: { symbol, price, delta }
288+
}
289+
]
261290
}
262291
]
263292
})
@@ -286,11 +315,42 @@ async function submitUserMessage(content: string) {
286315
)
287316
}),
288317
generate: async function* ({ symbol, price, numberOfShares = 100 }) {
318+
const toolCallId = nanoid()
319+
289320
if (numberOfShares <= 0 || numberOfShares > 1000) {
290321
aiState.done({
291322
...aiState.get(),
292323
messages: [
293324
...aiState.get().messages,
325+
{
326+
id: nanoid(),
327+
role: 'assistant',
328+
content: [
329+
{
330+
type: 'tool-call',
331+
toolName: 'showStockPurchase',
332+
toolCallId,
333+
args: { symbol, price, numberOfShares }
334+
}
335+
]
336+
},
337+
{
338+
id: nanoid(),
339+
role: 'tool',
340+
content: [
341+
{
342+
type: 'tool-result',
343+
toolName: 'showStockPurchase',
344+
toolCallId,
345+
result: {
346+
symbol,
347+
price,
348+
numberOfShares,
349+
status: 'expired'
350+
}
351+
}
352+
]
353+
},
294354
{
295355
id: nanoid(),
296356
role: 'system',
@@ -300,37 +360,55 @@ async function submitUserMessage(content: string) {
300360
})
301361

302362
return <BotMessage content={'Invalid amount'} />
303-
}
304-
305-
aiState.done({
306-
...aiState.get(),
307-
messages: [
308-
...aiState.get().messages,
309-
{
310-
id: nanoid(),
311-
role: 'function',
312-
name: 'showStockPurchase',
313-
content: JSON.stringify({
314-
symbol,
315-
price,
316-
numberOfShares
317-
})
318-
}
319-
]
320-
})
363+
} else {
364+
aiState.done({
365+
...aiState.get(),
366+
messages: [
367+
...aiState.get().messages,
368+
{
369+
id: nanoid(),
370+
role: 'assistant',
371+
content: [
372+
{
373+
type: 'tool-call',
374+
toolName: 'showStockPurchase',
375+
toolCallId,
376+
args: { symbol, price, numberOfShares }
377+
}
378+
]
379+
},
380+
{
381+
id: nanoid(),
382+
role: 'tool',
383+
content: [
384+
{
385+
type: 'tool-result',
386+
toolName: 'showStockPurchase',
387+
toolCallId,
388+
result: {
389+
symbol,
390+
price,
391+
numberOfShares
392+
}
393+
}
394+
]
395+
}
396+
]
397+
})
321398

322-
return (
323-
<BotCard>
324-
<Purchase
325-
props={{
326-
numberOfShares,
327-
symbol,
328-
price: +price,
329-
status: 'requires_action'
330-
}}
331-
/>
332-
</BotCard>
333-
)
399+
return (
400+
<BotCard>
401+
<Purchase
402+
props={{
403+
numberOfShares,
404+
symbol,
405+
price: +price,
406+
status: 'requires_action'
407+
}}
408+
/>
409+
</BotCard>
410+
)
411+
}
334412
}
335413
},
336414
getEvents: {
@@ -356,15 +434,35 @@ async function submitUserMessage(content: string) {
356434

357435
await sleep(1000)
358436

437+
const toolCallId = nanoid()
438+
359439
aiState.done({
360440
...aiState.get(),
361441
messages: [
362442
...aiState.get().messages,
363443
{
364444
id: nanoid(),
365-
role: 'function',
366-
name: 'getEvents',
367-
content: JSON.stringify(events)
445+
role: 'assistant',
446+
content: [
447+
{
448+
type: 'tool-call',
449+
toolName: 'getEvents',
450+
toolCallId,
451+
args: { events }
452+
}
453+
]
454+
},
455+
{
456+
id: nanoid(),
457+
role: 'tool',
458+
content: [
459+
{
460+
type: 'tool-result',
461+
toolName: 'getEvents',
462+
toolCallId,
463+
result: events
464+
}
465+
]
368466
}
369467
]
370468
})
@@ -385,13 +483,6 @@ async function submitUserMessage(content: string) {
385483
}
386484
}
387485

388-
export type Message = {
389-
role: 'user' | 'assistant' | 'system' | 'function' | 'data' | 'tool'
390-
content: string
391-
id: string
392-
name?: string
393-
}
394-
395486
export type AIState = {
396487
chatId: string
397488
messages: Message[]
@@ -425,7 +516,7 @@ export const AI = createAI<AIState, UIState>({
425516
return
426517
}
427518
},
428-
onSetAIState: async ({ state, done }) => {
519+
onSetAIState: async ({ state }) => {
429520
'use server'
430521

431522
const session = await auth()
@@ -436,7 +527,9 @@ export const AI = createAI<AIState, UIState>({
436527
const createdAt = new Date()
437528
const userId = session.user.id as string
438529
const path = `/chat/${chatId}`
439-
const title = messages[0].content.substring(0, 100)
530+
531+
const firstMessageContent = messages[0].content as string
532+
const title = firstMessageContent.substring(0, 100)
440533

441534
const chat: Chat = {
442535
id: chatId,
@@ -460,28 +553,36 @@ export const getUIStateFromAIState = (aiState: Chat) => {
460553
.map((message, index) => ({
461554
id: `${aiState.chatId}-${index}`,
462555
display:
463-
message.role === 'function' ? (
464-
message.name === 'listStocks' ? (
465-
<BotCard>
466-
<Stocks props={JSON.parse(message.content)} />
467-
</BotCard>
468-
) : message.name === 'showStockPrice' ? (
469-
<BotCard>
470-
<Stock props={JSON.parse(message.content)} />
471-
</BotCard>
472-
) : message.name === 'showStockPurchase' ? (
473-
<BotCard>
474-
<Purchase props={JSON.parse(message.content)} />
475-
</BotCard>
476-
) : message.name === 'getEvents' ? (
477-
<BotCard>
478-
<Events props={JSON.parse(message.content)} />
479-
</BotCard>
480-
) : null
556+
message.role === 'tool' ? (
557+
message.content.map(tool => {
558+
return tool.toolName === 'listStocks' ? (
559+
<BotCard>
560+
{/* TODO: Infer types based on the tool result*/}
561+
{/* @ts-expect-error */}
562+
<Stocks props={tool.result} />
563+
</BotCard>
564+
) : tool.toolName === 'showStockPrice' ? (
565+
<BotCard>
566+
{/* @ts-expect-error */}
567+
<Stock props={tool.result} />
568+
</BotCard>
569+
) : tool.toolName === 'showStockPurchase' ? (
570+
<BotCard>
571+
{/* @ts-expect-error */}
572+
<Purchase props={tool.result} />
573+
</BotCard>
574+
) : tool.toolName === 'getEvents' ? (
575+
<BotCard>
576+
{/* @ts-expect-error */}
577+
<Events props={tool.result} />
578+
</BotCard>
579+
) : null
580+
})
481581
) : message.role === 'user' ? (
482-
<UserMessage>{message.content}</UserMessage>
483-
) : (
582+
<UserMessage>{message.content as string}</UserMessage>
583+
) : message.role === 'assistant' &&
584+
typeof message.content === 'string' ? (
484585
<BotMessage content={message.content} />
485-
)
586+
) : null
486587
}))
487588
}

‎lib/types.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,8 @@
1-
import { Message } from 'ai'
1+
import { CoreMessage } from 'ai'
2+
3+
export type Message = CoreMessage & {
4+
id: string
5+
}
26

37
export interface Chat extends Record<string, any> {
48
id: string

‎package.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
"@vercel/analytics": "^1.1.2",
2828
"@vercel/kv": "^1.0.1",
2929
"@vercel/og": "^0.6.2",
30-
"ai": "^3.1.1",
30+
"ai": "^3.1.5",
3131
"class-variance-authority": "^0.7.0",
3232
"clsx": "^2.1.0",
3333
"d3-scale": "^4.0.2",

‎pnpm-lock.yaml

Lines changed: 25 additions & 5 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)
Please sign in to comment.