ai-functions
Version:
A powerful TypeScript library for building AI-powered applications with template literals and structured outputs
176 lines • 6.86 kB
JavaScript
import { describe, expect, it, beforeEach } from 'vitest';
import { openai } from '@ai-sdk/openai';
import { streamObject, streamText } from 'ai';
import { z } from 'zod';
describe('AI SDK Examples', () => {
beforeEach(() => {
if (!process.env.OPENAI_API_KEY) {
throw new Error('OPENAI_API_KEY environment variable is required');
}
});
const model = openai('gpt-4o-mini');
describe('Streaming Text', () => {
it('should stream text using streamText', async () => {
const chunks = [];
const { textStream } = streamText({
model,
prompt: 'Write a short story about a robot learning to paint',
onChunk: ({ chunk }) => {
if (chunk.type === 'text-delta' && chunk.text) {
chunks.push(chunk.text);
}
}
});
let fullText = '';
for await (const text of textStream) {
fullText += text;
}
expect(chunks.length).toBeGreaterThan(0);
expect(fullText).toBeDefined();
expect(fullText.length).toBeGreaterThan(0);
}, 30000);
});
describe('Streaming Objects', () => {
it('should stream array elements using streamObject', async () => {
const { elementStream } = streamObject({
model,
output: 'array',
schema: z.object({
name: z.string(),
type: z.string().describe('Type of fruit'),
color: z.string(),
taste: z.string().describe('Description of taste')
}),
prompt: 'Generate descriptions of 3 different fruits'
});
const fruits = [];
for await (const fruit of elementStream) {
fruits.push(fruit);
}
expect(fruits.length).toBe(3);
fruits.forEach(fruit => {
expect(fruit).toHaveProperty('name');
expect(fruit).toHaveProperty('type');
expect(fruit).toHaveProperty('color');
expect(fruit).toHaveProperty('taste');
});
}, 30000);
it('should stream partial objects using streamObject', async () => {
const { partialObjectStream } = streamObject({
model,
schema: z.object({
recipe: z.object({
name: z.string(),
ingredients: z.array(z.object({
name: z.string(),
amount: z.string()
})),
steps: z.array(z.string())
})
}),
prompt: 'Generate a recipe for chocolate chip cookies'
});
const updates = [];
for await (const partial of partialObjectStream) {
updates.push(partial);
}
expect(updates.length).toBeGreaterThan(0);
const finalUpdate = updates[updates.length - 1];
expect(finalUpdate.recipe).toBeDefined();
expect(finalUpdate.recipe.name).toBeDefined();
expect(Array.isArray(finalUpdate.recipe.ingredients)).toBe(true);
expect(Array.isArray(finalUpdate.recipe.steps)).toBe(true);
}, 30000);
it('should handle streaming errors gracefully', async () => {
let error;
try {
const { elementStream } = streamObject({
model: undefined, // Force an error
output: 'array',
schema: z.string(),
prompt: 'Generate a list of items'
});
for await (const item of elementStream) {
console.log(item);
}
}
catch (e) {
error = e;
}
expect(error).toBeDefined();
expect(error?.message).toBeDefined();
}, 30000);
});
describe('Advanced Features', () => {
it('should support streaming with tool calls', async () => {
const toolCalls = [];
const toolResults = [];
const { fullStream } = streamText({
model,
tools: {
cityAttractions: {
parameters: z.object({ city: z.string() }),
execute: async ({ city }) => ({
attractions: ['attraction1', 'attraction2', 'attraction3']
})
}
},
prompt: 'What are some San Francisco tourist attractions?'
});
for await (const part of fullStream) {
switch (part.type) {
case 'tool-call':
toolCalls.push(part);
break;
case 'tool-result':
toolResults.push(part);
break;
}
}
expect(toolCalls.length).toBeGreaterThan(0);
expect(toolResults.length).toBeGreaterThan(0);
expect(toolCalls[0].toolName).toBe('cityAttractions');
}, 30000);
it('should support predicted outputs', async () => {
const existingCode = `
interface User {
Username: string;
Password: string;
}
`;
const { textStream } = streamText({
model,
messages: [
{
role: 'user',
content: 'Replace the Username property with an Email property.'
},
{
role: 'user',
content: existingCode
}
],
experimental_providerMetadata: {
openai: {
prediction: {
type: 'content',
content: existingCode
}
}
},
seed: 12345
});
let text = '';
for await (const chunk of textStream) {
text += chunk;
}
// Extract just the interface code
const interfaceMatch = text.match(/interface User {[^}]*}/s);
const interfaceCode = interfaceMatch ? interfaceMatch[0] : '';
expect(interfaceCode).toContain('interface User');
expect(interfaceCode).toContain('Email: string');
expect(interfaceCode).not.toContain('Username');
}, 30000);
});
});
//# sourceMappingURL=streaming.test.js.map