diff --git a/.gitignore b/.gitignore index dd0c5f66..c76b9460 100644 --- a/.gitignore +++ b/.gitignore @@ -31,6 +31,7 @@ lerna-debug.log* !.vscode/extensions.json !.vscode/settings.json .claude/ +CLAUDE.md .idea/ .DS_Store *.suo diff --git a/docs/components/ai.mdx b/docs/components/ai.mdx index 2cfde73f..f8afb31b 100644 --- a/docs/components/ai.mdx +++ b/docs/components/ai.mdx @@ -103,10 +103,16 @@ An autonomous agent that uses reasoning steps and tool-calling to solve complex | `temperature` | Number | Reasoning creativity (default 0.7) | | `stepLimit` | Number | Max "Think -> Act -> Observe" loops (1-12) | | `memorySize` | Number | Number of previous turns to retain in context | +| `structuredOutputEnabled` | Toggle | Enable to enforce a specific JSON output structure | +| `schemaType` | Select | How to define the schema: `json-example` or `json-schema` | +| `jsonExample` | JSON | Example JSON object for schema inference (all properties become required) | +| `jsonSchema` | JSON | Full JSON Schema definition for precise validation | +| `autoFixFormat` | Toggle | Attempt to extract valid JSON from malformed responses | | Output | Type | Description | |--------|------|-------------| | `responseText`| Text | Final answer after reasoning is complete | +| `structuredOutput` | JSON | Parsed structured output (when enabled) | | `conversationState` | JSON | Updated state to pass to the next agent node | | `reasoningTrace` | JSON | Detailed step-by-step logs of the agent's thoughts | | `agentRunId` | Text | Unique session ID for tracking and streaming | @@ -168,6 +174,19 @@ Analyze incoming security alerts to filter out false positives. An agent that searches through logs and performs lookups to investigate a specific IP address. **Task:** "Investigate the IP {{ip}} using the available Splunk and VirusTotal tools." +### Structured Output for Data Extraction +**Flow:** `Provider` → `AI Agent` (with Structured Output enabled) + +Extract structured data from unstructured security reports. Enable **Structured Output** and provide a JSON example: +```json +{ + "severity": "high", + "affected_systems": ["web-server-01"], + "remediation_steps": ["Patch CVE-2024-1234", "Restart service"] +} +``` +The agent will always return validated JSON matching this schema, ready for downstream processing. + --- ## Best Practices @@ -177,7 +196,7 @@ An agent that searches through logs and performs lookups to investigate a specif ### Prompt Engineering -1. **Format Outputs**: If you need JSON for a downstream node, ask for it explicitly in the prompt: "Return only valid JSON with fields 'risk' and 'reason'." +1. **Use Structured Output**: When you need consistent JSON for downstream nodes, enable **Structured Output** instead of relying on prompt instructions. This guarantees schema compliance and eliminates parsing errors. 2. **Use System Prompts**: Set high-level rules (e.g., "You are a senior security researcher") in the System Prompt parameter instead of the User Input. 3. **Variable Injection**: Use `{{variableName}}` syntax to inject data from upstream nodes into your prompts. diff --git a/frontend/src/components/workflow/ConfigPanel.tsx b/frontend/src/components/workflow/ConfigPanel.tsx index 007c955f..337cdd60 100644 --- a/frontend/src/components/workflow/ConfigPanel.tsx +++ b/frontend/src/components/workflow/ConfigPanel.tsx @@ -987,22 +987,19 @@ export function ConfigPanel({ defaultOpen={true} >
- {/* Sort parameters: select types first, then others */} - {componentParameters - .slice() - .sort((a, b) => { - // Select parameters go first - const aIsSelect = a.type === 'select' - const bIsSelect = b.type === 'select' - if (aIsSelect && !bIsSelect) return -1 - if (!aIsSelect && bIsSelect) return 1 - return 0 - }) - .map((param, index) => ( + {/* Render parameters in component definition order to preserve hierarchy */} + {componentParameters.map((param, index) => { + // Only show border between top-level parameters (not nested ones) + const isTopLevel = !param.visibleWhen + const prevParam = index > 0 ? componentParameters[index - 1] : null + const prevIsTopLevel = prevParam ? !prevParam.visibleWhen : false + const showBorder = index > 0 && isTopLevel && prevIsTopLevel + + return (
0 && "border-t border-border pt-3" + showBorder && "border-t border-border pt-3" )} >
- ))} + ) + })}
)} diff --git a/frontend/src/components/workflow/ParameterField.tsx b/frontend/src/components/workflow/ParameterField.tsx index e49b6534..ba608688 100644 --- a/frontend/src/components/workflow/ParameterField.tsx +++ b/frontend/src/components/workflow/ParameterField.tsx @@ -2,6 +2,7 @@ import { useCallback, useEffect, useMemo, useRef, useState } from 'react' import { Input } from '@/components/ui/input' import { Badge } from '@/components/ui/badge' import { Checkbox } from '@/components/ui/checkbox' +import { Switch } from '@/components/ui/switch' import { Button } from '@/components/ui/button' import { useNavigate } from 'react-router-dom' import { RuntimeInputsEditor } from './RuntimeInputsEditor' @@ -399,18 +400,18 @@ export function ParameterField({ case 'boolean': return ( -
- onChange(checked)} - /> +
+ onChange(checked)} + />
) @@ -695,13 +696,17 @@ export function ParameterField({ if (!jsonTextareaRef.current) return let textValue = '' + let needsNormalization = false + if (value === undefined || value === null || value === '') { textValue = '' } else if (typeof value === 'string') { textValue = value } else { + // If Value is an object - normalize to string try { textValue = JSON.stringify(value, null, 2) + needsNormalization = true } catch (error) { console.error('Failed to serialize JSON parameter value', error) return @@ -714,7 +719,12 @@ export function ParameterField({ setJsonError(null) isExternalJsonUpdateRef.current = false } - }, [value]) + + // Normalize object values to string + if (needsNormalization) { + onChange(textValue) + } + }, [value, onChange]) // Sync to parent only on blur for native undo behavior const handleJsonBlur = useCallback(() => { @@ -728,9 +738,9 @@ export function ParameterField({ } try { - const parsed = JSON.parse(nextValue) + JSON.parse(nextValue) // Validate JSON syntax setJsonError(null) - onChange(parsed) + onChange(nextValue) // Pass string, not parsed object } catch (error) { setJsonError('Invalid JSON') // Keep showing error, don't update parent @@ -1010,6 +1020,52 @@ interface ParameterFieldWrapperProps { componentId?: string parameters?: Record | undefined onUpdateParameter?: (paramId: string, value: any) => void + allComponentParameters?: Parameter[] +} + +/** + * Checks if a parameter should be visible based on its visibleWhen conditions. + * Returns true if all conditions are met or if no conditions exist. + */ +function shouldShowParameter( + parameter: Parameter, + allParameters: Record | undefined +): boolean { + // If no visibleWhen conditions, always show + if (!parameter.visibleWhen) { + return true + } + + // If we have conditions but no parameter values to check against, hide by default + if (!allParameters) { + return false + } + + // Check all conditions in visibleWhen object + for (const [key, expectedValue] of Object.entries(parameter.visibleWhen)) { + const actualValue = allParameters[key] + if (actualValue !== expectedValue) { + return false + } + } + + return true +} + +/** + * Checks if a boolean parameter acts as a header toggle (controls visibility of other params). + * Returns true if other parameters have visibleWhen conditions referencing this parameter. + */ +function isHeaderToggleParameter( + parameter: Parameter, + allComponentParameters: Parameter[] | undefined +): boolean { + if (parameter.type !== 'boolean' || !allComponentParameters) return false + + // Check if any other parameter has visibleWhen referencing this param + return allComponentParameters.some( + (p) => p.visibleWhen && parameter.id in p.visibleWhen + ) } /** @@ -1023,7 +1079,13 @@ export function ParameterFieldWrapper({ componentId, parameters, onUpdateParameter, + allComponentParameters, }: ParameterFieldWrapperProps) { + // Check visibility conditions + if (!shouldShowParameter(parameter, parameters)) { + return null + } + // Special case: Runtime Inputs Editor for Entry Point if (parameter.id === 'runtimeInputs') { return ( @@ -1137,19 +1199,54 @@ export function ParameterFieldWrapper({ ) } - // Standard parameter field rendering - return ( -
-
- - {parameter.required && ( - *required + // Check if this is a nested/conditional parameter (has visibleWhen) + const isNestedParameter = Boolean(parameter.visibleWhen) + + // Check if this is a header toggle (boolean that controls other params' visibility) + const isHeaderToggle = isHeaderToggleParameter(parameter, allComponentParameters) + + // Header toggle rendering + if (isHeaderToggle) { + return ( +
+
+ + onChange(checked)} + /> +
+ {parameter.description && ( +

+ {parameter.description} +

)}
+ ) + } + + // Standard parameter field rendering + const isBooleanParameter = parameter.type === 'boolean' + + return ( +
+ {/* Label and required indicator - skip for boolean (label is inside) */} + {!isBooleanParameter && ( +
+ + {parameter.required && ( + *required + )} +
+ )} - {parameter.description && ( + {/* Description before the input field - for non-boolean parameters */} + {!isBooleanParameter && parameter.description && (

{parameter.description}

@@ -1165,6 +1262,13 @@ export function ParameterFieldWrapper({ onUpdateParameter={onUpdateParameter} /> + {/* Description after field (toggle control) - for boolean parameters */} + {isBooleanParameter && parameter.description && ( +

+ {parameter.description} +

+ )} + {parameter.helpText && (

💡 {parameter.helpText} diff --git a/frontend/src/components/workflow/WorkflowNode.tsx b/frontend/src/components/workflow/WorkflowNode.tsx index d6bff8d1..de71f3fc 100644 --- a/frontend/src/components/workflow/WorkflowNode.tsx +++ b/frontend/src/components/workflow/WorkflowNode.tsx @@ -331,8 +331,9 @@ function ParametersDisplay({ position = 'bottom' }: ParametersDisplayProps) { // Show required parameters and important select parameters (like mode) + // Exclude nested parameters (those with visibleWhen) like schemaType const selectParams = componentParameters.filter( - param => param.type === 'select' && !param.required + param => param.type === 'select' && !param.required && !param.visibleWhen ) const paramsToShow = [...requiredParams, ...selectParams] @@ -1412,7 +1413,6 @@ export const WorkflowNode = ({ data, selected, id }: NodeProps) => { const branchingOutputs = effectiveOutputs.filter((o: any) => o.isBranching) if (branchingOutputs.length === 0) return null - // Determine which branch is active (for execution mode) // Determine which branches are active (for execution mode) const data = isTimelineActive ? visualState.lastEvent?.data : undefined const activatedPorts = data?.activatedPorts @@ -1511,57 +1511,64 @@ export const WorkflowNode = ({ data, selected, id }: NodeProps) => {

) - })()} + })() + } {/* Enhanced Execution Status Messages */} - {isTimelineActive && ( -
- {visualState.lastEvent && ( -
-
- Last: {visualState.lastEvent.type.replace('_', ' ')} -
- {visualState.lastEvent.message && ( -
- {visualState.lastEvent.message} + { + isTimelineActive && ( +
+ {visualState.lastEvent && ( +
+
+ Last: {visualState.lastEvent.type.replace('_', ' ')}
- )} -
- )} + {visualState.lastEvent.message && ( +
+ {visualState.lastEvent.message} +
+ )} +
+ )} + + {/* Legacy status messages */} + {!isTimelineActive && nodeData.status === 'success' && nodeData.executionTime && ( + + {nodeData.executionTime}ms + + )} - {/* Legacy status messages */} - {!isTimelineActive && nodeData.status === 'success' && nodeData.executionTime && ( - - {nodeData.executionTime}ms + {!isTimelineActive && nodeData.status === 'error' && nodeData.error && ( + + ✗ {nodeData.error} + + )} +
+ ) + } + + {/* Legacy status messages (when not in timeline mode) */} + { + !isTimelineActive && nodeData.status === 'success' && nodeData.executionTime && ( +
+ + ✓ {nodeData.executionTime}ms - )} +
+ ) + } - {!isTimelineActive && nodeData.status === 'error' && nodeData.error && ( + { + !isTimelineActive && nodeData.status === 'error' && nodeData.error && ( +
✗ {nodeData.error} - )} -
- )} - - {/* Legacy status messages (when not in timeline mode) */} - {!isTimelineActive && nodeData.status === 'success' && nodeData.executionTime && ( -
- - ✓ {nodeData.executionTime}ms - -
- )} - - {!isTimelineActive && nodeData.status === 'error' && nodeData.error && ( -
- - ✗ {nodeData.error} - -
- )} -
-
+
+ ) + } +
+ ) } diff --git a/frontend/src/schemas/component.ts b/frontend/src/schemas/component.ts index 19e68477..14806ad5 100644 --- a/frontend/src/schemas/component.ts +++ b/frontend/src/schemas/component.ts @@ -101,6 +101,8 @@ export const ParameterSchema = z.object({ placeholder: z.string().optional(), description: z.string().optional(), helpText: z.string().optional(), + /** Conditional visibility: parameter is shown only when all conditions are met */ + visibleWhen: z.record(z.string(), z.any()).optional(), }) export type Parameter = z.infer diff --git a/packages/component-sdk/src/types.ts b/packages/component-sdk/src/types.ts index 55008f67..4386ef41 100644 --- a/packages/component-sdk/src/types.ts +++ b/packages/component-sdk/src/types.ts @@ -191,6 +191,8 @@ export interface ComponentParameterMetadata { min?: number; max?: number; rows?: number; + /** Conditional visibility: parameter is shown only when all conditions are met */ + visibleWhen?: Record; } export type ComponentAuthorType = 'shipsecai' | 'community'; diff --git a/worker/src/components/ai/__tests__/ai-agent.test.ts b/worker/src/components/ai/__tests__/ai-agent.test.ts index b2ecaa22..0b5f0db1 100644 --- a/worker/src/components/ai/__tests__/ai-agent.test.ts +++ b/worker/src/components/ai/__tests__/ai-agent.test.ts @@ -2,7 +2,7 @@ import { beforeAll, beforeEach, describe, expect, test, vi } from 'bun:test'; import '../../index'; import type { ExecutionContext } from '@shipsec/component-sdk'; import { componentRegistry, runComponentWithRunner } from '@shipsec/component-sdk'; -import type { ToolLoopAgentClass, StepCountIsFn, ToolFn, CreateOpenAIFn, CreateGoogleGenerativeAIFn } from '../ai-agent'; +import type { ToolLoopAgentClass, StepCountIsFn, ToolFn, CreateOpenAIFn, CreateGoogleGenerativeAIFn, GenerateObjectFn, GenerateTextFn } from '../ai-agent'; const makeAgentResult = (overrides: Record = {}) => ({ text: 'Agent final answer', @@ -463,4 +463,262 @@ describe('core.ai.agent component', () => { agentRunId: expect.any(String), }); }); + + describe('Structured Output', () => { + const generateObjectMock = vi.fn(); + const generateTextMock = vi.fn(); + + beforeEach(() => { + generateObjectMock.mockReset(); + generateTextMock.mockReset(); + }); + + test('generates structured output from JSON example', async () => { + const component = componentRegistry.get('core.ai.agent'); + expect(component).toBeDefined(); + + generateObjectMock.mockResolvedValue({ + object: { name: 'Test User', age: 30 }, + usage: { promptTokens: 10, completionTokens: 20, totalTokens: 30 }, + }); + + const params = { + userInput: 'Generate user data', + conversationState: undefined, + chatModel: { + provider: 'openai', + modelId: 'gpt-4o-mini', + }, + modelApiKey: 'sk-openai-from-secret', + systemPrompt: '', + temperature: 0.7, + maxTokens: 256, + memorySize: 8, + stepLimit: 4, + structuredOutputEnabled: true, + schemaType: 'json-example', + jsonExample: '{"name": "example", "age": 0}', + autoFixFormat: false, + }; + + const result = (await runComponentWithRunner( + component!.runner, + (p: any, ctx: any) => + (component!.execute as any)(p, ctx, { + ToolLoopAgent: MockToolLoopAgent as unknown as ToolLoopAgentClass, + stepCountIs: stepCountIsMock as unknown as StepCountIsFn, + tool: ((definition: any) => definition) as unknown as ToolFn, + createOpenAI: openAiFactoryMock as unknown as CreateOpenAIFn, + createGoogleGenerativeAI: googleFactoryMock as unknown as CreateGoogleGenerativeAIFn, + generateObject: generateObjectMock as unknown as GenerateObjectFn, + generateText: generateTextMock as unknown as GenerateTextFn, + }), + params, + workflowContext, + )) as any; + + expect(generateObjectMock).toHaveBeenCalledTimes(1); + expect(result.structuredOutput).toEqual({ name: 'Test User', age: 30 }); + expect(result.responseText).toBe(JSON.stringify({ name: 'Test User', age: 30 }, null, 2)); + // ToolLoopAgent should NOT be called when structured output is enabled + expect(toolLoopAgentConstructorMock).not.toHaveBeenCalled(); + }); + + test('generates structured output from JSON Schema', async () => { + const component = componentRegistry.get('core.ai.agent'); + expect(component).toBeDefined(); + + generateObjectMock.mockResolvedValue({ + object: { title: 'Hello World', count: 42 }, + usage: { promptTokens: 15, completionTokens: 25, totalTokens: 40 }, + }); + + const params = { + userInput: 'Generate article data', + conversationState: undefined, + chatModel: { + provider: 'gemini', + modelId: 'gemini-2.5-flash', + }, + modelApiKey: 'gm-gemini-from-secret', + systemPrompt: '', + temperature: 0.5, + maxTokens: 512, + memorySize: 8, + stepLimit: 4, + structuredOutputEnabled: true, + schemaType: 'json-schema', + jsonSchema: JSON.stringify({ + type: 'object', + properties: { + title: { type: 'string' }, + count: { type: 'integer' }, + }, + required: ['title', 'count'], + }), + autoFixFormat: false, + }; + + const result = (await runComponentWithRunner( + component!.runner, + (p: any, ctx: any) => + (component!.execute as any)(p, ctx, { + ToolLoopAgent: MockToolLoopAgent as unknown as ToolLoopAgentClass, + stepCountIs: stepCountIsMock as unknown as StepCountIsFn, + tool: ((definition: any) => definition) as unknown as ToolFn, + createOpenAI: openAiFactoryMock as unknown as CreateOpenAIFn, + createGoogleGenerativeAI: googleFactoryMock as unknown as CreateGoogleGenerativeAIFn, + generateObject: generateObjectMock as unknown as GenerateObjectFn, + generateText: generateTextMock as unknown as GenerateTextFn, + }), + params, + workflowContext, + )) as any; + + expect(generateObjectMock).toHaveBeenCalledTimes(1); + expect(result.structuredOutput).toEqual({ title: 'Hello World', count: 42 }); + }); + + test('uses auto-fix when generateObject fails', async () => { + const component = componentRegistry.get('core.ai.agent'); + expect(component).toBeDefined(); + + generateObjectMock.mockRejectedValue(new Error('Schema validation failed')); + generateTextMock.mockResolvedValue({ + text: '```json\n{"name": "Fixed User", "age": 25}\n```', + usage: { promptTokens: 20, completionTokens: 30, totalTokens: 50 }, + }); + + const params = { + userInput: 'Generate user data', + conversationState: undefined, + chatModel: { + provider: 'openai', + modelId: 'gpt-4o-mini', + }, + modelApiKey: 'sk-openai-from-secret', + systemPrompt: '', + temperature: 0.7, + maxTokens: 256, + memorySize: 8, + stepLimit: 4, + structuredOutputEnabled: true, + schemaType: 'json-example', + jsonExample: '{"name": "example", "age": 0}', + autoFixFormat: true, + }; + + const result = (await runComponentWithRunner( + component!.runner, + (p: any, ctx: any) => + (component!.execute as any)(p, ctx, { + ToolLoopAgent: MockToolLoopAgent as unknown as ToolLoopAgentClass, + stepCountIs: stepCountIsMock as unknown as StepCountIsFn, + tool: ((definition: any) => definition) as unknown as ToolFn, + createOpenAI: openAiFactoryMock as unknown as CreateOpenAIFn, + createGoogleGenerativeAI: googleFactoryMock as unknown as CreateGoogleGenerativeAIFn, + generateObject: generateObjectMock as unknown as GenerateObjectFn, + generateText: generateTextMock as unknown as GenerateTextFn, + }), + params, + workflowContext, + )) as any; + + expect(generateObjectMock).toHaveBeenCalledTimes(1); + expect(generateTextMock).toHaveBeenCalledTimes(1); + expect(result.structuredOutput).toEqual({ name: 'Fixed User', age: 25 }); + }); + + test('returns null structuredOutput when not enabled', async () => { + const component = componentRegistry.get('core.ai.agent'); + expect(component).toBeDefined(); + + nextAgentResult = makeAgentResult(); + + const params = { + userInput: 'Regular text query', + conversationState: undefined, + chatModel: { + provider: 'openai', + modelId: 'gpt-4o-mini', + }, + modelApiKey: 'sk-openai-from-secret', + systemPrompt: '', + temperature: 0.7, + maxTokens: 256, + memorySize: 8, + stepLimit: 4, + structuredOutputEnabled: false, + }; + + const result = (await runComponentWithRunner( + component!.runner, + (p: any, ctx: any) => + (component!.execute as any)(p, ctx, { + ToolLoopAgent: MockToolLoopAgent as unknown as ToolLoopAgentClass, + stepCountIs: stepCountIsMock as unknown as StepCountIsFn, + tool: ((definition: any) => definition) as unknown as ToolFn, + createOpenAI: openAiFactoryMock as unknown as CreateOpenAIFn, + createGoogleGenerativeAI: googleFactoryMock as unknown as CreateGoogleGenerativeAIFn, + generateObject: generateObjectMock as unknown as GenerateObjectFn, + generateText: generateTextMock as unknown as GenerateTextFn, + }), + params, + workflowContext, + )) as any; + + expect(generateObjectMock).not.toHaveBeenCalled(); + expect(toolLoopAgentConstructorMock).toHaveBeenCalled(); + expect(result.structuredOutput).toBeNull(); + expect(result.responseText).toBe('Agent final answer'); + }); + + test('throws error when auto-fix fails to parse', async () => { + const component = componentRegistry.get('core.ai.agent'); + expect(component).toBeDefined(); + + generateObjectMock.mockRejectedValue(new Error('Schema validation failed')); + generateTextMock.mockResolvedValue({ + text: 'This is not valid JSON at all', + usage: { promptTokens: 20, completionTokens: 30, totalTokens: 50 }, + }); + + const params = { + userInput: 'Generate user data', + conversationState: undefined, + chatModel: { + provider: 'openai', + modelId: 'gpt-4o-mini', + }, + modelApiKey: 'sk-openai-from-secret', + systemPrompt: '', + temperature: 0.7, + maxTokens: 256, + memorySize: 8, + stepLimit: 4, + structuredOutputEnabled: true, + schemaType: 'json-example', + jsonExample: '{"name": "example", "age": 0}', + autoFixFormat: true, + }; + + await expect( + runComponentWithRunner( + component!.runner, + (p: any, ctx: any) => + (component!.execute as any)(p, ctx, { + ToolLoopAgent: MockToolLoopAgent as unknown as ToolLoopAgentClass, + stepCountIs: stepCountIsMock as unknown as StepCountIsFn, + tool: ((definition: any) => definition) as unknown as ToolFn, + createOpenAI: openAiFactoryMock as unknown as CreateOpenAIFn, + createGoogleGenerativeAI: googleFactoryMock as unknown as CreateGoogleGenerativeAIFn, + generateObject: generateObjectMock as unknown as GenerateObjectFn, + generateText: generateTextMock as unknown as GenerateTextFn, + }), + params, + workflowContext, + ), + ).rejects.toThrow('auto-fix could not parse'); + }); + }); }); diff --git a/worker/src/components/ai/ai-agent.ts b/worker/src/components/ai/ai-agent.ts index 76a728c9..0e8effcb 100644 --- a/worker/src/components/ai/ai-agent.ts +++ b/worker/src/components/ai/ai-agent.ts @@ -4,6 +4,9 @@ import { ToolLoopAgent as ToolLoopAgentImpl, stepCountIs as stepCountIsImpl, tool as toolImpl, + generateObject as generateObjectImpl, + generateText as generateTextImpl, + jsonSchema as createJsonSchema, type Tool, } from 'ai'; import { createOpenAI as createOpenAIImpl } from '@ai-sdk/openai'; @@ -29,6 +32,8 @@ export type StepCountIsFn = typeof stepCountIsImpl; export type ToolFn = typeof toolImpl; export type CreateOpenAIFn = typeof createOpenAIImpl; export type CreateGoogleGenerativeAIFn = typeof createGoogleGenerativeAIImpl; +export type GenerateObjectFn = typeof generateObjectImpl; +export type GenerateTextFn = typeof generateTextImpl; type ModelProvider = 'openai' | 'gemini' | 'openrouter'; @@ -154,6 +159,26 @@ const inputSchema = z.object({ .max(12) .default(DEFAULT_STEP_LIMIT) .describe('Maximum sequential reasoning/tool steps before the agent stops.'), + structuredOutputEnabled: z + .boolean() + .default(false) + .describe('Enable structured JSON output that adheres to a defined schema.'), + schemaType: z + .enum(['json-example', 'json-schema']) + .default('json-example') + .describe('How to define the output schema: from a JSON example or a full JSON Schema.'), + jsonExample: z + .string() + .optional() + .describe('Example JSON object to generate schema from. All properties become required.'), + jsonSchema: z + .string() + .optional() + .describe('Full JSON Schema definition for structured output validation.'), + autoFixFormat: z + .boolean() + .default(false) + .describe('Attempt to fix malformed JSON responses from the model.'), }); type Input = z.infer; @@ -167,6 +192,7 @@ type ReasoningStep = z.infer; type Output = { responseText: string; + structuredOutput: unknown; conversationState: ConversationState; toolInvocations: ToolInvocationEntry[]; reasoningTrace: ReasoningStep[]; @@ -177,6 +203,7 @@ type Output = { const outputSchema = z.object({ responseText: z.string(), + structuredOutput: z.unknown().nullable(), conversationState: conversationStateSchema, toolInvocations: z.array(toolInvocationSchema), reasoningTrace: z.array(reasoningStepSchema), @@ -639,6 +666,119 @@ function mapStepToReasoning(step: any, index: number, sessionId: string): Reason }; } +/** + * Converts a JSON example object to a JSON Schema. + * All properties are treated as required (matching n8n behavior). + */ +function jsonExampleToJsonSchema(example: unknown): object { + if (example === null) { + return { type: 'null' }; + } + + if (Array.isArray(example)) { + const items = example.length > 0 + ? jsonExampleToJsonSchema(example[0]) + : {}; + return { type: 'array', items }; + } + + if (typeof example === 'object') { + const properties: Record = {}; + const required: string[] = []; + + for (const [key, value] of Object.entries(example as Record)) { + properties[key] = jsonExampleToJsonSchema(value); + required.push(key); + } + + return { + type: 'object', + properties, + required, + additionalProperties: false, + }; + } + + if (typeof example === 'string') return { type: 'string' }; + if (typeof example === 'number') { + return Number.isInteger(example) ? { type: 'integer' } : { type: 'number' }; + } + if (typeof example === 'boolean') return { type: 'boolean' }; + + return {}; +} + +/** + * Resolves the structured output schema from user input. + * Returns null if structured output is disabled or no valid schema provided. + */ +function resolveStructuredOutputSchema(params: { + structuredOutputEnabled?: boolean; + schemaType?: 'json-example' | 'json-schema'; + jsonExample?: string; + jsonSchema?: string; +}): object | null { + if (!params.structuredOutputEnabled) { + return null; + } + + if (params.schemaType === 'json-example' && params.jsonExample) { + try { + const example = JSON.parse(params.jsonExample); + return jsonExampleToJsonSchema(example); + } catch { + throw new Error('Invalid JSON example: unable to parse JSON.'); + } + } + + if (params.schemaType === 'json-schema' && params.jsonSchema) { + try { + return JSON.parse(params.jsonSchema); + } catch { + throw new Error('Invalid JSON Schema: unable to parse JSON.'); + } + } + + return null; +} + +/** + * Attempts to fix malformed JSON by extracting valid JSON from text. + * Handles common issues like markdown code blocks, extra text before/after JSON. + */ +function attemptJsonFix(text: string): unknown | null { + try { + return JSON.parse(text); + } catch { + // Continue to fixes + } + + let cleaned = text.replace(/```json\s*/gi, '').replace(/```\s*/g, ''); + + const objectMatch = cleaned.match(/\{[\s\S]*\}/); + const arrayMatch = cleaned.match(/\[[\s\S]*\]/); + const jsonCandidate = objectMatch?.[0] ?? arrayMatch?.[0]; + + if (jsonCandidate) { + try { + return JSON.parse(jsonCandidate); + } catch { + // Continue + } + } + + cleaned = cleaned + .trim() + .replace(/^(Here'?s?|The|Output:?|Result:?|Response:?)\s*/i, '') + .trim(); + + try { + return JSON.parse(cleaned); + } catch { + return null; + } +} + const definition: ComponentDefinition = { id: 'core.ai.agent', label: 'AI SDK Agent', @@ -707,6 +847,12 @@ Loop the Conversation State output back into the next agent invocation to keep m dataType: port.text(), description: 'Final assistant message produced by the agent.', }, + { + id: 'structuredOutput', + label: 'Structured Output', + dataType: port.json(), + description: 'Parsed JSON object when structured output is enabled. Null otherwise.', + }, { id: 'conversationState', label: 'Conversation State', @@ -782,10 +928,59 @@ Loop the Conversation State output back into the next agent invocation to keep m max: 12, description: 'Maximum reasoning/tool steps before the agent stops automatically.', }, + { + id: 'structuredOutputEnabled', + label: 'Structured Output', + type: 'boolean', + required: false, + default: false, + description: 'Enable to enforce a specific JSON output structure from the AI model.', + }, + { + id: 'schemaType', + label: 'Schema Type', + type: 'select', + required: false, + default: 'json-example', + options: [ + { label: 'Generate From JSON Example', value: 'json-example' }, + { label: 'Define Using JSON Schema', value: 'json-schema' }, + ], + description: 'Choose how to define the output structure.', + visibleWhen: { structuredOutputEnabled: true }, + }, + { + id: 'jsonExample', + label: 'JSON Example', + type: 'json', + required: false, + description: 'Provide an example JSON object. Property types and names will be used to generate the schema. All fields are treated as required.', + helpText: 'Example: { "name": "John", "age": 30, "skills": ["security", "architecture"] }', + visibleWhen: { structuredOutputEnabled: true, schemaType: 'json-example' }, + }, + { + id: 'jsonSchema', + label: 'JSON Schema', + type: 'json', + required: false, + description: 'Provide a full JSON Schema definition. Refer to json-schema.org for syntax.', + helpText: 'Example: { "type": "object", "properties": { "name": { "type": "string" } }, "required": ["name"] }', + visibleWhen: { structuredOutputEnabled: true, schemaType: 'json-schema' }, + }, + { + id: 'autoFixFormat', + label: 'Auto-Fix Format', + type: 'boolean', + required: false, + default: false, + description: 'Attempt to fix malformed JSON responses from the model.', + helpText: 'When enabled, tries to extract valid JSON from responses that contain extra text or formatting issues.', + visibleWhen: { structuredOutputEnabled: true }, + }, ], }, async execute( - params, + params, context, // Optional dependencies for testing - in production these will use the default implementations dependencies?: { @@ -794,6 +989,8 @@ Loop the Conversation State output back into the next agent invocation to keep m tool?: ToolFn; createOpenAI?: CreateOpenAIFn; createGoogleGenerativeAI?: CreateGoogleGenerativeAIFn; + generateObject?: GenerateObjectFn; + generateText?: GenerateTextFn; } ) { const { @@ -806,6 +1003,11 @@ Loop the Conversation State output back into the next agent invocation to keep m maxTokens, memorySize, stepLimit, + structuredOutputEnabled, + schemaType, + jsonExample, + jsonSchema, + autoFixFormat, } = params; const debugLog = (...args: unknown[]) => context.logger.debug(`[AIAgent Debug] ${args.join(' ')}`); @@ -966,53 +1168,138 @@ Loop the Conversation State output back into the next agent invocation to keep m stepLimit, }); - const ToolLoopAgent = dependencies?.ToolLoopAgent ?? ToolLoopAgentImpl; - const stepCountIs = dependencies?.stepCountIs ?? stepCountIsImpl; - let streamedStepCount = 0; - const agent = new ToolLoopAgent({ - id: `${sessionId}-agent`, - model, - instructions: resolvedSystemPrompt || undefined, - ...(toolsConfig ? { tools: toolsConfig } : {}), - temperature, - maxOutputTokens: maxTokens, - stopWhen: stepCountIs(stepLimit), - onStepFinish: (stepResult: unknown) => { - const mappedStep = mapStepToReasoning(stepResult, streamedStepCount, sessionId); - streamedStepCount += 1; - agentStream.emitReasoningStep(mappedStep); - }, - }); - debugLog('ToolLoopAgent instantiated', { - id: `${sessionId}-agent`, - temperature, - maxTokens, - stepLimit, - toolKeys: toolsConfig ? Object.keys(toolsConfig) : [], + // Resolve structured output schema if enabled + const structuredSchema = resolveStructuredOutputSchema({ + structuredOutputEnabled, + schemaType, + jsonExample, + jsonSchema, }); - context.logger.info( - `[AIAgent] Using ${effectiveProvider} model "${effectiveModel}" with ${availableToolsCount} connected tool(s).`, - ); - context.emitProgress({ - level: 'info', - message: 'AI agent reasoning in progress...', - data: { - agentRunId, - agentStatus: 'running', - }, - }); - debugLog('Invoking ToolLoopAgent.generate with payload', { - messages: messagesForModel, - }); + let responseText: string; + let structuredOutput: unknown = null; + let generationResult: any; + + if (structuredSchema) { + // Use generateObject for structured output mode + context.logger.info('[AIAgent] Using structured output mode with JSON Schema.'); + context.emitProgress({ + level: 'info', + message: 'AI agent generating structured output...', + data: { + agentRunId, + agentStatus: 'running', + }, + }); - const generationResult = await agent.generate({ - messages: messagesForModel as any, - }); - debugLog('Generation result', generationResult); + const generateObject = dependencies?.generateObject ?? generateObjectImpl; + const generateText = dependencies?.generateText ?? generateTextImpl; + + try { + const objectResult = await generateObject({ + model, + schema: createJsonSchema(structuredSchema), + system: resolvedSystemPrompt || undefined, + messages: messagesForModel as any, + temperature, + maxOutputTokens: maxTokens, + }); + + structuredOutput = objectResult.object; + responseText = JSON.stringify(structuredOutput, null, 2); + generationResult = { + text: responseText, + steps: [], + toolResults: [], + finishReason: 'stop', + usage: objectResult.usage, + }; + debugLog('Structured output generated successfully', structuredOutput); + } catch (error) { + // If generateObject fails and auto-fix is enabled, try text generation + fix + if (autoFixFormat) { + context.logger.warn('[AIAgent] Structured output failed, attempting auto-fix via text generation.'); + + const textResult = await generateText({ + model, + system: resolvedSystemPrompt || undefined, + messages: [ + ...messagesForModel, + { role: 'user' as const, content: `Respond with valid JSON matching this schema: ${JSON.stringify(structuredSchema)}` } + ] as any, + temperature, + maxOutputTokens: maxTokens, + }); + + const fixedOutput = attemptJsonFix(textResult.text); + if (fixedOutput !== null) { + structuredOutput = fixedOutput; + responseText = JSON.stringify(fixedOutput, null, 2); + generationResult = { + text: responseText, + steps: [], + toolResults: [], + finishReason: 'stop', + usage: textResult.usage, + }; + debugLog('Auto-fix succeeded', fixedOutput); + } else { + throw new Error(`Structured output failed and auto-fix could not parse response: ${textResult.text.slice(0, 200)}`); + } + } else { + throw error; + } + } + } else { + // Use ToolLoopAgent for standard text generation with tools + const ToolLoopAgent = dependencies?.ToolLoopAgent ?? ToolLoopAgentImpl; + const stepCountIs = dependencies?.stepCountIs ?? stepCountIsImpl; + let streamedStepCount = 0; + const agent = new ToolLoopAgent({ + id: `${sessionId}-agent`, + model, + instructions: resolvedSystemPrompt || undefined, + ...(toolsConfig ? { tools: toolsConfig } : {}), + temperature, + maxOutputTokens: maxTokens, + stopWhen: stepCountIs(stepLimit), + onStepFinish: (stepResult: unknown) => { + const mappedStep = mapStepToReasoning(stepResult, streamedStepCount, sessionId); + streamedStepCount += 1; + agentStream.emitReasoningStep(mappedStep); + }, + }); + debugLog('ToolLoopAgent instantiated', { + id: `${sessionId}-agent`, + temperature, + maxTokens, + stepLimit, + toolKeys: toolsConfig ? Object.keys(toolsConfig) : [], + }); + + context.logger.info( + `[AIAgent] Using ${effectiveProvider} model "${effectiveModel}" with ${availableToolsCount} connected tool(s).`, + ); + context.emitProgress({ + level: 'info', + message: 'AI agent reasoning in progress...', + data: { + agentRunId, + agentStatus: 'running', + }, + }); + debugLog('Invoking ToolLoopAgent.generate with payload', { + messages: messagesForModel, + }); - const responseText = - typeof generationResult.text === 'string' ? generationResult.text : String(generationResult.text ?? ''); + generationResult = await agent.generate({ + messages: messagesForModel as any, + }); + debugLog('Generation result', generationResult); + + responseText = + typeof generationResult.text === 'string' ? generationResult.text : String(generationResult.text ?? ''); + } debugLog('Response text', responseText); const currentTimestamp = new Date().toISOString(); @@ -1096,6 +1383,7 @@ Loop the Conversation State output back into the next agent invocation to keep m return { responseText, + structuredOutput, conversationState: nextState, toolInvocations: toolLogEntries, reasoningTrace,