Skip to content

Commit 744d73d

Browse files
jeremyphilemonjeremyphilemon
committed
bump to beta
1 parent b0e1d96 commit 744d73d

30 files changed

+580
-631
lines changed

‎app/(chat)/api/chat/[id]/stream/route.ts

Lines changed: 109 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,109 @@
1+
import { auth } from '@/app/(auth)/auth';
2+
import {
3+
getChatById,
4+
getMessagesByChatId,
5+
getStreamIdsByChatId,
6+
} from '@/lib/db/queries';
7+
import type { Chat } from '@/lib/db/schema';
8+
import { ChatSDKError } from '@/lib/errors';
9+
import { createUIMessageStream, JsonToSseTransformStream } from 'ai';
10+
import { differenceInSeconds } from 'date-fns';
11+
import { after } from 'next/server';
12+
import { createResumableStreamContext } from 'resumable-stream';
13+
14+
export const maxDuration = 60;
15+
16+
export async function GET(
17+
_: Request,
18+
{ params }: { params: Promise<{ id: string }> },
19+
) {
20+
const { id } = await params;
21+
22+
if (!id) {
23+
return new ChatSDKError('bad_request:api').toResponse();
24+
}
25+
26+
const resumeRequestedAt = new Date();
27+
28+
const session = await auth();
29+
30+
if (!session?.user) {
31+
return new ChatSDKError('unauthorized:chat').toResponse();
32+
}
33+
34+
let chat: Chat;
35+
36+
try {
37+
chat = await getChatById({ id });
38+
} catch {
39+
return new ChatSDKError('not_found:chat').toResponse();
40+
}
41+
42+
if (!chat) {
43+
return new ChatSDKError('not_found:chat').toResponse();
44+
}
45+
46+
if (chat.visibility === 'private' && chat.userId !== session.user.id) {
47+
return new ChatSDKError('forbidden:chat').toResponse();
48+
}
49+
50+
const streamIds = await getStreamIdsByChatId({ chatId: id });
51+
52+
if (!streamIds.length) {
53+
return new ChatSDKError('not_found:stream').toResponse();
54+
}
55+
56+
const recentStreamId = streamIds.at(-1);
57+
58+
if (!recentStreamId) {
59+
return new ChatSDKError('not_found:stream').toResponse();
60+
}
61+
62+
const emptyDataStream = createUIMessageStream({
63+
execute: () => {},
64+
});
65+
66+
const streamContext = createResumableStreamContext({
67+
waitUntil: after,
68+
});
69+
70+
/*
71+
* For when the generation is streaming during SSR
72+
* but the resumable stream has concluded at this point.
73+
*/
74+
if (!streamContext) {
75+
const messages = await getMessagesByChatId({ id: id });
76+
const mostRecentMessage = messages.at(-1);
77+
78+
if (!mostRecentMessage) {
79+
return new Response(emptyDataStream, { status: 200 });
80+
}
81+
82+
if (mostRecentMessage.role !== 'assistant') {
83+
return new Response(emptyDataStream, { status: 200 });
84+
}
85+
86+
const messageCreatedAt = new Date(mostRecentMessage.createdAt);
87+
88+
if (differenceInSeconds(resumeRequestedAt, messageCreatedAt) > 15) {
89+
return new Response(emptyDataStream, { status: 200 });
90+
}
91+
92+
const restoredStream = createUIMessageStream({
93+
execute: ({ writer }) => {
94+
writer.write({
95+
type: 'data-append-in-flight-message',
96+
data: mostRecentMessage,
97+
});
98+
},
99+
});
100+
101+
return new Response(restoredStream, { status: 200 });
102+
}
103+
104+
return new Response(
105+
await streamContext.resumableStream(recentStreamId, () =>
106+
emptyDataStream.pipeThrough(new JsonToSseTransformStream()),
107+
),
108+
);
109+
}

‎app/(chat)/api/chat/route.ts

