🎄 merry christmas: ai sdk v6 beta + tool approval (#1361)

This commit is contained in:
josh 2025-12-19 23:24:24 +00:00 committed by GitHub
parent 6e5b883cf2
commit 4d3ba8d9fe
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
21 changed files with 429 additions and 202 deletions

View file

@ -9,7 +9,7 @@ export const entitlementsByUserType: Record<UserType, Entitlements> = {
* For users without an account
*/
guest: {
maxMessagesPerDay: 10,
maxMessagesPerDay: 20,
},
/*

View file

@ -6,6 +6,11 @@ const mockResponses: Record<string, string> = {
greeting: "Hello! How can I help you today?",
};
const mockUsage = {
inputTokens: { total: 10, noCache: 10, cacheRead: 0, cacheWrite: 0 },
outputTokens: { total: 20, text: 20, reasoning: 0 },
};
function getResponseForPrompt(prompt: unknown): string {
const promptStr = JSON.stringify(prompt).toLowerCase();
@ -25,17 +30,14 @@ function getResponseForPrompt(prompt: unknown): string {
const createMockModel = (): LanguageModel => {
return {
specificationVersion: "v2",
specificationVersion: "v3",
provider: "mock",
modelId: "mock-model",
defaultObjectGenerationMode: "tool",
supportedUrls: [],
supportsImageUrls: false,
supportsStructuredOutputs: false,
supportedUrls: {},
doGenerate: async ({ prompt }: { prompt: unknown }) => ({
rawCall: { rawPrompt: null, rawSettings: {} },
finishReason: "stop",
usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
usage: mockUsage,
content: [{ type: "text", text: getResponseForPrompt(prompt) }],
warnings: [],
}),
@ -46,24 +48,26 @@ const createMockModel = (): LanguageModel => {
return {
stream: new ReadableStream({
async start(controller) {
controller.enqueue({ type: "text-start", id: "t1" });
for (const word of words) {
controller.enqueue({
type: "text-delta",
textDelta: `${word} `,
id: "t1",
delta: `${word} `,
});
await new Promise((resolve) => {
setTimeout(resolve, 10);
});
}
controller.enqueue({ type: "text-end", id: "t1" });
controller.enqueue({
type: "finish",
finishReason: "stop",
usage: { inputTokens: 10, outputTokens: 20 },
usage: mockUsage,
});
controller.close();
},
}),
rawCall: { rawPrompt: null, rawSettings: {} },
};
},
} as unknown as LanguageModel;
@ -71,17 +75,14 @@ const createMockModel = (): LanguageModel => {
const createMockReasoningModel = (): LanguageModel => {
return {
specificationVersion: "v2",
specificationVersion: "v3",
provider: "mock",
modelId: "mock-reasoning-model",
defaultObjectGenerationMode: "tool",
supportedUrls: [],
supportsImageUrls: false,
supportsStructuredOutputs: false,
supportedUrls: {},
doGenerate: async () => ({
rawCall: { rawPrompt: null, rawSettings: {} },
finishReason: "stop",
usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
usage: mockUsage,
content: [{ type: "text", text: "This is a reasoned response." }],
reasoning: [
{ type: "text", text: "Let me think through this step by step..." },
@ -91,62 +92,77 @@ const createMockReasoningModel = (): LanguageModel => {
doStream: () => ({
stream: new ReadableStream({
async start(controller) {
controller.enqueue({ type: "reasoning-start", id: "r1" });
controller.enqueue({
type: "reasoning",
textDelta: "Let me think through this step by step... ",
type: "reasoning-delta",
id: "r1",
delta: "Let me think through this step by step... ",
});
controller.enqueue({ type: "reasoning-end", id: "r1" });
await new Promise((resolve) => {
setTimeout(resolve, 10);
});
controller.enqueue({ type: "text-start", id: "t1" });
controller.enqueue({
type: "text-delta",
textDelta: "This is a reasoned response.",
id: "t1",
delta: "This is a reasoned response.",
});
controller.enqueue({ type: "text-end", id: "t1" });
controller.enqueue({
type: "finish",
finishReason: "stop",
usage: { inputTokens: 10, outputTokens: 20 },
usage: mockUsage,
});
controller.close();
},
}),
rawCall: { rawPrompt: null, rawSettings: {} },
}),
} as unknown as LanguageModel;
};
const createMockTitleModel = (): LanguageModel => {
return {
specificationVersion: "v2",
specificationVersion: "v3",
provider: "mock",
modelId: "mock-title-model",
defaultObjectGenerationMode: "tool",
supportedUrls: [],
supportsImageUrls: false,
supportsStructuredOutputs: false,
supportedUrls: {},
doGenerate: async () => ({
rawCall: { rawPrompt: null, rawSettings: {} },
finishReason: "stop",
usage: { inputTokens: 5, outputTokens: 5, totalTokens: 10 },
usage: {
inputTokens: { total: 5, noCache: 5, cacheRead: 0, cacheWrite: 0 },
outputTokens: { total: 5, text: 5, reasoning: 0 },
},
content: [{ type: "text", text: "Test Conversation" }],
warnings: [],
}),
doStream: () => ({
stream: new ReadableStream({
start(controller) {
controller.enqueue({ type: "text-start", id: "t1" });
controller.enqueue({
type: "text-delta",
textDelta: "Test Conversation",
id: "t1",
delta: "Test Conversation",
});
controller.enqueue({ type: "text-end", id: "t1" });
controller.enqueue({
type: "finish",
finishReason: "stop",
usage: { inputTokens: 5, outputTokens: 5 },
usage: {
inputTokens: {
total: 5,
noCache: 5,
cacheRead: 0,
cacheWrite: 0,
},
outputTokens: { total: 5, text: 5, reasoning: 0 },
},
});
controller.close();
},
}),
rawCall: { rawPrompt: null, rawSettings: {} },
}),
} as unknown as LanguageModel;
};

View file

@ -1,12 +1,16 @@
import { simulateReadableStream } from "ai";
import { MockLanguageModelV2 } from "ai/test";
import { MockLanguageModelV3 } from "ai/test";
import { getResponseChunksByPrompt } from "@/tests/prompts/utils";
export const chatModel = new MockLanguageModelV2({
const mockUsage = {
inputTokens: { total: 10, noCache: 10, cacheRead: 0, cacheWrite: 0 },
outputTokens: { total: 20, text: 20, reasoning: 0 },
};
export const chatModel = new MockLanguageModelV3({
doGenerate: async () => ({
rawCall: { rawPrompt: null, rawSettings: {} },
finishReason: "stop",
usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
usage: mockUsage,
content: [{ type: "text", text: "Hello, world!" }],
warnings: [],
}),
@ -16,15 +20,13 @@ export const chatModel = new MockLanguageModelV2({
initialDelayInMs: 1000,
chunks: getResponseChunksByPrompt(prompt),
}),
rawCall: { rawPrompt: null, rawSettings: {} },
}),
});
export const reasoningModel = new MockLanguageModelV2({
export const reasoningModel = new MockLanguageModelV3({
doGenerate: async () => ({
rawCall: { rawPrompt: null, rawSettings: {} },
finishReason: "stop",
usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
usage: mockUsage,
content: [{ type: "text", text: "Hello, world!" }],
warnings: [],
}),
@ -34,15 +36,13 @@ export const reasoningModel = new MockLanguageModelV2({
initialDelayInMs: 1000,
chunks: getResponseChunksByPrompt(prompt, true),
}),
rawCall: { rawPrompt: null, rawSettings: {} },
}),
});
export const titleModel = new MockLanguageModelV2({
export const titleModel = new MockLanguageModelV3({
doGenerate: async () => ({
rawCall: { rawPrompt: null, rawSettings: {} },
finishReason: "stop",
usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
usage: mockUsage,
content: [{ type: "text", text: "This is a test title" }],
warnings: [],
}),
@ -57,19 +57,17 @@ export const titleModel = new MockLanguageModelV2({
{
type: "finish",
finishReason: "stop",
usage: { inputTokens: 3, outputTokens: 10, totalTokens: 13 },
usage: mockUsage,
},
],
}),
rawCall: { rawPrompt: null, rawSettings: {} },
}),
});
export const artifactModel = new MockLanguageModelV2({
export const artifactModel = new MockLanguageModelV3({
doGenerate: async () => ({
rawCall: { rawPrompt: null, rawSettings: {} },
finishReason: "stop",
usage: { inputTokens: 10, outputTokens: 20, totalTokens: 30 },
usage: mockUsage,
content: [{ type: "text", text: "Hello, world!" }],
warnings: [],
}),
@ -79,6 +77,5 @@ export const artifactModel = new MockLanguageModelV2({
initialDelayInMs: 100,
chunks: getResponseChunksByPrompt(prompt),
}),
rawCall: { rawPrompt: null, rawSettings: {} },
}),
});

View file

@ -40,6 +40,7 @@ export const getWeather = tool({
.describe("City name (e.g., 'San Francisco', 'New York', 'London')")
.optional(),
}),
needsApproval: true,
execute: async (input) => {
let latitude: number;
let longitude: number;

View file

@ -9,7 +9,8 @@ config({
const runMigrate = async () => {
if (!process.env.POSTGRES_URL) {
throw new Error("POSTGRES_URL is not defined");
console.log("⏭️ POSTGRES_URL not defined, skipping migrations");
process.exit(0);
}
const connection = postgres(process.env.POSTGRES_URL, { max: 1 });

View file

@ -250,6 +250,20 @@ export async function saveMessages({ messages }: { messages: DBMessage[] }) {
}
}
export async function updateMessage({
id,
parts,
}: {
id: string;
parts: DBMessage["parts"];
}) {
try {
return await db.update(message).set({ parts }).where(eq(message.id, id));
} catch (_error) {
throw new ChatSDKError("bad_request:database", "Failed to update message");
}
}
export async function getMessagesByChatId({ id }: { id: string }) {
try {
return await db

View file

@ -1,6 +1,6 @@
import type {
CoreAssistantMessage,
CoreToolMessage,
AssistantModelMessage,
ToolModelMessage,
UIMessage,
UIMessagePart,
} from 'ai';
@ -63,7 +63,7 @@ export function generateUUID(): string {
});
}
type ResponseMessageWithoutId = CoreToolMessage | CoreAssistantMessage;
type ResponseMessageWithoutId = ToolModelMessage | AssistantModelMessage;
type ResponseMessage = ResponseMessageWithoutId & { id: string };
export function getMostRecentUserMessage(messages: UIMessage[]) {