Skip to content

Commit 54bf408

Browse files
lgrammellgrammel
authoredJun 24, 2024··
feat (ai/react): control request body in useChat (#1991)
·
@ai-sdk/xai@2.0.0-canary.7ai@3.2.8
1 parent 5f3bac8 commit 54bf408

File tree

12 files changed

+303
-57
lines changed

12 files changed

+303
-57
lines changed
 

Diff for: ‎.changeset/thirty-badgers-hide.md

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
---
2+
'@ai-sdk/ui-utils': patch
3+
'@ai-sdk/react': patch
4+
'ai': patch
5+
---
6+
7+
feat (ai/react): control request body in useChat

Diff for: ‎content/docs/07-reference/ai-sdk-ui/01-use-chat.mdx

+7
Original file line numberDiff line numberDiff line change
@@ -127,6 +127,13 @@ Allows you to easily create a conversational user interface for your chatbot app
127127
description:
128128
'An optional literal that sets the mode of the stream to be used. Defaults to `stream-data`. If set to `text`, the stream will be treated as a text stream.',
129129
},
130+
{
131+
name: 'experimental_prepareRequestBody',
132+
type: '(options: { messages: Message[]; requestData?: Record<string, string>; requestBody?: object }) => JSONValue',
133+
optional: true,
134+
description:
135+
'Experimental (React only). When a function is provided, it will be used to prepare the request body for the chat API. This can be useful for customizing the request body based on the messages and data in the chat.',
136+
},
130137
]}
131138
/>
132139

Diff for: ‎content/examples/02-next-pages/05-chat/15-use-chat-custom-body.mdx

+92
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
---
2+
title: Custom body content for useChat
3+
description: Learn how to control the body content that useChat sends to the server
4+
---
5+
6+
# useChat: custom body content
7+
8+
<Note type="warning">
9+
`experimental_prepareRequestBody` is an experimental feature and only
10+
available in React.
11+
</Note>
12+
13+
By default, `useChat` sends all messages as well as information from the request to the server.
14+
However, it is often desirable to control the body content that is sent to the server, e.g. to:
15+
16+
- only send the last message
17+
- send additional data along with the message
18+
- change the structure of the request body
19+
20+
The `experimental_prepareRequestBody` option allows you to customize the body content that is sent to the server.
21+
The function receives the message list, the request data, and the request body from the append call.
22+
It should return the body content that will be sent to the server.
23+
24+
## Example
25+
26+
This example shows how to only send the text of the last message to the server.
27+
This can be useful if you want to reduce the amount of data sent to the server.
28+
29+
### Client
30+
31+
```typescript filename='app/page.tsx' highlight="7-10"
32+
'use client';
33+
34+
import { useChat } from 'ai/react';
35+
36+
export default function Chat() {
37+
const { messages, input, handleInputChange, handleSubmit } = useChat({
38+
experimental_prepareRequestBody: ({ messages }) => {
39+
// e.g. only the text of the last message:
40+
return messages[messages.length - 1].content;
41+
},
42+
});
43+
44+
return (
45+
<div>
46+
{messages.map(m => (
47+
<div key={m.id}>
48+
{m.role === 'user' ? 'User: ' : 'AI: '}
49+
{m.content}
50+
</div>
51+
))}
52+
53+
<form onSubmit={handleSubmit}>
54+
<input value={input} onChange={handleInputChange} />
55+
</form>
56+
</div>
57+
);
58+
}
59+
```
60+
61+
### Server
62+
63+
We need to adjust the server to only receive the text of the last message.
64+
The rest of the message history can be loaded from storage.
65+
66+
```tsx filename='app/api/chat/route.ts' highlight="8,9,23"
67+
import { openai } from '@ai-sdk/openai'
68+
import { streamText } from 'ai'
69+
70+
// Allow streaming responses up to 30 seconds
71+
export const maxDuration = 30
72+
73+
export async function POST(req: Request) {
74+
// we receive only the text from the last message
75+
const text = await req.json()
76+
77+
// e.g. load message history from storage
78+
const history = await loadHistory()
79+
80+
// Call the language model
81+
const result = await streamText({
82+
model: openai('gpt-4-turbo'),
83+
messages: [...history, { role: 'user', content: text }]
84+
onFinish({ text }) {
85+
// e.g. save the message and the response to storage
86+
}
87+
})
88+
89+
// Respond with the stream
90+
return result.toAIStreamResponse()
91+
}
92+
```