Lines changed: 23 additions & 142 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import {
22
convertToModelMessages,
33
createUIMessageStream,
4-
createUIMessageStreamResponse,
4+
JsonToSseTransformStream,
55
smoothStream,
66
stepCountIs,
77
streamText,
@@ -15,7 +15,6 @@ import {
1515
getChatById,
1616
getMessageCountByUserId,
1717
getMessagesByChatId,
18-
getStreamIdsByChatId,
1918
saveChat,
2019
saveMessages,
2120
} from '@/lib/db/queries';
@@ -24,45 +23,19 @@ import { isProductionEnvironment } from '@/lib/constants';
2423
import { myProvider } from '@/lib/ai/providers';
2524
import { postRequestBodySchema, type PostRequestBody } from './schema';
2625
import { geolocation } from '@vercel/functions';
27-
import {
28-
createResumableStreamContext,
29-
type ResumableStreamContext,
30-
} from 'resumable-stream';
31-
import { after } from 'next/server';
32-
import type { Chat } from '@/lib/db/schema';
33-
import { differenceInSeconds } from 'date-fns';
3426
import { ChatSDKError } from '@/lib/errors';
3527
import { createDocument } from '@/lib/ai/tools/create-document';
3628
import { updateDocument } from '@/lib/ai/tools/update-document';
3729
import { requestSuggestions } from '@/lib/ai/tools/request-suggestions';
3830
import { generateTitleFromUserMessage } from '../../actions';
39-
import { convertToStringStream, generateUUID } from '@/lib/utils';
31+
import { generateUUID } from '@/lib/utils';
4032
import { entitlementsByUserType } from '@/lib/ai/entitlements';
33+
import type { ChatMessage } from '@/lib/types';
34+
import { createResumableStreamContext } from 'resumable-stream';
35+
import { after } from 'next/server';
4136

4237
export const maxDuration = 60;
4338

44-
let globalStreamContext: ResumableStreamContext | null = null;
45-
46-
function getStreamContext() {
47-
if (!globalStreamContext) {
48-
try {
49-
globalStreamContext = createResumableStreamContext({
50-
waitUntil: after,
51-
});
52-
} catch (error: any) {
53-
if (error.message.includes('REDIS_URL')) {
54-
console.log(
55-
' > Resumable streams are disabled due to missing REDIS_URL',
56-
);
57-
} else {
58-
console.error(error);
59-
}
60-
}
61-
}
62-
63-
return globalStreamContext;
64-
}
65-
6639
export async function POST(request: Request) {
6740
let requestBody: PostRequestBody;
6841

@@ -74,6 +47,10 @@ export async function POST(request: Request) {
7447
}
7548

7649
try {
50+
const streamContext = createResumableStreamContext({
51+
waitUntil: after,
52+
});
53+
7754
const { id, message, selectedChatModel, selectedVisibilityType } =
7855
requestBody;
7956

@@ -164,17 +141,23 @@ export async function POST(request: Request) {
164141
isEnabled: isProductionEnvironment,
165142
functionId: 'stream-text',
166143
},
144+
_internal: {
145+
generateId: generateUUID,
146+
},
167147
});
168148

149+
result.consumeStream();
150+
169151
streamWriter.merge(
170-
result.toUIMessageStream({
152+
result.toUIMessageStream<ChatMessage>({
171153
sendReasoning: true,
172-
newMessageId: generateUUID(),
173154
onFinish: async ({ responseMessage }) => {
174155
await saveMessages({
175156
messages: [
176157
{
177-
...responseMessage,
158+
id: responseMessage.id,
159+
role: 'assistant',
160+
parts: responseMessage.parts,
178161
createdAt: new Date(),
179162
attachments: [],
180163
chatId: id,
@@ -184,126 +167,24 @@ export async function POST(request: Request) {
184167
},
185168
}),
186169
);
187-
188-
result.consumeStream();
189170
},
190171
onError: () => {
191172
return 'Oops! Something went wrong, please try again later.';
192173
},
193174
});
194175

195-
const streamContext = getStreamContext();
196-
197-
if (streamContext) {
198-
return new Response(
199-
await streamContext.resumableStream(streamId, () =>
200-
convertToStringStream(stream),
201-
),
202-
);
203-
} else {
204-
return createUIMessageStreamResponse({ stream });
205-
}
176+
return new Response(
177+
await streamContext.resumableStream(streamId, () =>
178+
stream.pipeThrough(new JsonToSseTransformStream()),
179+
),
180+
);
206181
} catch (error) {
207182
if (error instanceof ChatSDKError) {
208183
return error.toResponse();
209184
}
210185
}
211186
}
212187

213-
export async function GET(request: Request) {
214-
const streamContext = getStreamContext();
215-
const resumeRequestedAt = new Date();
216-
217-
if (!streamContext) {
218-
return new Response(null, { status: 204 });
219-
}
220-
221-
const { searchParams } = new URL(request.url);
222-
const chatId = searchParams.get('chatId');
223-
224-
if (!chatId) {
225-
return new ChatSDKError('bad_request:api').toResponse();
226-
}
227-
228-
const session = await auth();
229-
230-
if (!session?.user) {
231-
return new ChatSDKError('unauthorized:chat').toResponse();
232-
}
233-
234-
let chat: Chat;
235-
236-
try {
237-
chat = await getChatById({ id: chatId });
238-
} catch {
239-
return new ChatSDKError('not_found:chat').toResponse();
240-
}
241-
242-
if (!chat) {
243-
return new ChatSDKError('not_found:chat').toResponse();
244-
}
245-
246-
if (chat.visibility === 'private' && chat.userId !== session.user.id) {
247-
return new ChatSDKError('forbidden:chat').toResponse();
248-
}
249-
250-
const streamIds = await getStreamIdsByChatId({ chatId });
251-
252-
if (!streamIds.length) {
253-
return new ChatSDKError('not_found:stream').toResponse();
254-
}
255-
256-
const recentStreamId = streamIds.at(-1);
257-
258-
if (!recentStreamId) {
259-
return new ChatSDKError('not_found:stream').toResponse();
260-
}
261-
262-
const emptyDataStream = createUIMessageStream({
263-
execute: () => {},
264-
});
265-
266-
const stream = await streamContext.resumableStream(recentStreamId, () =>
267-
convertToStringStream(emptyDataStream),
268-
);
269-
270-
/*
271-
* For when the generation is streaming during SSR
272-
* but the resumable stream has concluded at this point.
273-
*/
274-
if (!stream) {
275-
const messages = await getMessagesByChatId({ id: chatId });
276-
const mostRecentMessage = messages.at(-1);
277-
278-
if (!mostRecentMessage) {
279-
return new Response(emptyDataStream, { status: 200 });
280-
}
281-
282-
if (mostRecentMessage.role !== 'assistant') {
283-
return new Response(emptyDataStream, { status: 200 });
284-
}
285-
286-
const messageCreatedAt = new Date(mostRecentMessage.createdAt);
287-
288-
if (differenceInSeconds(resumeRequestedAt, messageCreatedAt) > 15) {
289-
return new Response(emptyDataStream, { status: 200 });
290-
}
291-
292-
const restoredStream = createUIMessageStream({
293-
execute: ({ writer }) => {
294-
writer.write({
295-
type: 'data-append-in-flight-message',
296-
data: mostRecentMessage,
297-
});
298-
},
299-
});
300-
301-
return new Response(restoredStream, { status: 200 });
302-
}
303-
304-
return new Response(stream, { status: 200 });
305-
}
306-
307188
export async function DELETE(request: Request) {
308189
const { searchParams } = new URL(request.url);
309190
const id = searchParams.get('id');

0 commit comments

Comments
 (0)