From 8f41f048da960ff354ec4e7e5400228fc2e12d8c Mon Sep 17 00:00:00 2001 From: pietrozullo Date: Fri, 1 Aug 2025 15:20:35 -0700 Subject: [PATCH 1/2] Structured output --- examples/structured_output.ts | 110 +++++++++++++++ package.json | 3 +- src/agents/mcp_agent.ts | 245 ++++++++++++++++++++++++++++++++-- 3 files changed, 347 insertions(+), 11 deletions(-) create mode 100644 examples/structured_output.ts diff --git a/examples/structured_output.ts b/examples/structured_output.ts new file mode 100644 index 00000000..9a2b9b89 --- /dev/null +++ b/examples/structured_output.ts @@ -0,0 +1,110 @@ +/** + * 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) \ No newline at end of file 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..f68d4b4d 100644 --- a/src/agents/mcp_agent.ts +++ b/src/agents/mcp_agent.ts @@ -8,6 +8,7 @@ import type { AgentFinish, AgentStep } from 'langchain/agents' import type { MCPClient } from '../client.js' import type { BaseConnector } from '../connectors/base.js' import type { MCPSession } from '../session.js' +import type { z } from 'zod' import { AIMessage, HumanMessage, @@ -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,31 +304,85 @@ 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() const toolsUsedNames: string[] = [] 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) { @@ -406,7 +461,54 @@ 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 +563,41 @@ 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 +824,99 @@ 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 + } + } } From edbe528e4d1eac33aa01f8c7890f78aa6e578879 Mon Sep 17 00:00:00 2001 From: pietrozullo Date: Fri, 1 Aug 2025 15:20:56 -0700 Subject: [PATCH 2/2] Lint fix --- examples/structured_output.ts | 23 ++++++----- src/agents/mcp_agent.ts | 77 +++++++++++++++++++---------------- 2 files changed, 54 insertions(+), 46 deletions(-) diff --git a/examples/structured_output.ts b/examples/structured_output.ts index 9a2b9b89..b6ec7a70 100644 --- a/examples/structured_output.ts +++ b/examples/structured_output.ts @@ -2,7 +2,7 @@ * 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 + * The agent becomes schema-aware and will intelligently retry to gather missing * information until all required fields can be populated. */ @@ -66,10 +66,10 @@ async function main() { 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 + 50, // maxSteps + true, // manageConnector + [], // externalHistory + CityInfoSchema, // outputSchema - this enables structured output ) // Now you have strongly-typed, validated data! @@ -85,18 +85,19 @@ async function main() { 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) { + } + catch (error) { console.error('Error:', error) - } finally { + } + finally { await agent.close() } } @@ -107,4 +108,4 @@ process.on('unhandledRejection', (reason, promise) => { process.exit(1) }) -main().catch(console.error) \ No newline at end of file +main().catch(console.error) diff --git a/src/agents/mcp_agent.ts b/src/agents/mcp_agent.ts index f68d4b4d..e241b75a 100644 --- a/src/agents/mcp_agent.ts +++ b/src/agents/mcp_agent.ts @@ -5,10 +5,10 @@ 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' -import type { z } from 'zod' import { AIMessage, HumanMessage, @@ -305,7 +305,7 @@ export class MCPAgent { manageConnector?: boolean, externalHistory?: BaseMessage[], ): Promise - + /** * Runs the agent with structured output and returns a promise for the typed result. */ @@ -316,7 +316,7 @@ export class MCPAgent { externalHistory?: BaseMessage[], outputSchema?: z.ZodSchema, ): Promise - + public async run( query: string, maxSteps?: number, @@ -351,7 +351,7 @@ export class MCPAgent { const toolsUsedNames: string[] = [] let stepsTaken = 0 let success = false - + // Schema-aware setup for structured output let structuredLlm: BaseLanguageModelInterface | null = null let schemaDescription = '' @@ -360,7 +360,8 @@ export class MCPAgent { // 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 { + } + else { // Fallback: use the same LLM but we'll handle structure in our helper method structuredLlm = this.llm } @@ -378,7 +379,8 @@ export class MCPAgent { } schemaDescription = fields.join('\n') } - } catch (e) { + } + catch (e) { logger.warn(`Could not extract schema details: ${e}`) schemaDescription = `Schema: ${outputSchema.constructor.name}` } @@ -461,27 +463,28 @@ 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' - + // 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 + 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) { + } + catch (e) { logger.warn(`⚠️ Structured output failed: ${e}`) // Continue execution to gather missing information const missingInfoPrompt = ` @@ -495,17 +498,18 @@ export class MCPAgent { 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 { + } + else { // Regular execution without structured output break } @@ -568,21 +572,22 @@ export class MCPAgent { try { logger.info('🔧 Final attempt at structured output...') const structuredResult = await this._attemptStructuredOutput( - result, - structuredLlm, - outputSchema, - schemaDescription + 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) { + } + catch (e) { logger.error(`❌ Final structured output attempt failed: ${e}`) throw new Error(`Failed to generate structured output after ${steps} steps: ${e}`) } @@ -595,7 +600,7 @@ export class MCPAgent { logger.info('🎉 Agent execution complete') success = true - + // Return regular result return result as string | T } @@ -832,7 +837,7 @@ export class MCPAgent { rawResult: string, structuredLlm: BaseLanguageModelInterface, outputSchema: z.ZodSchema, - schemaDescription: string + schemaDescription: string, ): Promise { const formatPrompt = ` Please format the following information according to the specified schema. @@ -854,7 +859,7 @@ export class MCPAgent { 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) { @@ -864,17 +869,18 @@ export class MCPAgent { 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)) { + 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) { + } + catch (e) { logger.debug(`Validation details: ${e}`) throw e // Re-raise to trigger retry logic } @@ -900,7 +906,7 @@ export class MCPAgent { } const schemaDescription = schemaFields.join('\n') - + // Enhance the query with schema awareness const enhancedQuery = ` ${query} @@ -912,9 +918,10 @@ export class MCPAgent { 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) { + } + catch (e) { logger.warn(`Could not extract schema details: ${e}`) return query }