Diff for: ‎content/examples/02-next-pages/05-chat/index.mdx

+6
Original file line numberDiff line numberDiff line change
@@ -24,5 +24,11 @@ So far you've learned how to generate text and structured data using single prom
2424
description: 'Learn how to add image inputs to a chat completion.',
2525
href: '/examples/next-pages/chat/use-chat-image-input',
2626
},
27+
{
28+
title: 'Custom body content for useChat',
29+
description:
30+
'Learn how to control the body content that useChat sends to the server.',
31+
href: '/examples/next-pages/chat/use-chat-custom-body',
32+
},
2733
]}
2834
/>

Diff for: ‎packages/core/svelte/use-chat.ts

+2
Original file line numberDiff line numberDiff line change
@@ -117,6 +117,8 @@ const getStreamedResponse = async (
117117
api,
118118
messages: constructedMessagesPayload,
119119
body: {
120+
messages: constructedMessagesPayload,
121+
data: chatRequest.data,
120122
...extraMetadata.body,
121123
...chatRequest.options?.body,
122124
...(chatRequest.functions !== undefined && {

Diff for: ‎packages/react/src/use-chat.ts

+29-1
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,11 @@ const getStreamedResponse = async (
8787
onResponse?: (response: Response) => void | Promise<void>,
8888
onToolCall?: UseChatOptions['onToolCall'],
8989
sendExtraMessageFields?: boolean,
90+
experimental_prepareRequestBody?: (options: {
91+
messages: Message[];
92+
requestData?: Record<string, string>;
93+
requestBody?: object;
94+
}) => JSONValue,
9095
) => {
9196
// Do an optimistic update to the chat state to show the updated messages
9297
// immediately.
@@ -123,7 +128,12 @@ const getStreamedResponse = async (
123128
return await callChatApi({
124129
api,
125130
messages: constructedMessagesPayload,
126-
body: {
131+
body: experimental_prepareRequestBody?.({
132+
messages: chatRequest.messages,
133+
requestData: chatRequest.data,
134+
requestBody: chatRequest.options?.body,
135+
}) ?? {
136+
messages: constructedMessagesPayload,
127137
data: chatRequest.data,
128138
...extraMetadataRef.current.body,
129139
...chatRequest.options?.body,
@@ -170,6 +180,7 @@ export function useChat({
170180
experimental_onFunctionCall,
171181
experimental_onToolCall,
172182
onToolCall,
183+
experimental_prepareRequestBody,
173184
experimental_maxAutomaticRoundtrips = 0,
174185
maxAutomaticRoundtrips = experimental_maxAutomaticRoundtrips,
175186
maxToolRoundtrips = maxAutomaticRoundtrips,
@@ -194,6 +205,21 @@ export function useChat({
194205
*/
195206
maxAutomaticRoundtrips?: number;
196207

208+
/**
209+
* Experimental (React only). When a function is provided, it will be used
210+
* to prepare the request body for the chat API. This can be useful for
211+
* customizing the request body based on the messages and data in the chat.
212+
*
213+
* @param messages The current messages in the chat.
214+
* @param requestData The data object passed in the chat request.
215+
* @param requestBody The request body object passed in the chat request.
216+
*/
217+
experimental_prepareRequestBody?: (options: {
218+
messages: Message[];
219+
requestData?: Record<string, string>;
220+
requestBody?: object;
221+
}) => JSONValue;
222+
197223
/**
198224
Maximal number of automatic roundtrips for tool calls.
199225
@@ -308,6 +334,7 @@ By default, it's set to 0, which will disable the feature.
308334
onResponse,
309335
onToolCall,
310336
sendExtraMessageFields,
337+
experimental_prepareRequestBody,
311338
),
312339
experimental_onFunctionCall,
313340
experimental_onToolCall,
@@ -367,6 +394,7 @@ By default, it's set to 0, which will disable the feature.
367394
sendExtraMessageFields,
368395
experimental_onFunctionCall,
369396
experimental_onToolCall,
397+
experimental_prepareRequestBody,
370398
onToolCall,
371399
maxToolRoundtrips,
372400
messagesRef,

Diff for: ‎packages/react/src/use-chat.ui.test.tsx

+75
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,81 @@ describe('text stream', () => {
209209
});
210210
});
211211

212+
describe('prepareRequestBody', () => {
213+
let bodyOptions: any;
214+
215+
const TestComponent = () => {
216+
const { messages, append, isLoading } = useChat({
217+
experimental_prepareRequestBody(options) {
218+
bodyOptions = options;
219+
return 'test-request-body';
220+
},
221+
});
222+
223+
return (
224+
<div>
225+
<div data-testid="loading">{isLoading.toString()}</div>
226+
{messages.map((m, idx) => (
227+
<div data-testid={`message-${idx}`} key={m.id}>
228+
{m.role === 'user' ? 'User: ' : 'AI: '}
229+
{m.content}
230+
</div>
231+
))}
232+
233+
<button
234+
data-testid="do-append"
235+
onClick={() => {
236+
append(
237+
{ role: 'user', content: 'hi' },
238+
{
239+
data: { 'test-data-key': 'test-data-value' },
240+
options: {
241+
body: { 'request-body-key': 'request-body-value' },
242+
},
243+
},
244+
);
245+
}}
246+
/>
247+
</div>
248+
);
249+
};
250+
251+
beforeEach(() => {
252+
render(<TestComponent />);
253+
});
254+
255+
afterEach(() => {
256+
bodyOptions = undefined;
257+
vi.restoreAllMocks();
258+
cleanup();
259+
});
260+
261+
it('should show streamed response', async () => {
262+
const { requestBody } = mockFetchDataStream({
263+
url: 'https://example.com/api/chat',
264+
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
265+
});
266+
267+
await userEvent.click(screen.getByTestId('do-append'));
268+
269+
await screen.findByTestId('message-0');
270+
expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
271+
272+
expect(bodyOptions).toStrictEqual({
273+
messages: [{ role: 'user', content: 'hi', id: expect.any(String) }],
274+
requestData: { 'test-data-key': 'test-data-value' },
275+
requestBody: { 'request-body-key': 'request-body-value' },
276+
});
277+
278+
expect(await requestBody).toBe('"test-request-body"');
279+
280+
await screen.findByTestId('message-1');
281+
expect(screen.getByTestId('message-1')).toHaveTextContent(
282+
'AI: Hello, world.',
283+
);
284+
});
285+
});
286+
212287
describe('onToolCall', () => {
213288
const TestComponent = () => {
214289
const { messages, append } = useChat({

Diff for: ‎packages/solid/src/use-chat.ts

+23-20
Original file line numberDiff line numberDiff line change
@@ -145,29 +145,32 @@ export function useChat({
145145
getStreamedResponse: async () => {
146146
const existingData = streamData() ?? [];
147147

148+
const constructedMessagesPayload = sendExtraMessageFields
149+
? chatRequest.messages
150+
: chatRequest.messages.map(
151+
({
152+
role,
153+
content,
154+
name,
155+
data,
156+
annotations,
157+
function_call,
158+
}) => ({
159+
role,
160+
content,
161+
...(name !== undefined && { name }),
162+
...(data !== undefined && { data }),
163+
...(annotations !== undefined && { annotations }),
164+
// outdated function/tool call handling (TODO deprecate):
165+
...(function_call !== undefined && { function_call }),
166+
}),
167+
);
168+
148169
return await callChatApi({
149170
api,
150-
messages: sendExtraMessageFields
151-
? chatRequest.messages
152-
: chatRequest.messages.map(
153-
({
154-
role,
155-
content,
156-
name,
157-
data,
158-
annotations,
159-
function_call,
160-
}) => ({
161-
role,
162-
content,
163-
...(name !== undefined && { name }),
164-
...(data !== undefined && { data }),
165-
...(annotations !== undefined && { annotations }),
166-
// outdated function/tool call handling (TODO deprecate):
167-
...(function_call !== undefined && { function_call }),
168-
}),
169-
),
171+
messages: constructedMessagesPayload,
170172
body: {
173+
messages: constructedMessagesPayload,
171174
data: chatRequest.data,
172175
...body,
173176
...options?.body,

Diff for: ‎packages/svelte/src/use-chat.ts

+2
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,8 @@ const getStreamedResponse = async (
115115
api,
116116
messages: constructedMessagesPayload,
117117
body: {
118+
messages: constructedMessagesPayload,
119+
data: chatRequest.data,
118120
...extraMetadata.body,
119121
...chatRequest.options?.body,
120122
...(chatRequest.functions !== undefined && {

Diff for: ‎packages/ui-utils/src/call-chat-api.ts

+1-4
Original file line numberDiff line numberDiff line change
@@ -33,10 +33,7 @@ export async function callChatApi({
3333
}) {
3434
const response = await fetch(api, {
3535
method: 'POST',
36-
body: JSON.stringify({
37-
messages,
38-
...body,
39-
}),
36+
body: JSON.stringify(body),
4037
headers: {
4138
'Content-Type': 'application/json',
4239
...headers,

Diff for: ‎packages/ui-utils/src/types.ts

+36-12
Original file line numberDiff line numberDiff line change
@@ -27,18 +27,28 @@ export interface FunctionCall {
2727
* The tool calls generated by the model, such as function calls.
2828
*/
2929
export interface ToolCall {
30-
// The ID of the tool call.
30+
/**
31+
* The ID of the tool call.
32+
*/
3133
id: string;
3234

33-
// The type of the tool. Currently, only `function` is supported.
35+
/**
36+
* The type of the tool. Currently, only `function` is supported.
37+
*/
3438
type: string;
3539

36-
// The function that the model called.
40+
/**
41+
* The function that the model called.
42+
*/
3743
function: {
38-
// The name of the function.
44+
/**
45+
* The name of the function.
46+
*/
3947
name: string;
4048

41-
// The arguments to call the function with, as generated by the model in JSON
49+
/**
50+
* The arguments to call the function with, as generated by the model in JSON
51+
*/
4252
arguments: string;
4353
};
4454
}
@@ -116,7 +126,9 @@ export interface Message {
116126

117127
content: string;
118128

119-
// @deprecated
129+
/**
130+
* @deprecated Use AI SDK 3.1 `toolInvocations` instead.
131+
*/
120132
tool_call_id?: string;
121133

122134
/**
@@ -180,14 +192,26 @@ export type CreateMessage = Omit<Message, 'id'> & {
180192
export type ChatRequest = {
181193
messages: Message[];
182194
options?: RequestOptions;
183-
// @deprecated
195+
data?: Record<string, string>;
196+
197+
/**
198+
* @deprecated
199+
*/
184200
functions?: Array<Function>;
185-
// @deprecated
201+
202+
/**
203+
* @deprecated
204+
*/
186205
function_call?: FunctionCall;
187-
data?: Record<string, string>;
188-
// @deprecated
206+
207+
/**
208+
* @deprecated
209+
*/
189210
tools?: Array<Tool>;
190-
// @deprecated
211+
212+
/**
213+
* @deprecated
214+
*/
191215
tool_choice?: ToolChoice;
192216
};
193217

@@ -437,7 +461,7 @@ export type JSONValue =
437461
| string
438462
| number
439463
| boolean
440-
| { [x: string]: JSONValue }
464+
| { [value: string]: JSONValue }
441465
| Array<JSONValue>;
442466

443467
export type AssistantMessage = {

Diff for: ‎packages/vue/src/use-chat.ts

+23-20
Original file line numberDiff line numberDiff line change
@@ -138,29 +138,32 @@ export function useChat({
138138
getStreamedResponse: async () => {
139139
const existingData = (streamData.value ?? []) as JSONValue[];
140140

141+
const constructedMessagesPayload = sendExtraMessageFields
142+
? chatRequest.messages
143+
: chatRequest.messages.map(
144+
({
145+
role,
146+
content,
147+
name,
148+
data,
149+
annotations,
150+
function_call,
151+
}) => ({
152+
role,
153+
content,
154+
...(name !== undefined && { name }),
155+
...(data !== undefined && { data }),
156+
...(annotations !== undefined && { annotations }),
157+
// outdated function/tool call handling (TODO deprecate):
158+
...(function_call !== undefined && { function_call }),
159+
}),
160+
);
161+
141162
return await callChatApi({
142163
api,
143-
messages: sendExtraMessageFields
144-
? chatRequest.messages
145-
: chatRequest.messages.map(
146-
({
147-
role,
148-
content,
149-
name,
150-
data,
151-
annotations,
152-
function_call,
153-
}) => ({
154-
role,
155-
content,
156-
...(name !== undefined && { name }),
157-
...(data !== undefined && { data }),
158-
...(annotations !== undefined && { annotations }),
159-
// outdated function/tool call handling (TODO deprecate):
160-
...(function_call !== undefined && { function_call }),
161-
}),
162-
),
164+
messages: constructedMessagesPayload,
163165
body: {
166+
messages: constructedMessagesPayload,
164167
data: chatRequest.data,
165168
...unref(body), // Use unref to unwrap the ref value
166169
...options?.body,

0 commit comments

Comments
 (0)
Please sign in to comment.