Skip to content

Commit 66b5892

Browse files
lgrammellgrammel
authoredApr 16, 2024··
Introduce streamMode for useChat / useCompletion. (#1350)
·
@ai-sdk/xai@1.2.3ai@3.0.23
1 parent f272b01 commit 66b5892

23 files changed

+879
-455
lines changed
 

‎.changeset/empty-windows-think.md

+5
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'ai': patch
3+
---
4+
5+
Add streamMode parameter to useChat and useCompletion.

‎examples/solidstart-openai/package.json

+1
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
"vite": "^4.1.4"
1717
},
1818
"dependencies": {
19+
"@ai-sdk/openai": "latest",
1920
"@solidjs/meta": "0.29.3",
2021
"@solidjs/router": "0.8.2",
2122
"ai": "latest",

‎examples/solidstart-openai/src/routes/api/chat/index.ts

+13-19
Original file line numberDiff line numberDiff line change
@@ -1,25 +1,19 @@
1-
import { OpenAIStream, StreamingTextResponse } from 'ai';
2-
import OpenAI from 'openai';
1+
import { openai } from '@ai-sdk/openai';
2+
import { StreamingTextResponse, experimental_streamText } from 'ai';
33
import { APIEvent } from 'solid-start/api';
44

5-
// Create an OpenAI API client
6-
const openai = new OpenAI({
7-
apiKey: process.env['OPENAI_API_KEY'] || '',
8-
});
9-
105
export const POST = async (event: APIEvent) => {
11-
// Extract the `prompt` from the body of the request
12-
const { messages } = await event.request.json();
6+
try {
7+
const { messages } = await event.request.json();
138

14-
// Ask OpenAI for a streaming chat completion given the prompt
15-
const response = await openai.chat.completions.create({
16-
model: 'gpt-3.5-turbo',
17-
stream: true,
18-
messages,
19-
});
9+
const result = await experimental_streamText({
10+
model: openai.chat('gpt-4-turbo-preview'),
11+
messages,
12+
});
2013

21-
// Convert the response into a friendly text-stream
22-
const stream = OpenAIStream(response);
23-
// Respond with the stream
24-
return new StreamingTextResponse(stream);
14+
return new StreamingTextResponse(result.toAIStream());
15+
} catch (error) {
16+
console.error(error);
17+
throw error;
18+
}
2519
};

‎packages/core/react/use-chat.ts

+4
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ const getStreamedResponse = async (
8888
messagesRef: React.MutableRefObject<Message[]>,
8989
abortControllerRef: React.MutableRefObject<AbortController | null>,
9090
generateId: IdGenerator,
91+
streamMode?: 'stream-data' | 'text',
9192
onFinish?: (message: Message) => void,
9293
onResponse?: (response: Response) => void | Promise<void>,
9394
sendExtraMessageFields?: boolean,
@@ -179,6 +180,7 @@ const getStreamedResponse = async (
179180
tool_choice: chatRequest.tool_choice,
180181
}),
181182
},
183+
streamMode,
182184
credentials: extraMetadataRef.current.credentials,
183185
headers: {
184186
...extraMetadataRef.current.headers,
@@ -206,6 +208,7 @@ export function useChat({
206208
sendExtraMessageFields,
207209
experimental_onFunctionCall,
208210
experimental_onToolCall,
211+
streamMode,
209212
onResponse,
210213
onFinish,
211214
onError,
@@ -292,6 +295,7 @@ export function useChat({
292295
messagesRef,
293296
abortControllerRef,
294297
generateId,
298+
streamMode,
295299
onFinish,
296300
onResponse,
297301
sendExtraMessageFields,

‎packages/core/react/use-chat.ui.test.tsx

+160-104
Original file line numberDiff line numberDiff line change
@@ -9,145 +9,201 @@ import {
99
} from '../tests/utils/mock-fetch';
1010
import { useChat } from './use-chat';
1111

12-
const TestComponent = () => {
13-
const [id, setId] = React.useState<string>('first-id');
14-
const { messages, append, error, data, isLoading } = useChat({ id });
15-
16-
return (
17-
<div>
18-
<div data-testid="loading">{isLoading.toString()}</div>
19-
{error && <div data-testid="error">{error.toString()}</div>}
20-
{data && <div data-testid="data">{JSON.stringify(data)}</div>}
21-
{messages.map((m, idx) => (
22-
<div data-testid={`message-${idx}`} key={m.id}>
23-
{m.role === 'user' ? 'User: ' : 'AI: '}
24-
{m.content}
25-
</div>
26-
))}
27-
28-
<button
29-
data-testid="do-append"
30-
onClick={() => {
31-
append({ role: 'user', content: 'hi' });
32-
}}
33-
/>
34-
<button
35-
data-testid="do-change-id"
36-
onClick={() => {
37-
setId('second-id');
38-
}}
39-
/>
40-
</div>
41-
);
42-
};
43-
44-
beforeEach(() => {
45-
render(<TestComponent />);
46-
});
12+
describe('stream data stream', () => {
13+
const TestComponent = () => {
14+
const [id, setId] = React.useState<string>('first-id');
15+
const { messages, append, error, data, isLoading } = useChat({ id });
16+
17+
return (
18+
<div>
19+
<div data-testid="loading">{isLoading.toString()}</div>
20+
{error && <div data-testid="error">{error.toString()}</div>}
21+
{data && <div data-testid="data">{JSON.stringify(data)}</div>}
22+
{messages.map((m, idx) => (
23+
<div data-testid={`message-${idx}`} key={m.id}>
24+
{m.role === 'user' ? 'User: ' : 'AI: '}
25+
{m.content}
26+
</div>
27+
))}
28+
29+
<button
30+
data-testid="do-append"
31+
onClick={() => {
32+
append({ role: 'user', content: 'hi' });
33+
}}
34+
/>
35+
<button
36+
data-testid="do-change-id"
37+
onClick={() => {
38+
setId('second-id');
39+
}}
40+
/>
41+
</div>
42+
);
43+
};
4744

48-
afterEach(() => {
49-
vi.restoreAllMocks();
50-
cleanup();
51-
});
45+
beforeEach(() => {
46+
render(<TestComponent />);
47+
});
5248

53-
test('Shows streamed complex text response', async () => {
54-
mockFetchDataStream({
55-
url: 'https://example.com/api/chat',
56-
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
49+
afterEach(() => {
50+
vi.restoreAllMocks();
51+
cleanup();
5752
});
5853

59-
await userEvent.click(screen.getByTestId('do-append'));
54+
it('should show streamed response', async () => {
55+
mockFetchDataStream({
56+
url: 'https://example.com/api/chat',
57+
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
58+
});
6059

61-
await screen.findByTestId('message-0');
62-
expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
60+
await userEvent.click(screen.getByTestId('do-append'));
6361

64-
await screen.findByTestId('message-1');
65-
expect(screen.getByTestId('message-1')).toHaveTextContent(
66-
'AI: Hello, world.',
67-
);
68-
});
62+
await screen.findByTestId('message-0');
63+
expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
6964

70-
test('Shows streamed complex text response with data', async () => {
71-
mockFetchDataStream({
72-
url: 'https://example.com/api/chat',
73-
chunks: ['2:[{"t1":"v1"}]\n', '0:"Hello"\n'],
65+
await screen.findByTestId('message-1');
66+
expect(screen.getByTestId('message-1')).toHaveTextContent(
67+
'AI: Hello, world.',
68+
);
7469
});
7570

76-
await userEvent.click(screen.getByTestId('do-append'));
71+
it('should show streamed response with data', async () => {
72+
mockFetchDataStream({
73+
url: 'https://example.com/api/chat',
74+
chunks: ['2:[{"t1":"v1"}]\n', '0:"Hello"\n'],
75+
});
7776

78-
await screen.findByTestId('data');
79-
expect(screen.getByTestId('data')).toHaveTextContent('[{"t1":"v1"}]');
77+
await userEvent.click(screen.getByTestId('do-append'));
8078

81-
await screen.findByTestId('message-1');
82-
expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello');
83-
});
79+
await screen.findByTestId('data');
80+
expect(screen.getByTestId('data')).toHaveTextContent('[{"t1":"v1"}]');
8481

85-
test('Shows error response', async () => {
86-
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
82+
await screen.findByTestId('message-1');
83+
expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello');
84+
});
8785

88-
await userEvent.click(screen.getByTestId('do-append'));
86+
it('should show error response', async () => {
87+
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
8988

90-
// TODO bug? the user message does not show up
91-
// await screen.findByTestId('message-0');
92-
// expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
89+
await userEvent.click(screen.getByTestId('do-append'));
9390

94-
await screen.findByTestId('error');
95-
expect(screen.getByTestId('error')).toHaveTextContent('Error: Not found');
96-
});
91+
// TODO bug? the user message does not show up
92+
// await screen.findByTestId('message-0');
93+
// expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
94+
95+
await screen.findByTestId('error');
96+
expect(screen.getByTestId('error')).toHaveTextContent('Error: Not found');
97+
});
98+
99+
describe('loading state', () => {
100+
it('should show loading state', async () => {
101+
let finishGeneration: ((value?: unknown) => void) | undefined;
102+
const finishGenerationPromise = new Promise(resolve => {
103+
finishGeneration = resolve;
104+
});
97105

98-
describe('loading state', () => {
99-
test('should show loading state', async () => {
100-
let finishGeneration: ((value?: unknown) => void) | undefined;
101-
const finishGenerationPromise = new Promise(resolve => {
102-
finishGeneration = resolve;
106+
mockFetchDataStreamWithGenerator({
107+
url: 'https://example.com/api/chat',
108+
chunkGenerator: (async function* generate() {
109+
const encoder = new TextEncoder();
110+
yield encoder.encode('0:"Hello"\n');
111+
await finishGenerationPromise;
112+
})(),
113+
});
114+
115+
await userEvent.click(screen.getByTestId('do-append'));
116+
117+
await screen.findByTestId('loading');
118+
expect(screen.getByTestId('loading')).toHaveTextContent('true');
119+
120+
finishGeneration?.();
121+
122+
await findByText(await screen.findByTestId('loading'), 'false');
123+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
103124
});
104125

105-
mockFetchDataStreamWithGenerator({
106-
url: 'https://example.com/api/chat',
107-
chunkGenerator: (async function* generate() {
108-
const encoder = new TextEncoder();
109-
yield encoder.encode('0:"Hello"\n');
110-
await finishGenerationPromise;
111-
})(),
126+
it('should reset loading state on error', async () => {
127+
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
128+
129+
await userEvent.click(screen.getByTestId('do-append'));
130+
131+
await screen.findByTestId('loading');
132+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
112133
});
134+
});
113135

114-
await userEvent.click(screen.getByTestId('do-append'));
136+
describe('id', () => {
137+
it('should clear out messages when the id changes', async () => {
138+
mockFetchDataStream({
139+
url: 'https://example.com/api/chat',
140+
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
141+
});
115142

116-
await screen.findByTestId('loading');
117-
expect(screen.getByTestId('loading')).toHaveTextContent('true');
143+
await userEvent.click(screen.getByTestId('do-append'));
118144

119-
finishGeneration?.();
145+
await screen.findByTestId('message-1');
146+
expect(screen.getByTestId('message-1')).toHaveTextContent(
147+
'AI: Hello, world.',
148+
);
120149

121-
await findByText(await screen.findByTestId('loading'), 'false');
122-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
150+
await userEvent.click(screen.getByTestId('do-change-id'));
151+
152+
expect(screen.queryByTestId('message-0')).not.toBeInTheDocument();
153+
});
123154
});
155+
});
124156

125-
test('should reset loading state on error', async () => {
126-
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
157+
describe('text stream', () => {
158+
const TestComponent = () => {
159+
const { messages, append } = useChat({
160+
streamMode: 'text',
161+
});
127162

128-
await userEvent.click(screen.getByTestId('do-append'));
163+
return (
164+
<div>
165+
{messages.map((m, idx) => (
166+
<div data-testid={`message-${idx}-text-stream`} key={m.id}>
167+
{m.role === 'user' ? 'User: ' : 'AI: '}
168+
{m.content}
169+
</div>
170+
))}
171+
172+
<button
173+
data-testid="do-append-text-stream"
174+
onClick={() => {
175+
append({ role: 'user', content: 'hi' });
176+
}}
177+
/>
178+
</div>
179+
);
180+
};
129181

130-
await screen.findByTestId('loading');
131-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
182+
beforeEach(() => {
183+
render(<TestComponent />);
132184
});
133-
});
134185

135-
describe('id', () => {
136-
it('should clear out messages when the id changes', async () => {
186+
afterEach(() => {
187+
vi.restoreAllMocks();
188+
cleanup();
189+
});
190+
191+
it('should show streamed response', async () => {
137192
mockFetchDataStream({
138193
url: 'https://example.com/api/chat',
139-
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
194+
chunks: ['Hello', ',', ' world', '.'],
140195
});
141196

142-
await userEvent.click(screen.getByTestId('do-append'));
197+
await userEvent.click(screen.getByTestId('do-append-text-stream'));
143198

144-
await screen.findByTestId('message-1');
145-
expect(screen.getByTestId('message-1')).toHaveTextContent(
146-
'AI: Hello, world.',
199+
await screen.findByTestId('message-0-text-stream');
200+
expect(screen.getByTestId('message-0-text-stream')).toHaveTextContent(
201+
'User: hi',
147202
);
148203

149-
await userEvent.click(screen.getByTestId('do-change-id'));
150-
151-
expect(screen.queryByTestId('message-0')).not.toBeInTheDocument();
204+
await screen.findByTestId('message-1-text-stream');
205+
expect(screen.getByTestId('message-1-text-stream')).toHaveTextContent(
206+
'AI: Hello, world.',
207+
);
152208
});
153209
});

‎packages/core/react/use-completion.ts

+2
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@
6969
credentials,
7070
headers,
7171
body,
72+
streamMode,
7273
onResponse,
7374
onFinish,
7475
onError,
@@ -122,6 +123,7 @@
122123
...extraMetadataRef.current.body,
123124
...options?.body,
124125
},
126+
streamMode,
125127
setCompletion: completion => mutate(completion, false),
126128
setLoading: mutateLoading,
127129
setError,
@@ -133,7 +135,7 @@
133135
mutateStreamData([...(streamData || []), ...(data || [])], false);
134136
},
135137
}),
136138
[
137139
mutate,
138140
mutateLoading,
139141
api,

‎packages/core/react/use-completion.ui.test.tsx

+114-66
Original file line numberDiff line numberDiff line change
@@ -8,87 +8,135 @@ import {
88
} from '../tests/utils/mock-fetch';
99
import { useCompletion } from './use-completion';
1010

11-
const TestComponent = () => {
12-
const {
13-
completion,
14-
handleSubmit,
15-
error,
16-
handleInputChange,
17-
input,
18-
isLoading,
19-
} = useCompletion();
20-
21-
return (
22-
<div>
23-
<div data-testid="loading">{isLoading.toString()}</div>
24-
<div data-testid="error">{error?.toString()}</div>
25-
<div data-testid="completion">{completion}</div>
26-
<form onSubmit={handleSubmit}>
27-
<input
28-
data-testid="input"
29-
value={input}
30-
placeholder="Say something..."
31-
onChange={handleInputChange}
32-
/>
33-
</form>
34-
</div>
35-
);
36-
};
37-
38-
beforeEach(() => {
39-
render(<TestComponent />);
40-
});
11+
describe('stream data stream', () => {
12+
const TestComponent = () => {
13+
const {
14+
completion,
15+
handleSubmit,
16+
error,
17+
handleInputChange,
18+
input,
19+
isLoading,
20+
} = useCompletion();
21+
22+
return (
23+
<div>
24+
<div data-testid="loading">{isLoading.toString()}</div>
25+
<div data-testid="error">{error?.toString()}</div>
26+
<div data-testid="completion">{completion}</div>
27+
<form onSubmit={handleSubmit}>
28+
<input
29+
data-testid="input"
30+
value={input}
31+
placeholder="Say something..."
32+
onChange={handleInputChange}
33+
/>
34+
</form>
35+
</div>
36+
);
37+
};
38+
39+
beforeEach(() => {
40+
render(<TestComponent />);
41+
});
4142

42-
afterEach(() => {
43-
vi.restoreAllMocks();
44-
cleanup();
45-
});
43+
afterEach(() => {
44+
vi.restoreAllMocks();
45+
cleanup();
46+
});
47+
48+
it('should render stream', async () => {
49+
mockFetchDataStream({
50+
url: 'https://example.com/api/completion',
51+
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
52+
});
53+
54+
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
4655

47-
it('should render complex text stream', async () => {
48-
mockFetchDataStream({
49-
url: 'https://example.com/api/completion',
50-
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
56+
await screen.findByTestId('completion');
57+
expect(screen.getByTestId('completion')).toHaveTextContent('Hello, world.');
5158
});
5259

53-
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
60+
describe('loading state', () => {
61+
it('should show loading state', async () => {
62+
let finishGeneration: ((value?: unknown) => void) | undefined;
63+
const finishGenerationPromise = new Promise(resolve => {
64+
finishGeneration = resolve;
65+
});
5466

55-
await screen.findByTestId('completion');
56-
expect(screen.getByTestId('completion')).toHaveTextContent('Hello, world.');
57-
});
67+
mockFetchDataStreamWithGenerator({
68+
url: 'https://example.com/api/chat',
69+
chunkGenerator: (async function* generate() {
70+
const encoder = new TextEncoder();
71+
yield encoder.encode('0:"Hello"\n');
72+
await finishGenerationPromise;
73+
})(),
74+
});
5875

59-
describe('loading state', () => {
60-
it('should show loading state', async () => {
61-
let finishGeneration: ((value?: unknown) => void) | undefined;
62-
const finishGenerationPromise = new Promise(resolve => {
63-
finishGeneration = resolve;
64-
});
76+
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
6577

66-
mockFetchDataStreamWithGenerator({
67-
url: 'https://example.com/api/chat',
68-
chunkGenerator: (async function* generate() {
69-
const encoder = new TextEncoder();
70-
yield encoder.encode('0:"Hello"\n');
71-
await finishGenerationPromise;
72-
})(),
78+
await screen.findByTestId('loading');
79+
expect(screen.getByTestId('loading')).toHaveTextContent('true');
80+
81+
finishGeneration?.();
82+
83+
await findByText(await screen.findByTestId('loading'), 'false');
84+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
7385
});
7486

75-
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
87+
it('should reset loading state on error', async () => {
88+
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
7689

77-
await screen.findByTestId('loading');
78-
expect(screen.getByTestId('loading')).toHaveTextContent('true');
90+
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
91+
92+
await screen.findByTestId('loading');
93+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
94+
});
95+
});
96+
});
7997

80-
finishGeneration?.();
98+
describe('text stream', () => {
99+
const TestComponent = () => {
100+
const { completion, handleSubmit, handleInputChange, input } =
101+
useCompletion({
102+
streamMode: 'text',
103+
});
104+
105+
return (
106+
<div>
107+
<div data-testid="completion-text-stream">{completion}</div>
108+
<form onSubmit={handleSubmit}>
109+
<input
110+
data-testid="input-text-stream"
111+
value={input}
112+
placeholder="Say something..."
113+
onChange={handleInputChange}
114+
/>
115+
</form>
116+
</div>
117+
);
118+
};
119+
120+
beforeEach(() => {
121+
render(<TestComponent />);
122+
});
81123

82-
await findByText(await screen.findByTestId('loading'), 'false');
83-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
124+
afterEach(() => {
125+
vi.restoreAllMocks();
126+
cleanup();
84127
});
85128

86-
it('should reset loading state on error', async () => {
87-
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
129+
it('should render stream', async () => {
130+
mockFetchDataStream({
131+
url: 'https://example.com/api/completion',
132+
chunks: ['Hello', ',', ' world', '.'],
133+
});
88134

89-
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
135+
await userEvent.type(screen.getByTestId('input-text-stream'), 'hi{enter}');
90136

91-
await screen.findByTestId('loading');
92-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
137+
await screen.findByTestId('completion-text-stream');
138+
expect(screen.getByTestId('completion-text-stream')).toHaveTextContent(
139+
'Hello, world.',
140+
);
93141
});
94142
});

‎packages/core/shared/call-chat-api.ts

+60-11
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,12 @@
11
import { parseComplexResponse } from './parse-complex-response';
22
import { IdGenerator, JSONValue, Message } from './types';
3+
import { createChunkDecoder } from './utils';
34

45
export async function callChatApi({
56
api,
67
messages,
78
body,
9+
streamMode = 'stream-data',
810
credentials,
911
headers,
1012
abortController,
@@ -17,6 +19,7 @@ export async function callChatApi({
1719
api: string;
1820
messages: Omit<Message, 'id'>[];
1921
body: Record<string, any>;
22+
streamMode?: 'stream-data' | 'text';
2023
credentials?: RequestCredentials;
2124
headers?: HeadersInit;
2225
abortController?: () => AbortController | null;
@@ -64,16 +67,62 @@ export async function callChatApi({
6467

6568
const reader = response.body.getReader();
6669

67-
return await parseComplexResponse({
68-
reader,
69-
abortControllerRef:
70-
abortController != null ? { current: abortController() } : undefined,
71-
update: onUpdate,
72-
onFinish(prefixMap) {
73-
if (onFinish && prefixMap.text != null) {
74-
onFinish(prefixMap.text);
70+
switch (streamMode) {
71+
case 'text': {
72+
const decoder = createChunkDecoder();
73+
74+
const resultMessage = {
75+
id: generateId(),
76+
createdAt: new Date(),
77+
role: 'assistant' as const,
78+
content: '',
79+
};
80+
81+
while (true) {
82+
const { done, value } = await reader.read();
83+
if (done) {
84+
break;
85+
}
86+
87+
resultMessage.content += decoder(value);
88+
resultMessage.id = generateId();
89+
90+
// note: creating a new message object is required for Solid.js streaming
91+
onUpdate([{ ...resultMessage }], []);
92+
93+
// The request has been aborted, stop reading the stream.
94+
if (abortController?.() === null) {
95+
reader.cancel();
96+
break;
97+
}
7598
}
76-
},
77-
generateId,
78-
});
99+
100+
onFinish?.(resultMessage);
101+
102+
return {
103+
messages: [resultMessage],
104+
data: [],
105+
};
106+
}
107+
108+
case 'stream-data': {
109+
return await parseComplexResponse({
110+
reader,
111+
abortControllerRef:
112+
abortController != null ? { current: abortController() } : undefined,
113+
update: onUpdate,
114+
onFinish(prefixMap) {
115+
if (onFinish && prefixMap.text != null) {
116+
onFinish(prefixMap.text);
117+
}
118+
},
119+
generateId,
120+
});
121+
}
122+
123+
default: {
124+
const exhaustiveCheck: never = streamMode;
125+
throw new Error(`Unknown stream mode: ${exhaustiveCheck}`);
126+
}
127+
}
79128
}

‎packages/core/shared/call-completion-api.ts

+46-10
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,14 @@
11
import { readDataStream } from './read-data-stream';
22
import { JSONValue } from './types';
3+
import { createChunkDecoder } from './utils';
34

45
export async function callCompletionApi({
56
api,
67
prompt,
78
credentials,
89
headers,
910
body,
11+
streamMode = 'stream-data',
1012
setCompletion,
1113
setLoading,
1214
setError,
@@ -21,6 +23,7 @@ export async function callCompletionApi({
2123
credentials?: RequestCredentials;
2224
headers?: HeadersInit;
2325
body: Record<string, any>;
26+
streamMode?: 'stream-data' | 'text';
2427
setCompletion: (completion: string) => void;
2528
setLoading: (loading: boolean) => void;
2629
setError: (error: Error | undefined) => void;
@@ -77,19 +80,52 @@ export async function callCompletionApi({
7780
let result = '';
7881
const reader = res.body.getReader();
7982

80-
for await (const { type, value } of readDataStream(reader, {
81-
isAborted: () => abortController === null,
82-
})) {
83-
switch (type) {
84-
case 'text': {
85-
result += value;
83+
switch (streamMode) {
84+
case 'text': {
85+
const decoder = createChunkDecoder();
86+
87+
while (true) {
88+
const { done, value } = await reader.read();
89+
if (done) {
90+
break;
91+
}
92+
93+
// Update the completion state with the new message tokens.
94+
result += decoder(value);
8695
setCompletion(result);
87-
break;
96+
97+
// The request has been aborted, stop reading the stream.
98+
if (abortController === null) {
99+
reader.cancel();
100+
break;
101+
}
88102
}
89-
case 'data': {
90-
onData?.(value);
91-
break;
103+
104+
break;
105+
}
106+
107+
case 'stream-data': {
108+
for await (const { type, value } of readDataStream(reader, {
109+
isAborted: () => abortController === null,
110+
})) {
111+
switch (type) {
112+
case 'text': {
113+
result += value;
114+
setCompletion(result);
115+
break;
116+
}
117+
case 'data': {
118+
onData?.(value);
119+
break;
120+
}
121+
}
92122
}
123+
break;
124+
}
125+
126+
default: {
127+
const exhaustiveCheck: never = streamMode;
128+
throw new Error(`Unknown stream mode: ${exhaustiveCheck}`);
93129
}
94130
}
95131

‎packages/core/shared/types.ts

+6
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,9 @@ export type UseChatOptions = {
248248
* handle the extra fields before forwarding the request to the AI service.
249249
*/
250250
sendExtraMessageFields?: boolean;
251+
252+
/** Stream mode (default to "stream-data") */
253+
streamMode?: 'stream-data' | 'text';
251254
};
252255

253256
export type UseCompletionOptions = {
@@ -313,6 +316,9 @@ export type UseCompletionOptions = {
313316
* ```
314317
*/
315318
body?: object;
319+
320+
/** Stream mode (default to "stream-data") */
321+
streamMode?: 'stream-data' | 'text';
316322
};
317323

318324
export type JSONValue =

‎packages/core/solid/use-chat.ts

+2
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ export function useChat({
8282
credentials,
8383
headers,
8484
body,
85+
streamMode,
8586
generateId = generateIdFunc,
8687
}: UseChatOptions = {}): UseChatHelpers {
8788
// Generate a unique ID for the chat if not provided.
@@ -158,6 +159,7 @@ export function useChat({
158159
...body,
159160
...options?.body,
160161
},
162+
streamMode,
161163
headers: {
162164
...headers,
163165
...options?.headers,

‎packages/core/solid/use-chat.ui.test.tsx

+145-87
Original file line numberDiff line numberDiff line change
@@ -10,117 +10,175 @@ import {
1010
} from '../tests/utils/mock-fetch';
1111
import { useChat } from './use-chat';
1212

13-
const TestComponent = () => {
14-
const { messages, append, error, data, isLoading } = useChat();
15-
16-
return (
17-
<div>
18-
<div data-testid="loading">{isLoading().toString()}</div>
19-
<div data-testid="error">{error()?.toString()}</div>
20-
<div data-testid="data">{JSON.stringify(data())}</div>
21-
22-
<For each={messages()}>
23-
{(m, idx) => (
24-
<div data-testid={`message-${idx()}`} class="whitespace-pre-wrap">
25-
{m.role === 'user' ? 'User: ' : 'AI: '}
26-
{m.content}
27-
</div>
28-
)}
29-
</For>
30-
31-
<button
32-
data-testid="button"
33-
onClick={() => {
34-
append({ role: 'user', content: 'hi' });
35-
}}
36-
/>
37-
</div>
38-
);
39-
};
40-
41-
beforeEach(() => {
42-
render(() => <TestComponent />);
43-
});
13+
describe('stream data stream', () => {
14+
const TestComponent = () => {
15+
const { messages, append, error, data, isLoading } = useChat();
16+
17+
return (
18+
<div>
19+
<div data-testid="loading">{isLoading().toString()}</div>
20+
<div data-testid="error">{error()?.toString()}</div>
21+
<div data-testid="data">{JSON.stringify(data())}</div>
22+
23+
<For each={messages()}>
24+
{(m, idx) => (
25+
<div data-testid={`message-${idx()}`}>
26+
{m.role === 'user' ? 'User: ' : 'AI: '}
27+
{m.content}
28+
</div>
29+
)}
30+
</For>
31+
32+
<button
33+
data-testid="button"
34+
onClick={() => {
35+
append({ role: 'user', content: 'hi' });
36+
}}
37+
/>
38+
</div>
39+
);
40+
};
41+
42+
beforeEach(() => {
43+
render(() => <TestComponent />);
44+
});
4445

45-
afterEach(() => {
46-
vi.restoreAllMocks();
47-
cleanup();
48-
});
46+
afterEach(() => {
47+
vi.restoreAllMocks();
48+
cleanup();
49+
});
50+
51+
it('should return messages', async () => {
52+
mockFetchDataStream({
53+
url: 'https://example.com/api/chat',
54+
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
55+
});
56+
57+
await userEvent.click(screen.getByTestId('button'));
4958

50-
it('should return messages', async () => {
51-
mockFetchDataStream({
52-
url: 'https://example.com/api/chat',
53-
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
59+
await screen.findByTestId('message-0');
60+
expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
61+
62+
await screen.findByTestId('message-1');
63+
expect(screen.getByTestId('message-1')).toHaveTextContent(
64+
'AI: Hello, world.',
65+
);
5466
});
5567

56-
await userEvent.click(screen.getByTestId('button'));
68+
it('should return messages and data', async () => {
69+
mockFetchDataStream({
70+
url: 'https://example.com/api/chat',
71+
chunks: ['2:[{"t1":"v1"}]\n', '0:"Hello"\n'],
72+
});
5773

58-
await screen.findByTestId('message-0');
59-
expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
74+
await userEvent.click(screen.getByTestId('button'));
6075

61-
await screen.findByTestId('message-1');
62-
expect(screen.getByTestId('message-1')).toHaveTextContent(
63-
'AI: Hello, world.',
64-
);
65-
});
76+
await screen.findByTestId('data');
77+
expect(screen.getByTestId('data')).toHaveTextContent('[{"t1":"v1"}]');
6678

67-
it('should return messages and data', async () => {
68-
mockFetchDataStream({
69-
url: 'https://example.com/api/chat',
70-
chunks: ['2:[{"t1":"v1"}]\n', '0:"Hello"\n'],
79+
await screen.findByTestId('message-1');
80+
expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello');
7181
});
7282

73-
await userEvent.click(screen.getByTestId('button'));
83+
it('should return error', async () => {
84+
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
7485

75-
await screen.findByTestId('data');
76-
expect(screen.getByTestId('data')).toHaveTextContent('[{"t1":"v1"}]');
86+
await userEvent.click(screen.getByTestId('button'));
7787

78-
await screen.findByTestId('message-1');
79-
expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello');
80-
});
88+
await screen.findByTestId('error');
89+
expect(screen.getByTestId('error')).toHaveTextContent('Error: Not found');
90+
});
8191

82-
it('should return error', async () => {
83-
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
92+
describe('loading state', () => {
93+
it('should show loading state', async () => {
94+
let finishGeneration: ((value?: unknown) => void) | undefined;
95+
const finishGenerationPromise = new Promise(resolve => {
96+
finishGeneration = resolve;
97+
});
8498

85-
await userEvent.click(screen.getByTestId('button'));
99+
mockFetchDataStreamWithGenerator({
100+
url: 'https://example.com/api/chat',
101+
chunkGenerator: (async function* generate() {
102+
const encoder = new TextEncoder();
103+
yield encoder.encode('0:"Hello"\n');
104+
await finishGenerationPromise;
105+
})(),
106+
});
86107

87-
await screen.findByTestId('error');
88-
expect(screen.getByTestId('error')).toHaveTextContent('Error: Not found');
89-
});
108+
await userEvent.click(screen.getByTestId('button'));
90109

91-
describe('loading state', () => {
92-
it('should show loading state', async () => {
93-
let finishGeneration: ((value?: unknown) => void) | undefined;
94-
const finishGenerationPromise = new Promise(resolve => {
95-
finishGeneration = resolve;
96-
});
110+
await screen.findByTestId('loading');
111+
expect(screen.getByTestId('loading')).toHaveTextContent('true');
97112

98-
mockFetchDataStreamWithGenerator({
99-
url: 'https://example.com/api/chat',
100-
chunkGenerator: (async function* generate() {
101-
const encoder = new TextEncoder();
102-
yield encoder.encode('0:"Hello"\n');
103-
await finishGenerationPromise;
104-
})(),
113+
finishGeneration?.();
114+
115+
await findByText(await screen.findByTestId('loading'), 'false');
116+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
105117
});
106118

107-
await userEvent.click(screen.getByTestId('button'));
119+
it('should reset loading state on error', async () => {
120+
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
108121

109-
await screen.findByTestId('loading');
110-
expect(screen.getByTestId('loading')).toHaveTextContent('true');
122+
await userEvent.click(screen.getByTestId('button'));
111123

112-
finishGeneration?.();
124+
await screen.findByTestId('loading');
125+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
126+
});
127+
});
128+
});
113129

114-
await findByText(await screen.findByTestId('loading'), 'false');
115-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
130+
describe('text stream', () => {
131+
const TestComponent = () => {
132+
const { messages, append } = useChat({
133+
streamMode: 'text',
134+
});
135+
136+
return (
137+
<div>
138+
<For each={messages()}>
139+
{(m, idx) => (
140+
<div data-testid={`message-${idx()}-text-stream`}>
141+
{m.role === 'user' ? 'User: ' : 'AI: '}
142+
{m.content}
143+
</div>
144+
)}
145+
</For>
146+
147+
<button
148+
data-testid="do-append-text-stream"
149+
onClick={() => {
150+
append({ role: 'user', content: 'hi' });
151+
}}
152+
/>
153+
</div>
154+
);
155+
};
156+
157+
beforeEach(() => {
158+
render(() => <TestComponent />);
116159
});
117160

118-
it('should reset loading state on error', async () => {
119-
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
161+
afterEach(() => {
162+
vi.restoreAllMocks();
163+
cleanup();
164+
});
120165

121-
await userEvent.click(screen.getByTestId('button'));
166+
it('should show streamed response', async () => {
167+
mockFetchDataStream({
168+
url: 'https://example.com/api/chat',
169+
chunks: ['Hello', ',', ' world', '.'],
170+
});
171+
172+
await userEvent.click(screen.getByTestId('do-append-text-stream'));
173+
174+
await screen.findByTestId('message-0-text-stream');
175+
expect(screen.getByTestId('message-0-text-stream')).toHaveTextContent(
176+
'User: hi',
177+
);
122178

123-
await screen.findByTestId('loading');
124-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
179+
await screen.findByTestId('message-1-text-stream');
180+
expect(screen.getByTestId('message-1-text-stream')).toHaveTextContent(
181+
'AI: Hello, world.',
182+
);
125183
});
126184
});

‎packages/core/solid/use-completion.ts

+2
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ export function useCompletion({
6767
credentials,
6868
headers,
6969
body,
70+
streamMode,
7071
onResponse,
7172
onFinish,
7273
onError,
@@ -115,6 +116,7 @@ export function useCompletion({
115116
...body,
116117
...options?.body,
117118
},
119+
streamMode,
118120
setCompletion: mutate,
119121
setLoading: setIsLoading,
120122
setError,

‎packages/core/solid/use-completion.ui.test.tsx

+103-59
Original file line numberDiff line numberDiff line change
@@ -9,80 +9,124 @@ import {
99
} from '../tests/utils/mock-fetch';
1010
import { useCompletion } from './use-completion';
1111

12-
const TestComponent = () => {
13-
const { completion, complete, error, isLoading } = useCompletion();
14-
15-
return (
16-
<div>
17-
<div data-testid="loading">{isLoading().toString()}</div>
18-
<div data-testid="error">{error()?.toString()}</div>
19-
20-
<div data-testid="completion">{completion()}</div>
21-
22-
<button
23-
data-testid="button"
24-
onClick={() => {
25-
complete('hi');
26-
}}
27-
/>
28-
</div>
29-
);
30-
};
31-
32-
beforeEach(() => {
33-
render(() => <TestComponent />);
34-
});
12+
describe('stream data stream', () => {
13+
const TestComponent = () => {
14+
const { completion, complete, error, isLoading } = useCompletion();
15+
16+
return (
17+
<div>
18+
<div data-testid="loading">{isLoading().toString()}</div>
19+
<div data-testid="error">{error()?.toString()}</div>
20+
21+
<div data-testid="completion">{completion()}</div>
22+
23+
<button
24+
data-testid="button"
25+
onClick={() => {
26+
complete('hi');
27+
}}
28+
/>
29+
</div>
30+
);
31+
};
32+
33+
beforeEach(() => {
34+
render(() => <TestComponent />);
35+
});
3536

36-
afterEach(() => {
37-
vi.restoreAllMocks();
38-
cleanup();
39-
});
37+
afterEach(() => {
38+
vi.restoreAllMocks();
39+
cleanup();
40+
});
41+
42+
it('should render complex text stream', async () => {
43+
mockFetchDataStream({
44+
url: 'https://example.com/api/completion',
45+
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
46+
});
47+
48+
await userEvent.click(screen.getByTestId('button'));
4049

41-
it('should render complex text stream', async () => {
42-
mockFetchDataStream({
43-
url: 'https://example.com/api/completion',
44-
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
50+
await screen.findByTestId('completion');
51+
expect(screen.getByTestId('completion')).toHaveTextContent('Hello, world.');
4552
});
4653

47-
await userEvent.click(screen.getByTestId('button'));
54+
describe('loading state', () => {
55+
it('should show loading state', async () => {
56+
let finishGeneration: ((value?: unknown) => void) | undefined;
57+
const finishGenerationPromise = new Promise(resolve => {
58+
finishGeneration = resolve;
59+
});
4860

49-
await screen.findByTestId('completion');
50-
expect(screen.getByTestId('completion')).toHaveTextContent('Hello, world.');
51-
});
61+
mockFetchDataStreamWithGenerator({
62+
url: 'https://example.com/api/chat',
63+
chunkGenerator: (async function* generate() {
64+
const encoder = new TextEncoder();
65+
yield encoder.encode('0:"Hello"\n');
66+
await finishGenerationPromise;
67+
})(),
68+
});
5269

53-
describe('loading state', () => {
54-
it('should show loading state', async () => {
55-
let finishGeneration: ((value?: unknown) => void) | undefined;
56-
const finishGenerationPromise = new Promise(resolve => {
57-
finishGeneration = resolve;
58-
});
70+
await userEvent.click(screen.getByTestId('button'));
5971

60-
mockFetchDataStreamWithGenerator({
61-
url: 'https://example.com/api/chat',
62-
chunkGenerator: (async function* generate() {
63-
const encoder = new TextEncoder();
64-
yield encoder.encode('0:"Hello"\n');
65-
await finishGenerationPromise;
66-
})(),
72+
await screen.findByTestId('loading');
73+
expect(screen.getByTestId('loading')).toHaveTextContent('true');
74+
75+
finishGeneration?.();
76+
77+
await findByText(await screen.findByTestId('loading'), 'false');
78+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
6779
});
6880

69-
await userEvent.click(screen.getByTestId('button'));
81+
it('should reset loading state on error', async () => {
82+
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
7083

71-
await screen.findByTestId('loading');
72-
expect(screen.getByTestId('loading')).toHaveTextContent('true');
84+
await userEvent.click(screen.getByTestId('button'));
85+
86+
await screen.findByTestId('loading');
87+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
88+
});
89+
});
90+
});
7391

74-
finishGeneration?.();
92+
describe('text stream', () => {
93+
const TestComponent = () => {
94+
const { completion, complete } = useCompletion({ streamMode: 'text' });
95+
96+
return (
97+
<div>
98+
<div data-testid="completion-text-stream">{completion()}</div>
99+
100+
<button
101+
data-testid="button-text-stream"
102+
onClick={() => {
103+
complete('hi');
104+
}}
105+
/>
106+
</div>
107+
);
108+
};
109+
110+
beforeEach(() => {
111+
render(() => <TestComponent />);
112+
});
75113

76-
await findByText(await screen.findByTestId('loading'), 'false');
77-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
114+
afterEach(() => {
115+
vi.restoreAllMocks();
116+
cleanup();
78117
});
79118

80-
it('should reset loading state on error', async () => {
81-
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
119+
it('should render stream', async () => {
120+
mockFetchDataStream({
121+
url: 'https://example.com/api/completion',
122+
chunks: ['Hello', ',', ' world', '.'],
123+
});
82124

83-
await userEvent.click(screen.getByTestId('button'));
125+
await userEvent.click(screen.getByTestId('button-text-stream'));
84126

85-
await screen.findByTestId('loading');
86-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
127+
await screen.findByTestId('completion-text-stream');
128+
expect(screen.getByTestId('completion-text-stream')).toHaveTextContent(
129+
'Hello, world.',
130+
);
87131
});
88132
});

‎packages/core/svelte/use-chat.ts

+4
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ const getStreamedResponse = async (
7272
previousMessages: Message[],
7373
abortControllerRef: AbortController | null,
7474
generateId: IdGenerator,
75+
streamMode?: 'stream-data' | 'text',
7576
onFinish?: (message: Message) => void,
7677
onResponse?: (response: Response) => void | Promise<void>,
7778
sendExtraMessageFields?: boolean,
@@ -116,6 +117,7 @@ const getStreamedResponse = async (
116117
tool_choice: chatRequest.tool_choice,
117118
}),
118119
},
120+
streamMode,
119121
credentials: extraMetadata.credentials,
120122
headers: {
121123
...extraMetadata.headers,
@@ -147,6 +149,7 @@ export function useChat({
147149
sendExtraMessageFields,
148150
experimental_onFunctionCall,
149151
experimental_onToolCall,
152+
streamMode,
150153
onResponse,
151154
onFinish,
152155
onError,
@@ -216,6 +219,7 @@ export function useChat({
216219
get(messages),
217220
abortController,
218221
generateId,
222+
streamMode,
219223
onFinish,
220224
onResponse,
221225
sendExtraMessageFields,

‎packages/core/svelte/use-completion.ts

+2
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ export function useCompletion({
6060
credentials,
6161
headers,
6262
body,
63+
streamMode,
6364
onResponse,
6465
onFinish,
6566
onError,
@@ -113,6 +114,7 @@ export function useCompletion({
113114
...body,
114115
...options?.body,
115116
},
117+
streamMode,
116118
setCompletion: mutate,
117119
setLoading: loadingState => loading.set(loadingState),
118120
setError: err => error.set(err),

‎packages/core/vue/TestChatTextStreamComponent.vue

+28
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
<script setup lang="ts">
2+
import { useChat } from './use-chat';
3+
4+
const { messages, append, data, error, isLoading } = useChat({
5+
streamMode: 'text',
6+
});
7+
</script>
8+
9+
<template>
10+
<div class="flex flex-col w-full max-w-md py-24 mx-auto stretch">
11+
<div data-testid="loading">{{ isLoading?.toString() }}</div>
12+
<div data-testid="error">{{ error?.toString() }}</div>
13+
<div data-testid="data">{{ JSON.stringify(data) }}</div>
14+
<div
15+
v-for="(m, idx) in messages"
16+
key="m.id"
17+
:data-testid="`message-${idx}`"
18+
>
19+
{{ m.role === 'user' ? 'User: ' : 'AI: ' }}
20+
{{ m.content }}
21+
</div>
22+
23+
<button
24+
data-testid="button"
25+
@click="append({ role: 'user', content: 'hi' })"
26+
/>
27+
</div>
28+
</template>

‎packages/core/vue/TestCompletionTextStreamComponent.vue

+19
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
<script setup lang="ts">
2+
import { useCompletion } from './use-completion';
3+
4+
const { completion, handleSubmit, input, isLoading, error } = useCompletion({
5+
streamMode: 'text',
6+
});
7+
</script>
8+
9+
<template>
10+
<div class="flex flex-col w-full max-w-md py-24 mx-auto stretch">
11+
<div data-testid="loading">{{ isLoading?.toString() }}</div>
12+
<div data-testid="error">{{ error?.toString() }}</div>
13+
<div data-testid="completion">{{ completion }}</div>
14+
15+
<form @submit="handleSubmit">
16+
<input data-testid="input" v-model="input" />
17+
</form>
18+
</div>
19+
</template>

‎packages/core/vue/use-chat.ts

+2
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,7 @@ export function useChat({
7070
initialInput = '',
7171
sendExtraMessageFields,
7272
experimental_onFunctionCall,
73+
streamMode,
7374
onResponse,
7475
onFinish,
7576
onError,
@@ -154,6 +155,7 @@ export function useChat({
154155
...unref(body), // Use unref to unwrap the ref value
155156
...options?.body,
156157
},
158+
streamMode,
157159
headers: {
158160
...headers,
159161
...options?.headers,

‎packages/core/vue/use-chat.ui.test.tsx

+92-61
Original file line numberDiff line numberDiff line change
@@ -7,95 +7,126 @@ import {
77
mockFetchError,
88
} from '../tests/utils/mock-fetch';
99
import TestChatComponent from './TestChatComponent.vue';
10+
import TestChatTextStreamComponent from './TestChatTextStreamComponent.vue';
1011

11-
beforeEach(() => {
12-
render(TestChatComponent);
13-
});
12+
describe('stream data stream', () => {
13+
beforeEach(() => {
14+
render(TestChatComponent);
15+
});
1416

15-
afterEach(() => {
16-
vi.restoreAllMocks();
17-
cleanup();
18-
});
17+
afterEach(() => {
18+
vi.restoreAllMocks();
19+
cleanup();
20+
});
21+
22+
it('should show streamed response', async () => {
23+
mockFetchDataStream({
24+
url: 'https://example.com/api/chat',
25+
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
26+
});
27+
28+
await userEvent.click(screen.getByTestId('button'));
29+
30+
await screen.findByTestId('message-0');
31+
expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
1932

20-
test('Shows streamed complex text response', async () => {
21-
mockFetchDataStream({
22-
url: 'https://example.com/api/chat',
23-
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
33+
await screen.findByTestId('message-1');
34+
expect(screen.getByTestId('message-1')).toHaveTextContent(
35+
'AI: Hello, world.',
36+
);
2437
});
2538

26-
await userEvent.click(screen.getByTestId('button'));
39+
it('should show streamed response with data', async () => {
40+
mockFetchDataStream({
41+
url: 'https://example.com/api/chat',
42+
chunks: ['2:[{"t1":"v1"}]\n', '0:"Hello"\n'],
43+
});
2744

28-
await screen.findByTestId('message-0');
29-
expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
45+
await userEvent.click(screen.getByTestId('button'));
3046

31-
await screen.findByTestId('message-1');
32-
expect(screen.getByTestId('message-1')).toHaveTextContent(
33-
'AI: Hello, world.',
34-
);
35-
});
47+
await screen.findByTestId('data');
48+
expect(screen.getByTestId('data')).toHaveTextContent('[{"t1":"v1"}]');
3649

37-
test('Shows streamed complex text response with data', async () => {
38-
mockFetchDataStream({
39-
url: 'https://example.com/api/chat',
40-
chunks: ['2:[{"t1":"v1"}]\n', '0:"Hello"\n'],
50+
await screen.findByTestId('message-1');
51+
expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello');
4152
});
4253

43-
await userEvent.click(screen.getByTestId('button'));
54+
it('should show error response', async () => {
55+
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
56+
57+
await userEvent.click(screen.getByTestId('button'));
4458

45-
await screen.findByTestId('data');
46-
expect(screen.getByTestId('data')).toHaveTextContent('[{"t1":"v1"}]');
59+
// TODO bug? the user message does not show up
60+
// await screen.findByTestId('message-0');
61+
// expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
4762

48-
await screen.findByTestId('message-1');
49-
expect(screen.getByTestId('message-1')).toHaveTextContent('AI: Hello');
50-
});
63+
await screen.findByTestId('error');
64+
expect(screen.getByTestId('error')).toHaveTextContent('Error: Not found');
65+
});
5166

52-
test('Shows error response', async () => {
53-
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
67+
describe('loading state', () => {
68+
it('should show loading state', async () => {
69+
let finishGeneration: ((value?: unknown) => void) | undefined;
70+
const finishGenerationPromise = new Promise(resolve => {
71+
finishGeneration = resolve;
72+
});
5473

55-
await userEvent.click(screen.getByTestId('button'));
74+
mockFetchDataStreamWithGenerator({
75+
url: 'https://example.com/api/chat',
76+
chunkGenerator: (async function* generate() {
77+
const encoder = new TextEncoder();
78+
yield encoder.encode('0:"Hello"\n');
79+
await finishGenerationPromise;
80+
})(),
81+
});
5682

57-
// TODO bug? the user message does not show up
58-
// await screen.findByTestId('message-0');
59-
// expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
83+
await userEvent.click(screen.getByTestId('button'));
6084

61-
await screen.findByTestId('error');
62-
expect(screen.getByTestId('error')).toHaveTextContent('Error: Not found');
63-
});
85+
await screen.findByTestId('loading');
86+
expect(screen.getByTestId('loading')).toHaveTextContent('true');
6487

65-
describe('loading state', () => {
66-
test('should show loading state', async () => {
67-
let finishGeneration: ((value?: unknown) => void) | undefined;
68-
const finishGenerationPromise = new Promise(resolve => {
69-
finishGeneration = resolve;
70-
});
88+
finishGeneration?.();
7189

72-
mockFetchDataStreamWithGenerator({
73-
url: 'https://example.com/api/chat',
74-
chunkGenerator: (async function* generate() {
75-
const encoder = new TextEncoder();
76-
yield encoder.encode('0:"Hello"\n');
77-
await finishGenerationPromise;
78-
})(),
90+
await findByText(await screen.findByTestId('loading'), 'false');
91+
92+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
7993
});
8094

81-
await userEvent.click(screen.getByTestId('button'));
95+
it('should reset loading state on error', async () => {
96+
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
8297

83-
await screen.findByTestId('loading');
84-
expect(screen.getByTestId('loading')).toHaveTextContent('true');
98+
await userEvent.click(screen.getByTestId('button'));
8599

86-
finishGeneration?.();
100+
await screen.findByTestId('loading');
101+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
102+
});
103+
});
104+
});
87105

88-
await findByText(await screen.findByTestId('loading'), 'false');
106+
describe('text stream', () => {
107+
beforeEach(() => {
108+
render(TestChatTextStreamComponent);
109+
});
89110

90-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
111+
afterEach(() => {
112+
vi.restoreAllMocks();
113+
cleanup();
91114
});
92115

93-
test('should reset loading state on error', async () => {
94-
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
116+
it('should show streamed response', async () => {
117+
mockFetchDataStream({
118+
url: 'https://example.com/api/chat',
119+
chunks: ['Hello', ',', ' world', '.'],
120+
});
95121

96122
await userEvent.click(screen.getByTestId('button'));
97123

98-
await screen.findByTestId('loading');
99-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
124+
await screen.findByTestId('message-0');
125+
expect(screen.getByTestId('message-0')).toHaveTextContent('User: hi');
126+
127+
await screen.findByTestId('message-1');
128+
expect(screen.getByTestId('message-1')).toHaveTextContent(
129+
'AI: Hello, world.',
130+
);
100131
});
101132
});

‎packages/core/vue/use-completion.ts

+2
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,7 @@ export function useCompletion({
6363
credentials,
6464
headers,
6565
body,
66+
streamMode,
6667
onResponse,
6768
onFinish,
6869
onError,
@@ -116,6 +117,7 @@ export function useCompletion({
116117
...unref(body),
117118
...options?.body,
118119
},
120+
streamMode,
119121
setCompletion: mutate,
120122
setLoading: loading => mutateLoading(() => loading),
121123
setError: err => {

‎packages/core/vue/use-completion.ui.test.ts

+64-38
Original file line numberDiff line numberDiff line change
@@ -7,61 +7,87 @@ import {
77
mockFetchError,
88
} from '../tests/utils/mock-fetch';
99
import TestCompletionComponent from './TestCompletionComponent.vue';
10+
import TestCompletionTextStreamComponent from './TestCompletionTextStreamComponent.vue';
1011

11-
beforeEach(() => {
12-
render(TestCompletionComponent);
13-
});
12+
describe('stream data stream', () => {
13+
beforeEach(() => {
14+
render(TestCompletionComponent);
15+
});
1416

15-
afterEach(() => {
16-
vi.restoreAllMocks();
17-
cleanup();
18-
});
17+
afterEach(() => {
18+
vi.restoreAllMocks();
19+
cleanup();
20+
});
21+
22+
it('should show streamed response', async () => {
23+
mockFetchDataStream({
24+
url: 'https://example.com/api/completion',
25+
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
26+
});
27+
28+
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
1929

20-
it('should render complex text stream', async () => {
21-
mockFetchDataStream({
22-
url: 'https://example.com/api/completion',
23-
chunks: ['0:"Hello"\n', '0:","\n', '0:" world"\n', '0:"."\n'],
30+
await screen.findByTestId('completion');
31+
expect(screen.getByTestId('completion')).toHaveTextContent('Hello, world.');
2432
});
2533

26-
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
34+
describe('loading state', () => {
35+
it('should show loading state', async () => {
36+
let finishGeneration: ((value?: unknown) => void) | undefined;
37+
const finishGenerationPromise = new Promise(resolve => {
38+
finishGeneration = resolve;
39+
});
2740

28-
await screen.findByTestId('completion');
29-
expect(screen.getByTestId('completion')).toHaveTextContent('Hello, world.');
30-
});
41+
mockFetchDataStreamWithGenerator({
42+
url: 'https://example.com/api/chat',
43+
chunkGenerator: (async function* generate() {
44+
const encoder = new TextEncoder();
45+
yield encoder.encode('0:"Hello"\n');
46+
await finishGenerationPromise;
47+
})(),
48+
});
3149

32-
describe('loading state', () => {
33-
it('should show loading state', async () => {
34-
let finishGeneration: ((value?: unknown) => void) | undefined;
35-
const finishGenerationPromise = new Promise(resolve => {
36-
finishGeneration = resolve;
37-
});
50+
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
3851

39-
mockFetchDataStreamWithGenerator({
40-
url: 'https://example.com/api/chat',
41-
chunkGenerator: (async function* generate() {
42-
const encoder = new TextEncoder();
43-
yield encoder.encode('0:"Hello"\n');
44-
await finishGenerationPromise;
45-
})(),
52+
await screen.findByTestId('loading');
53+
expect(screen.getByTestId('loading')).toHaveTextContent('true');
54+
55+
finishGeneration?.();
56+
57+
await findByText(await screen.findByTestId('loading'), 'false');
58+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
4659
});
4760

48-
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
61+
it('should reset loading state on error', async () => {
62+
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
4963

50-
await screen.findByTestId('loading');
51-
expect(screen.getByTestId('loading')).toHaveTextContent('true');
64+
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
5265

53-
finishGeneration?.();
66+
await screen.findByTestId('loading');
67+
expect(screen.getByTestId('loading')).toHaveTextContent('false');
68+
});
69+
});
70+
});
71+
72+
describe('stream data stream', () => {
73+
beforeEach(() => {
74+
render(TestCompletionTextStreamComponent);
75+
});
5476

55-
await findByText(await screen.findByTestId('loading'), 'false');
56-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
77+
afterEach(() => {
78+
vi.restoreAllMocks();
79+
cleanup();
5780
});
5881

59-
it('should reset loading state on error', async () => {
60-
mockFetchError({ statusCode: 404, errorMessage: 'Not found' });
82+
it('should show streamed response', async () => {
83+
mockFetchDataStream({
84+
url: 'https://example.com/api/completion',
85+
chunks: ['Hello', ',', ' world', '.'],
86+
});
6187

6288
await userEvent.type(screen.getByTestId('input'), 'hi{enter}');
6389

64-
await screen.findByTestId('loading');
65-
expect(screen.getByTestId('loading')).toHaveTextContent('false');
90+
await screen.findByTestId('completion');
91+
expect(screen.getByTestId('completion')).toHaveTextContent('Hello, world.');
6692
});
6793
});

‎pnpm-lock.yaml

+3
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.