105 lines
3.0 KiB
TypeScript
105 lines
3.0 KiB
TypeScript
import { openai } from '@ai-sdk/openai';
|
|
import {
|
|
convertToModelMessages,
|
|
InferUITools,
|
|
stepCountIs,
|
|
streamText,
|
|
tool,
|
|
UIDataTypes,
|
|
UIMessage,
|
|
} from 'ai';
|
|
import { convertArrayToReadableStream, MockLanguageModelV2 } from 'ai/test';
|
|
import { z } from 'zod';
|
|
|
|
// Allow streaming responses up to 30 seconds
|
|
export const maxDuration = 30;
|
|
|
|
const getWeatherInformationTool = tool({
|
|
description: 'show the weather in a given city to the user',
|
|
inputSchema: z.object({ city: z.string() }),
|
|
execute: async ({ city }: { city: string }) => {
|
|
// Add artificial delay of 5 seconds
|
|
await new Promise(resolve => setTimeout(resolve, 5000));
|
|
|
|
const weatherOptions = ['sunny', 'cloudy', 'rainy', 'snowy', 'windy'];
|
|
return weatherOptions[Math.floor(Math.random() * weatherOptions.length)];
|
|
},
|
|
});
|
|
|
|
const tools = {
|
|
// server-side tool with execute function:
|
|
getWeatherInformation: getWeatherInformationTool,
|
|
} as const;
|
|
|
|
export type UseChatToolsMessage = UIMessage<
|
|
never,
|
|
UIDataTypes,
|
|
InferUITools<typeof tools>
|
|
>;
|
|
|
|
export async function POST(req: Request) {
|
|
const { messages } = await req.json();
|
|
|
|
console.log('messages', JSON.stringify(messages, null, 2));
|
|
|
|
const result = streamText({
|
|
model: openai('gpt-4o'),
|
|
messages: convertToModelMessages(messages),
|
|
stopWhen: stepCountIs(5), // multi-steps for server-side tools
|
|
tools,
|
|
prepareStep: async ({ stepNumber }) => {
|
|
// inject invalid tool call in first step:
|
|
if (stepNumber === 0) {
|
|
return {
|
|
model: new MockLanguageModelV2({
|
|
doStream: async () => ({
|
|
stream: convertArrayToReadableStream([
|
|
{ type: 'stream-start', warnings: [] },
|
|
{
|
|
type: 'tool-input-start',
|
|
id: 'call-1',
|
|
toolName: 'getWeatherInformation',
|
|
providerExecuted: true,
|
|
},
|
|
{
|
|
type: 'tool-input-delta',
|
|
id: 'call-1',
|
|
delta: `{ "cities": "San Francisco" }`,
|
|
},
|
|
{
|
|
type: 'tool-input-end',
|
|
id: 'call-1',
|
|
},
|
|
{
|
|
type: 'tool-call',
|
|
toolCallType: 'function',
|
|
toolCallId: 'call-1',
|
|
toolName: 'getWeatherInformation',
|
|
// wrong tool call arguments (city vs cities):
|
|
input: `{ "cities": "San Francisco" }`,
|
|
},
|
|
{
|
|
type: 'finish',
|
|
finishReason: 'stop',
|
|
usage: {
|
|
inputTokens: 10,
|
|
outputTokens: 20,
|
|
totalTokens: 30,
|
|
},
|
|
},
|
|
]),
|
|
}),
|
|
}),
|
|
};
|
|
}
|
|
},
|
|
});
|
|
|
|
return result.toUIMessageStreamResponse({
|
|
// originalMessages: messages, //add if you want to have correct ids
|
|
onFinish: options => {
|
|
console.log('onFinish', options);
|
|
},
|
|
});
|
|
}
|