diff --git a/examples/structured_output.ts b/examples/structured_output.ts new file mode 100644 index 00000000..b6ec7a70 --- /dev/null +++ b/examples/structured_output.ts @@ -0,0 +1,111 @@ +/** + * Structured Output Example - City Research with Playwright + * + * This example demonstrates intelligent structured output by researching Padova, Italy. + * The agent becomes schema-aware and will intelligently retry to gather missing + * information until all required fields can be populated. + */ + +import { ChatOpenAI } from '@langchain/openai' +import { config } from 'dotenv' +import { z } from 'zod' +import { MCPAgent, MCPClient } from '../index.js' + +// Load environment variables from .env file +config() + +// Define the structured output schema using Zod +const CityInfoSchema = z.object({ + name: z.string().describe('Official name of the city'), + country: z.string().describe('Country where the city is located'), + region: z.string().describe('Region or state within the country'), + population: z.number().describe('Current population count'), + area_km2: z.number().describe('Area in square kilometers'), + foundation_date: z.string().describe('When the city was founded (approximate year or period)'), + mayor: z.string().describe('Current mayor or city leader'), + famous_landmarks: z.array(z.string()).describe('List of famous landmarks, monuments, or attractions'), + universities: z.array(z.string()).describe('List of major universities or educational institutions'), + economy_sectors: z.array(z.string()).describe('Main economic sectors or industries'), + sister_cities: z.array(z.string()).describe('Twin cities or sister cities partnerships'), + historical_significance: z.string().describe('Brief description of historical importance'), + climate_type: z.string().nullable().describe('Type of climate (e.g., Mediterranean, Continental)'), + elevation_meters: z.number().nullable().describe('Elevation above sea level in meters'), +}) + +type CityInfo = z.infer + +async function main() { + const mcpConfig = { + mcpServers: { + playwright: { + command: 'npx', + args: ['@playwright/mcp@latest'], + env: { + DISPLAY: ':1', + }, + }, + }, + } + + const client = new MCPClient(mcpConfig) + const llm = new ChatOpenAI({ model: 'gpt-4o' }) + const agent = new MCPAgent({ llm, client, maxSteps: 50, memoryEnabled: true }) + + try { + // Use structured output with intelligent retry + // The agent will: + // 1. Know exactly what information it needs to collect + // 2. Attempt structured output at finish points + // 3. Continue execution if required information is missing + // 4. Only finish when all required fields can be populated + const result: CityInfo = await agent.run( + ` + Research comprehensive information about the city of Padova (also known as Padua) in Italy. + + Visit multiple reliable sources like Wikipedia, official city websites, tourism sites, + and university websites to gather detailed information including demographics, history, + governance, education, economy, landmarks, and international relationships. + `, + 50, // maxSteps + true, // manageConnector + [], // externalHistory + CityInfoSchema, // outputSchema - this enables structured output + ) + + // Now you have strongly-typed, validated data! + console.log(`Name: ${result.name}`) + console.log(`Country: ${result.country}`) + console.log(`Region: ${result.region}`) + console.log(`Population: ${result.population.toLocaleString()}`) + console.log(`Area: ${result.area_km2} km²`) + console.log(`Foundation: ${result.foundation_date}`) + console.log(`Mayor: ${result.mayor}`) + console.log(`Universities: ${result.universities.join(', ')}`) + console.log(`Economy: ${result.economy_sectors.join(', ')}`) + console.log(`Landmarks: ${result.famous_landmarks.join(', ')}`) + console.log(`Sister Cities: ${result.sister_cities.length > 0 ? result.sister_cities.join(', ') : 'None'}`) + console.log(`Historical Significance: ${result.historical_significance}`) + + if (result.climate_type) { + console.log(`Climate: ${result.climate_type}`) + } + + if (result.elevation_meters !== null) { + console.log(`Elevation: ${result.elevation_meters} meters`) + } + } + catch (error) { + console.error('Error:', error) + } + finally { + await agent.close() + } +} + +// Handle unhandled promise rejections +process.on('unhandledRejection', (reason, promise) => { + console.error('Unhandled Rejection at:', promise, 'reason:', reason) + process.exit(1) +}) + +main().catch(console.error) diff --git a/package.json b/package.json index 479bcfae..6e513c3e 100644 --- a/package.json +++ b/package.json @@ -73,7 +73,8 @@ "example:sandbox": "npm run build && node dist/examples/sandbox_everything.js", "example:oauth": "npm run build && node dist/examples/simple_oauth_example.js", "example:blender": "npm run build && node dist/examples/blender_use.js", - "example:add_server": "npm run build && node dist/examples/add_server_tool.js" + "example:add_server": "npm run build && node dist/examples/add_server_tool.js", + "example:structured": "npm run build && node dist/examples/structured_output.js" }, "dependencies": { "@dmitryrechkin/json-schema-to-zod": "^1.0.1", diff --git a/src/agents/mcp_agent.ts b/src/agents/mcp_agent.ts index 0cc56aed..e241b75a 100644 --- a/src/agents/mcp_agent.ts +++ b/src/agents/mcp_agent.ts @@ -5,6 +5,7 @@ import type { import type { StructuredToolInterface, ToolInterface } from '@langchain/core/tools' import type { StreamEvent } from '@langchain/core/tracers/log_stream' import type { AgentFinish, AgentStep } from 'langchain/agents' +import type { z } from 'zod' import type { MCPClient } from '../client.js' import type { BaseConnector } from '../connectors/base.js' import type { MCPSession } from '../session.js' @@ -281,9 +282,9 @@ export class MCPAgent { return this.disallowedTools } - private async _consumeAndReturn( - generator: AsyncGenerator, - ): Promise { + private async _consumeAndReturn( + generator: AsyncGenerator, + ): Promise { // Manually iterate through the generator to consume the steps. // The for-await-of loop is not used because it discards the generator's // final return value. We need to capture that value when `done` is true. @@ -303,25 +304,47 @@ export class MCPAgent { maxSteps?: number, manageConnector?: boolean, externalHistory?: BaseMessage[], - ): Promise { - const generator = this.stream( + ): Promise + + /** + * Runs the agent with structured output and returns a promise for the typed result. + */ + public async run( + query: string, + maxSteps?: number, + manageConnector?: boolean, + externalHistory?: BaseMessage[], + outputSchema?: z.ZodSchema, + ): Promise + + public async run( + query: string, + maxSteps?: number, + manageConnector?: boolean, + externalHistory?: BaseMessage[], + outputSchema?: z.ZodSchema, + ): Promise { + const generator = this.stream( query, maxSteps, manageConnector, externalHistory, + outputSchema, ) return this._consumeAndReturn(generator) } /** * Runs the agent and yields intermediate steps as an async generator. + * If outputSchema is provided, returns structured output of type T. */ - public async* stream( + public async* stream( query: string, maxSteps?: number, manageConnector = true, externalHistory?: BaseMessage[], - ): AsyncGenerator { + outputSchema?: z.ZodSchema, + ): AsyncGenerator { let result = '' let initializedHere = false const startTime = Date.now() @@ -329,6 +352,40 @@ export class MCPAgent { let stepsTaken = 0 let success = false + // Schema-aware setup for structured output + let structuredLlm: BaseLanguageModelInterface | null = null + let schemaDescription = '' + if (outputSchema) { + query = this._enhanceQueryWithSchema(query, outputSchema) + // Check if withStructuredOutput method exists + if ('withStructuredOutput' in this.llm && typeof (this.llm as any).withStructuredOutput === 'function') { + structuredLlm = (this.llm as any).withStructuredOutput(outputSchema) + } + else { + // Fallback: use the same LLM but we'll handle structure in our helper method + structuredLlm = this.llm + } + // Get schema description for feedback + try { + const schemaType = outputSchema as any + if (schemaType._def && schemaType._def.shape) { + const fields: string[] = [] + for (const [key, fieldSchema] of Object.entries(schemaType._def.shape)) { + const field = fieldSchema as any + const isOptional = field.isOptional?.() ?? field._def?.typeName === 'ZodOptional' + const isNullable = field.isNullable?.() ?? field._def?.typeName === 'ZodNullable' + const description = field._def?.description || field.description || key + fields.push(`- ${key}: ${description} ${(isOptional || isNullable) ? '(optional)' : '(required)'}`) + } + schemaDescription = fields.join('\n') + } + } + catch (e) { + logger.warn(`Could not extract schema details: ${e}`) + schemaDescription = `Schema: ${outputSchema.constructor.name}` + } + } + try { if (manageConnector && !this._initialized) { await this.initialize() @@ -406,7 +463,56 @@ export class MCPAgent { if ((nextStepOutput as AgentFinish).returnValues) { logger.info(`✅ Agent finished at step ${stepNum + 1}`) result = (nextStepOutput as AgentFinish).returnValues?.output ?? 'No output generated' - break + + // If structured output is requested, attempt to create it + if (outputSchema && structuredLlm) { + try { + logger.info('🔧 Attempting structured output...') + const structuredResult = await this._attemptStructuredOutput( + result, + structuredLlm, + outputSchema, + schemaDescription, + ) + + // Add the final response to conversation history if memory is enabled + if (this.memoryEnabled) { + this.addToHistory(new AIMessage(`Structured result: ${JSON.stringify(structuredResult)}`)) + } + + logger.info('✅ Structured output successful') + success = true + return structuredResult as string | T + } + catch (e) { + logger.warn(`⚠️ Structured output failed: ${e}`) + // Continue execution to gather missing information + const missingInfoPrompt = ` + The current result cannot be formatted into the required structure. + Error: ${String(e)} + + Current information: ${result} + + Please continue working to gather the missing information needed for: + ${schemaDescription} + + Focus on finding the specific missing details. + ` + + // Add this as feedback and continue the loop + inputs.input = missingInfoPrompt + if (this.memoryEnabled) { + this.addToHistory(new HumanMessage(missingInfoPrompt)) + } + + logger.info('🔄 Continuing execution to gather missing information...') + continue + } + } + else { + // Regular execution without structured output + break + } } const stepArray = nextStepOutput as AgentStep[] @@ -461,13 +567,42 @@ export class MCPAgent { result = `Agent stopped after reaching the maximum number of steps (${steps}).` } - if (this.memoryEnabled) { + // If structured output was requested but not achieved, attempt one final time + if (outputSchema && structuredLlm && !success) { + try { + logger.info('🔧 Final attempt at structured output...') + const structuredResult = await this._attemptStructuredOutput( + result, + structuredLlm, + outputSchema, + schemaDescription, + ) + + // Add the final response to conversation history if memory is enabled + if (this.memoryEnabled) { + this.addToHistory(new AIMessage(`Structured result: ${JSON.stringify(structuredResult)}`)) + } + + logger.info('✅ Final structured output successful') + success = true + return structuredResult as string | T + } + catch (e) { + logger.error(`❌ Final structured output attempt failed: ${e}`) + throw new Error(`Failed to generate structured output after ${steps} steps: ${e}`) + } + } + + // Add the final response to conversation history if memory is enabled (regular case) + if (this.memoryEnabled && !outputSchema) { this.addToHistory(new AIMessage(result)) } logger.info('🎉 Agent execution complete') success = true - return result + + // Return regular result + return result as string | T } catch (e) { logger.error(`❌ Error running query: ${e}`) @@ -694,4 +829,101 @@ export class MCPAgent { } } } + + /** + * Attempt to create structured output from raw result with validation. + */ + private async _attemptStructuredOutput( + rawResult: string, + structuredLlm: BaseLanguageModelInterface, + outputSchema: z.ZodSchema, + schemaDescription: string, + ): Promise { + const formatPrompt = ` + Please format the following information according to the specified schema. + Extract and structure the relevant information from the content below. + + Required schema fields: + ${schemaDescription} + + Content to format: + ${rawResult} + + Please provide the information in the requested structured format. + If any required information is missing, you must indicate this clearly. + ` + + const structuredResult = await structuredLlm.invoke(formatPrompt) + + // Validate that the result is complete (basic check) + try { + // Use Zod to validate the structured result + const validatedResult = outputSchema.parse(structuredResult) + + // Additional validation for required fields + const schemaType = outputSchema as any + if (schemaType._def && schemaType._def.shape) { + for (const [fieldName, fieldSchema] of Object.entries(schemaType._def.shape)) { + const field = fieldSchema as any + const isOptional = field.isOptional?.() ?? field._def?.typeName === 'ZodOptional' + const isNullable = field.isNullable?.() ?? field._def?.typeName === 'ZodNullable' + if (!isOptional && !isNullable) { + const value = (validatedResult as any)[fieldName] + if (value === null || value === undefined + || (typeof value === 'string' && !value.trim()) + || (Array.isArray(value) && value.length === 0)) { + throw new Error(`Required field '${fieldName}' is missing or empty`) + } + } + } + } + + return validatedResult + } + catch (e) { + logger.debug(`Validation details: ${e}`) + throw e // Re-raise to trigger retry logic + } + } + + /** + * Enhance the query with schema information to make the agent aware of required fields. + */ + private _enhanceQueryWithSchema(query: string, outputSchema: z.ZodSchema): string { + const schemaFields: string[] = [] + + try { + // Get field information from the schema + const schemaType = outputSchema as any + if (schemaType._def && schemaType._def.shape) { + for (const [fieldName, fieldSchema] of Object.entries(schemaType._def.shape)) { + const field = fieldSchema as any + const description = field._def?.description || field.description || fieldName + const isOptional = field.isOptional?.() ?? field._def?.typeName === 'ZodOptional' + const isNullable = field.isNullable?.() ?? field._def?.typeName === 'ZodNullable' + schemaFields.push(`- ${fieldName}: ${description} ${(isOptional || isNullable) ? '(optional)' : '(required)'}`) + } + } + + const schemaDescription = schemaFields.join('\n') + + // Enhance the query with schema awareness + const enhancedQuery = ` + ${query} + + IMPORTANT: Your response must include sufficient information to populate the following structured output: + + ${schemaDescription} + + Make sure you gather ALL the required information during your task execution. + If any required information is missing, continue working to find it. + ` + + return enhancedQuery + } + catch (e) { + logger.warn(`Could not extract schema details: ${e}`) + return query + } + } }