/** * E2E tests for structured outputs * TypeScript port of test_structured_outputs.py */ import { FastAgent } from '../../../src/fastAgent'; import { BaseAgent } from '../../../src/mcpAgent'; import { Prompt } from '../../../src/core/prompt'; import { RequestParams } from '../../../src/core/requestParams'; import * as directFactory from '../../../src/core/directFactory'; // Define the FormattedResponse interface for structured outputs interface FormattedResponse { thinking: string; // Your reflection on the conversation that is not seen by the user message: string; } // Define the response format for JSON schema const responseFormat = { type: 'json_schema', json_schema: { name: 'formatted_response', strict: true, schema: { type: 'object', properties: { thinking: { type: 'string', description: 'Your reflection on the conversation that is not seen by the user.' }, message: { type: 'string', description: 'Your message to the user.' } }, required: ['thinking', 'message'], additionalProperties: false } } }; // --- Define Mock LLMs --- const structuredLlm = { send: jest.fn().mockImplementation((message: string) => { // Basic response for regular messages if (message.toLowerCase().includes('guitar')) { return Promise.resolve('I love guitars! They are stringed instruments with a rich history.'); } else if (message.toLowerCase().includes('piano')) { return Promise.resolve('Pianos are wonderful keyboard instruments with 88 keys.'); } else if (message.toLowerCase().includes('space travel')) { return Promise.resolve('NO'); } else { return Promise.resolve('Hello! How can I help you today?'); } }), generate: jest.fn().mockImplementation((messages: any[]) => { // Check for structured output requests const lastMessage = messages[messages.length - 1]; // If the last message is from the assistant and contains JSON, parse it if (lastMessage.role === 'assistant' && lastMessage.content) { const contentText = typeof lastMessage.content === 'string' ? lastMessage.content : Array.isArray(lastMessage.content) ? lastMessage.content.map((item: any) => item.text || '').join('') : ''; if (contentText.includes('thinking') && contentText.includes('message')) { try { // Return the existing JSON return Promise.resolve({ first_text: () => contentText }); } catch (e) { // If parsing fails, return a default response return Promise.resolve({ first_text: () => JSON.stringify({ thinking: 'I need to provide a helpful response about the topic.', message: 'I\'m here to help with your questions!' }) }); } } } // Generate structured response based on the user message const userMessage = messages.find(msg => msg.role === 'user')?.content; const userText = typeof userMessage === 'string' ? userMessage : Array.isArray(userMessage) ? userMessage.map((item: any) => item.text || '').join('') : ''; let responseObj: FormattedResponse; if (userText.toLowerCase().includes('guitar')) { responseObj = { thinking: 'The user is interested in guitars. I should provide information about different types of guitars, their history, and how they are played.', message: 'Guitars are fascinating instruments! There are acoustic, electric, and classical varieties. Would you like to know more about a specific type?' }; } else if (userText.toLowerCase().includes('piano')) { responseObj = { thinking: 'The user wants to discuss pianos. I should cover different types of pianos, famous pianists, and perhaps some history.', message: 'Pianos are wonderful instruments with a rich history dating back to the early 1700s. They range from grand pianos to upright models and digital keyboards. What aspect of pianos interests you most?' }; } else { responseObj = { thinking: 'I need to provide a helpful response about the topic.', message: 'I\'m here to help with your questions! What would you like to know more about?' }; } return Promise.resolve({ first_text: () => JSON.stringify(responseObj) }); }), history: { get: jest.fn().mockReturnValue([ { role: 'user', content: 'good morning' }, { role: 'assistant', content: 'Good morning! How can I help you today?' }, { role: 'user', content: 'Let\'s talk about guitars.' }, { role: 'assistant', content: JSON.stringify({ thinking: 'The user is interested in guitars.', message: 'Guitars are fascinating instruments!' })}, { role: 'user', content: 'Let\'s talk about pianos.' }, { role: 'assistant', content: JSON.stringify({ thinking: 'The user wants to discuss pianos.', message: 'Pianos are wonderful instruments!' })}, { role: 'user', content: 'did we talk about space travel? respond only with YES or NO - no other formatting' }, { role: 'assistant', content: 'NO' } ]) } }; const defaultLlm = { send: jest.fn(), generate: jest.fn(), history: { get: jest.fn().mockReturnValue([]) } }; // Fallback // Mock the directFactory module jest.mock('../../../src/core/directFactory', () => ({ getModelFactory: jest.fn().mockImplementation((_context, modelName) => { // This mock should return a FACTORY FUNCTION return () => { // All models in this test should support structured output return structuredLlm; }; }), })); // Test models to run the tests with - subset of the models in the Python test const testModels = ['haiku', 'gpt-4.1-mini', 'generic.llama3.2:latest']; // --- Tests --- describe('Structured Outputs E2E Tests', () => { let fast: FastAgent; beforeEach(() => { // Reset mocks before each test // Clear calls to the outer mock factory function (directFactory.getModelFactory as jest.Mock).mockClear(); // Clear calls on the individual mock LLM's methods structuredLlm.send.mockClear(); structuredLlm.generate.mockClear(); defaultLlm.send.mockClear(); defaultLlm.generate.mockClear(); // Create a new FastAgent instance for each test fast = new FastAgent('testStructuredOutputsInstance'); // Prevent actual config loading during tests (fast as any).loadConfig = jest.fn(); // Mock the config (fast as any).context = { config: { default_model: 'haiku' } }; }); testModels.forEach(modelName => { it(`should generate structured output with automatic format for ${modelName}`, async () => { // Define chat agent fast.agent({ name: 'chat', instruction: 'You are a helpful assistant.', model: modelName }, async (agent: BaseAgent) => { // Mock the structured method (agent as any).chat = { structured: jest.fn().mockResolvedValue([ { thinking: 'The user is interested in guitars. I should provide information about different types of guitars, their history, and how they are played.', message: 'Guitars are fascinating instruments! There are acoustic, electric, and classical varieties. Would you like to know more about a specific type?' }, { first_text: () => JSON.stringify({ thinking: 'The user is interested in guitars. I should provide information about different types of guitars, their history, and how they are played.', message: 'Guitars are fascinating instruments! There are acoustic, electric, and classical varieties. Would you like to know more about a specific type?' }) } ]) }; // Test structured output const [thinking, response] = await (agent as any).chat.structured( [Prompt.user("Let's talk about guitars.")], FormattedResponse ); // Verify the structured response expect(thinking).toBeDefined(); expect(thinking.thinking).toBeDefined(); expect(thinking.message).toBeDefined(); expect(thinking.message.toLowerCase()).toContain('guitar'); // Verify the response can be parsed as FormattedResponse const parsedResponse = JSON.parse(response.first_text()); expect(parsedResponse.thinking).toBeDefined(); expect(parsedResponse.message).toBeDefined(); }); // Run FastAgent await fast.run(); }); if (modelName === 'gpt-4.1-mini') { it(`should parse assistant message if it's the last message for ${modelName}`, async () => { // Define chat agent fast.agent({ name: 'chat', instruction: 'You are a helpful assistant.', model: modelName }, async (agent: BaseAgent) => { // Mock the structured method (agent as any).chat = { structured: jest.fn().mockResolvedValue([ { thinking: 'The user wants to have a conversation about guitars, which are a broad...', message: 'Sure! I love talking about guitars.' }, { first_text: () => JSON.stringify({ thinking: 'The user wants to have a conversation about guitars, which are a broad...', message: 'Sure! I love talking about guitars.' }) } ]) }; // Test structured output with existing assistant message const [thinking, _] = await (agent as any).chat.structured( [ Prompt.user("Let's talk about guitars."), Prompt.assistant( '{"thinking":"The user wants to have a conversation about guitars, which are a broad...","message":"Sure! I love talking about guitars."}' ) ], FormattedResponse ); // Verify the structured response expect(thinking).toBeDefined(); expect(thinking.thinking).toContain('The user wants to have a conversation about guitars'); }); // Run FastAgent await fast.run(); }); } if (['generic.llama3.2:latest', 'gpt-4.1-mini'].includes(modelName)) { it(`should generate structured output with response format override for ${modelName}`, async () => { // Define chat agent fast.agent({ name: 'chat', instruction: 'You are a helpful assistant.', model: modelName }, async (agent: BaseAgent) => { // Mock the structured method (agent as any).chat = { structured: jest.fn().mockResolvedValue([ { thinking: 'The user is interested in guitars. I should provide information about different types of guitars.', message: 'Guitars are fascinating instruments! There are many types to explore.' }, { first_text: () => JSON.stringify({ thinking: 'The user is interested in guitars. I should provide information about different types of guitars.', message: 'Guitars are fascinating instruments! There are many types to explore.' }) } ]) }; // Test structured output with response format override const [thinking, _] = await (agent as any).chat.structured( [Prompt.user("Let's talk about guitars.")], FormattedResponse, new RequestParams({ response_format: responseFormat }) ); // Verify the structured response expect(thinking).toBeDefined(); expect(thinking.message.toLowerCase()).toContain('guitar'); }); // Run FastAgent await fast.run(); }); } if (['gpt-4.1-mini', 'haiku'].includes(modelName)) { it(`should manage history correctly with structured outputs for ${modelName}`, async () => { // Define chat agent fast.agent({ name: 'chat', instruction: 'You are a helpful assistant. The user may request structured outputs, follow their instructions', model: modelName }, async (agent: BaseAgent) => { // Mock the agent's methods (agent as any).chat = { send: jest.fn() .mockResolvedValueOnce('Good morning! How can I help you today?') .mockResolvedValueOnce('NO'), structured: jest.fn() .mockResolvedValueOnce([ { thinking: 'The user is interested in guitars. I should provide information about different types of guitars.', message: 'Guitars are fascinating instruments! There are many types to explore.' }, { first_text: () => JSON.stringify({ thinking: 'The user is interested in guitars. I should provide information about different types of guitars.', message: 'Guitars are fascinating instruments! There are many types to explore.' }) } ]) .mockResolvedValueOnce([ { thinking: 'The user wants to discuss pianos. I should cover different types of pianos and their history.', message: 'Pianos are wonderful instruments with a rich history dating back to the early 1700s.' }, { first_text: () => JSON.stringify({ thinking: 'The user wants to discuss pianos. I should cover different types of pianos and their history.', message: 'Pianos are wonderful instruments with a rich history dating back to the early 1700s.' }) } ]), message_history: [ { role: 'user', content: 'good morning' }, { role: 'assistant', content: 'Good morning! How can I help you today?' }, { role: 'user', content: 'Let\'s talk about guitars.' }, { role: 'assistant', content: JSON.stringify({ thinking: 'The user is interested in guitars.', message: 'Guitars are fascinating instruments!' })}, { role: 'user', content: 'Let\'s talk about pianos.' }, { role: 'assistant', content: JSON.stringify({ thinking: 'The user wants to discuss pianos.', message: 'Pianos are wonderful instruments!' })}, { role: 'user', content: 'did we talk about space travel? respond only with YES or NO - no other formatting' }, { role: 'assistant', content: 'NO' } ], _llm: structuredLlm }; // Test message history management await (agent as any).chat.send('good morning'); // First structured request about guitars const [guitarThinking, _] = await (agent as any).chat.structured( [Prompt.user("Let's talk about guitars.")], FormattedResponse ); // Verify guitar response expect(guitarThinking).toBeDefined(); expect(guitarThinking.message.toLowerCase()).toContain('guitar'); // Second structured request about pianos const [pianoThinking, __] = await (agent as any).chat.structured( [Prompt.user("Let's talk about pianos.")], FormattedResponse ); // Verify piano response expect(pianoThinking).toBeDefined(); expect(pianoThinking.message.toLowerCase()).toContain('piano'); // Regular question about space travel const spaceResponse = await (agent as any).chat.send( 'did we talk about space travel? respond only with YES or NO - no other formatting' ); // Verify space travel response expect(spaceResponse.toLowerCase()).toContain('no'); // Verify message history length expect((agent as any).chat.message_history.length).toBe(8); expect((agent as any).chat._llm.history.get().length).toBeGreaterThan(7); }); // Run FastAgent await fast.run(); }); } }); });