fix: title generation + ai sdk upgrade (#1392)
This commit is contained in:
josh
GitHub
parent
f19d3d4071
commit
9d5d8a3ea7
38 changed files with 1422 additions and 2360 deletions
|
|
@ -1,2 +1 @@
|
|||
// biome-ignore lint/performance/noBarrelFile: "Required"
|
||||
export { GET, POST } from "@/app/(auth)/auth";
|
||||
|
|
|
|||
|
|
@ -16,7 +16,6 @@ declare module "next-auth" {
|
|||
} & DefaultSession["user"];
|
||||
}
|
||||
|
||||
// biome-ignore lint/nursery/useConsistentTypeDefinitions: "Required"
|
||||
interface User {
|
||||
id?: string;
|
||||
email?: string | null;
|
||||
|
|
|
|||
|
|
@ -22,13 +22,15 @@ export async function generateTitleFromUserMessage({
|
|||
}: {
|
||||
message: UIMessage;
|
||||
}) {
|
||||
const { text: title } = await generateText({
|
||||
const { text } = await generateText({
|
||||
model: getTitleModel(),
|
||||
system: titlePrompt,
|
||||
prompt: getTextFromMessage(message),
|
||||
});
|
||||
|
||||
return title;
|
||||
return text
|
||||
.replace(/^[#*"\s]+/, "")
|
||||
.replace(/["]+$/, "")
|
||||
.trim();
|
||||
}
|
||||
|
||||
export async function deleteTrailingMessages({ id }: { id: string }) {
|
||||
|
|
|
|||
|
|
@ -1,113 +1,3 @@
|
|||
import { createUIMessageStream, JsonToSseTransformStream } from "ai";
|
||||
import { differenceInSeconds } from "date-fns";
|
||||
import { auth } from "@/app/(auth)/auth";
|
||||
import {
|
||||
getChatById,
|
||||
getMessagesByChatId,
|
||||
getStreamIdsByChatId,
|
||||
} from "@/lib/db/queries";
|
||||
import type { Chat } from "@/lib/db/schema";
|
||||
import { ChatSDKError } from "@/lib/errors";
|
||||
import type { ChatMessage } from "@/lib/types";
|
||||
import { getStreamContext } from "../../route";
|
||||
|
||||
export async function GET(
|
||||
_: Request,
|
||||
{ params }: { params: Promise<{ id: string }> }
|
||||
) {
|
||||
const { id: chatId } = await params;
|
||||
|
||||
const streamContext = getStreamContext();
|
||||
const resumeRequestedAt = new Date();
|
||||
|
||||
if (!streamContext) {
|
||||
return new Response(null, { status: 204 });
|
||||
}
|
||||
|
||||
if (!chatId) {
|
||||
return new ChatSDKError("bad_request:api").toResponse();
|
||||
}
|
||||
|
||||
const session = await auth();
|
||||
|
||||
if (!session?.user) {
|
||||
return new ChatSDKError("unauthorized:chat").toResponse();
|
||||
}
|
||||
|
||||
let chat: Chat | null;
|
||||
|
||||
try {
|
||||
chat = await getChatById({ id: chatId });
|
||||
} catch {
|
||||
return new ChatSDKError("not_found:chat").toResponse();
|
||||
}
|
||||
|
||||
if (!chat) {
|
||||
return new ChatSDKError("not_found:chat").toResponse();
|
||||
}
|
||||
|
||||
if (chat.visibility === "private" && chat.userId !== session.user.id) {
|
||||
return new ChatSDKError("forbidden:chat").toResponse();
|
||||
}
|
||||
|
||||
const streamIds = await getStreamIdsByChatId({ chatId });
|
||||
|
||||
if (!streamIds.length) {
|
||||
return new ChatSDKError("not_found:stream").toResponse();
|
||||
}
|
||||
|
||||
const recentStreamId = streamIds.at(-1);
|
||||
|
||||
if (!recentStreamId) {
|
||||
return new ChatSDKError("not_found:stream").toResponse();
|
||||
}
|
||||
|
||||
const emptyDataStream = createUIMessageStream<ChatMessage>({
|
||||
// biome-ignore lint/suspicious/noEmptyBlockStatements: "Needs to exist"
|
||||
execute: () => {},
|
||||
});
|
||||
|
||||
const stream = await streamContext.resumableStream(recentStreamId, () =>
|
||||
emptyDataStream.pipeThrough(new JsonToSseTransformStream())
|
||||
);
|
||||
|
||||
/*
|
||||
* For when the generation is streaming during SSR
|
||||
* but the resumable stream has concluded at this point.
|
||||
*/
|
||||
if (!stream) {
|
||||
const messages = await getMessagesByChatId({ id: chatId });
|
||||
const mostRecentMessage = messages.at(-1);
|
||||
|
||||
if (!mostRecentMessage) {
|
||||
return new Response(emptyDataStream, { status: 200 });
|
||||
}
|
||||
|
||||
if (mostRecentMessage.role !== "assistant") {
|
||||
return new Response(emptyDataStream, { status: 200 });
|
||||
}
|
||||
|
||||
const messageCreatedAt = new Date(mostRecentMessage.createdAt);
|
||||
|
||||
if (differenceInSeconds(resumeRequestedAt, messageCreatedAt) > 15) {
|
||||
return new Response(emptyDataStream, { status: 200 });
|
||||
}
|
||||
|
||||
const restoredStream = createUIMessageStream<ChatMessage>({
|
||||
execute: ({ writer }) => {
|
||||
writer.write({
|
||||
type: "data-appendMessage",
|
||||
data: JSON.stringify(mostRecentMessage),
|
||||
transient: true,
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
return new Response(
|
||||
restoredStream.pipeThrough(new JsonToSseTransformStream()),
|
||||
{ status: 200 }
|
||||
);
|
||||
}
|
||||
|
||||
return new Response(stream, { status: 200 });
|
||||
export function GET() {
|
||||
return new Response(null, { status: 204 });
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,16 +2,13 @@ import { geolocation } from "@vercel/functions";
|
|||
import {
|
||||
convertToModelMessages,
|
||||
createUIMessageStream,
|
||||
JsonToSseTransformStream,
|
||||
smoothStream,
|
||||
createUIMessageStreamResponse,
|
||||
generateId,
|
||||
stepCountIs,
|
||||
streamText,
|
||||
} from "ai";
|
||||
import { after } from "next/server";
|
||||
import {
|
||||
createResumableStreamContext,
|
||||
type ResumableStreamContext,
|
||||
} from "resumable-stream";
|
||||
import { createResumableStreamContext } from "resumable-stream";
|
||||
import { auth, type UserType } from "@/app/(auth)/auth";
|
||||
import { entitlementsByUserType } from "@/lib/ai/entitlements";
|
||||
import { type RequestHints, systemPrompt } from "@/lib/ai/prompts";
|
||||
|
|
@ -41,28 +38,16 @@ import { type PostRequestBody, postRequestBodySchema } from "./schema";
|
|||
|
||||
export const maxDuration = 60;
|
||||
|
||||
let globalStreamContext: ResumableStreamContext | null = null;
|
||||
|
||||
export function getStreamContext() {
|
||||
if (!globalStreamContext) {
|
||||
try {
|
||||
globalStreamContext = createResumableStreamContext({
|
||||
waitUntil: after,
|
||||
});
|
||||
} catch (error: any) {
|
||||
if (error.message.includes("REDIS_URL")) {
|
||||
console.log(
|
||||
" > Resumable streams are disabled due to missing REDIS_URL"
|
||||
);
|
||||
} else {
|
||||
console.error(error);
|
||||
}
|
||||
}
|
||||
function getStreamContext() {
|
||||
try {
|
||||
return createResumableStreamContext({ waitUntil: after });
|
||||
} catch (_) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return globalStreamContext;
|
||||
}
|
||||
|
||||
export { getStreamContext };
|
||||
|
||||
export async function POST(request: Request) {
|
||||
let requestBody: PostRequestBody;
|
||||
|
||||
|
|
@ -94,7 +79,6 @@ export async function POST(request: Request) {
|
|||
return new ChatSDKError("rate_limit:chat").toResponse();
|
||||
}
|
||||
|
||||
// Check if this is a tool approval flow (all messages sent)
|
||||
const isToolApprovalFlow = Boolean(messages);
|
||||
|
||||
const chat = await getChatById({ id });
|
||||
|
|
@ -105,24 +89,19 @@ export async function POST(request: Request) {
|
|||
if (chat.userId !== session.user.id) {
|
||||
return new ChatSDKError("forbidden:chat").toResponse();
|
||||
}
|
||||
// Only fetch messages if chat already exists and not tool approval
|
||||
if (!isToolApprovalFlow) {
|
||||
messagesFromDb = await getMessagesByChatId({ id });
|
||||
}
|
||||
} else if (message?.role === "user") {
|
||||
// Save chat immediately with placeholder title
|
||||
await saveChat({
|
||||
id,
|
||||
userId: session.user.id,
|
||||
title: "New chat",
|
||||
visibility: selectedVisibilityType,
|
||||
});
|
||||
|
||||
// Start title generation in parallel (don't await)
|
||||
titlePromise = generateTitleFromUserMessage({ message });
|
||||
}
|
||||
|
||||
// Use all messages for tool approval, otherwise DB messages + new message
|
||||
const uiMessages = isToolApprovalFlow
|
||||
? (messages as ChatMessage[])
|
||||
: [...convertToUIMessages(messagesFromDb), message as ChatMessage];
|
||||
|
|
@ -136,7 +115,6 @@ export async function POST(request: Request) {
|
|||
country,
|
||||
};
|
||||
|
||||
// Only save user messages to the database (not tool approval responses)
|
||||
if (message?.role === "user") {
|
||||
await saveMessages({
|
||||
messages: [
|
||||
|
|
@ -152,29 +130,19 @@ export async function POST(request: Request) {
|
|||
});
|
||||
}
|
||||
|
||||
const streamId = generateUUID();
|
||||
await createStreamId({ streamId, chatId: id });
|
||||
const isReasoningModel =
|
||||
selectedChatModel.includes("reasoning") ||
|
||||
selectedChatModel.includes("thinking");
|
||||
|
||||
const modelMessages = await convertToModelMessages(uiMessages);
|
||||
|
||||
const stream = createUIMessageStream({
|
||||
// Pass original messages for tool approval continuation
|
||||
originalMessages: isToolApprovalFlow ? uiMessages : undefined,
|
||||
execute: async ({ writer: dataStream }) => {
|
||||
// Handle title generation in parallel
|
||||
if (titlePromise) {
|
||||
titlePromise.then((title) => {
|
||||
updateChatTitleById({ chatId: id, title });
|
||||
dataStream.write({ type: "data-chat-title", data: title });
|
||||
});
|
||||
}
|
||||
|
||||
const isReasoningModel =
|
||||
selectedChatModel.includes("reasoning") ||
|
||||
selectedChatModel.includes("thinking");
|
||||
|
||||
const result = streamText({
|
||||
model: getLanguageModel(selectedChatModel),
|
||||
system: systemPrompt({ selectedChatModel, requestHints }),
|
||||
messages: await convertToModelMessages(uiMessages),
|
||||
messages: modelMessages,
|
||||
stopWhen: stepCountIs(5),
|
||||
experimental_activeTools: isReasoningModel
|
||||
? []
|
||||
|
|
@ -184,9 +152,6 @@ export async function POST(request: Request) {
|
|||
"updateDocument",
|
||||
"requestSuggestions",
|
||||
],
|
||||
experimental_transform: isReasoningModel
|
||||
? undefined
|
||||
: smoothStream({ chunking: "word" }),
|
||||
providerOptions: isReasoningModel
|
||||
? {
|
||||
anthropic: {
|
||||
|
|
@ -198,10 +163,7 @@ export async function POST(request: Request) {
|
|||
getWeather,
|
||||
createDocument: createDocument({ session, dataStream }),
|
||||
updateDocument: updateDocument({ session, dataStream }),
|
||||
requestSuggestions: requestSuggestions({
|
||||
session,
|
||||
dataStream,
|
||||
}),
|
||||
requestSuggestions: requestSuggestions({ session, dataStream }),
|
||||
},
|
||||
experimental_telemetry: {
|
||||
isEnabled: isProductionEnvironment,
|
||||
|
|
@ -209,28 +171,25 @@ export async function POST(request: Request) {
|
|||
},
|
||||
});
|
||||
|
||||
result.consumeStream();
|
||||
dataStream.merge(result.toUIMessageStream({ sendReasoning: true }));
|
||||
|
||||
dataStream.merge(
|
||||
result.toUIMessageStream({
|
||||
sendReasoning: true,
|
||||
})
|
||||
);
|
||||
if (titlePromise) {
|
||||
const title = await titlePromise;
|
||||
dataStream.write({ type: "data-chat-title", data: title });
|
||||
updateChatTitleById({ chatId: id, title });
|
||||
}
|
||||
},
|
||||
generateId: generateUUID,
|
||||
onFinish: async ({ messages: finishedMessages }) => {
|
||||
if (isToolApprovalFlow) {
|
||||
// For tool approval, update existing messages (tool state changed) and save new ones
|
||||
for (const finishedMsg of finishedMessages) {
|
||||
const existingMsg = uiMessages.find((m) => m.id === finishedMsg.id);
|
||||
if (existingMsg) {
|
||||
// Update existing message with new parts (tool state changed)
|
||||
await updateMessage({
|
||||
id: finishedMsg.id,
|
||||
parts: finishedMsg.parts,
|
||||
});
|
||||
} else {
|
||||
// Save new message
|
||||
await saveMessages({
|
||||
messages: [
|
||||
{
|
||||
|
|
@ -246,7 +205,6 @@ export async function POST(request: Request) {
|
|||
}
|
||||
}
|
||||
} else if (finishedMessages.length > 0) {
|
||||
// Normal flow - save all finished messages
|
||||
await saveMessages({
|
||||
messages: finishedMessages.map((currentMessage) => ({
|
||||
id: currentMessage.id,
|
||||
|
|
@ -259,28 +217,30 @@ export async function POST(request: Request) {
|
|||
});
|
||||
}
|
||||
},
|
||||
onError: () => {
|
||||
return "Oops, an error occurred!";
|
||||
},
|
||||
onError: () => "Oops, an error occurred!",
|
||||
});
|
||||
|
||||
const streamContext = getStreamContext();
|
||||
|
||||
if (streamContext) {
|
||||
try {
|
||||
const resumableStream = await streamContext.resumableStream(
|
||||
streamId,
|
||||
() => stream.pipeThrough(new JsonToSseTransformStream())
|
||||
);
|
||||
if (resumableStream) {
|
||||
return new Response(resumableStream);
|
||||
return createUIMessageStreamResponse({
|
||||
stream,
|
||||
async consumeSseStream({ stream: sseStream }) {
|
||||
if (!process.env.REDIS_URL) {
|
||||
return;
|
||||
}
|
||||
} catch (error) {
|
||||
console.error("Failed to create resumable stream:", error);
|
||||
}
|
||||
}
|
||||
|
||||
return new Response(stream.pipeThrough(new JsonToSseTransformStream()));
|
||||
try {
|
||||
const streamContext = getStreamContext();
|
||||
if (streamContext) {
|
||||
const streamId = generateId();
|
||||
await createStreamId({ streamId, chatId: id });
|
||||
await streamContext.createNewResumableStream(
|
||||
streamId,
|
||||
() => sseStream
|
||||
);
|
||||
}
|
||||
} catch (_) {
|
||||
// ignore redis errors
|
||||
}
|
||||
},
|
||||
});
|
||||
} catch (error) {
|
||||
const vercelId = request.headers.get("x-vercel-id");
|
||||
|
||||
|
|
@ -288,7 +248,6 @@ export async function POST(request: Request) {
|
|||
return error.toResponse();
|
||||
}
|
||||
|
||||
// Check for Vercel AI Gateway credit card error
|
||||
if (
|
||||
error instanceof Error &&
|
||||
error.message?.includes(
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
@import "katex/dist/katex.min.css";
|
||||
|
||||
/* include utility classes in streamdown */
|
||||
@source '../node_modules/streamdown/dist/index.js';
|
||||
@source "../node_modules/streamdown/dist/index.js";
|
||||
|
||||
/* custom variant for setting dark mode programmatically */
|
||||
@custom-variant dark (&:is(.dark, .dark *));
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue