From 7788fa8728f864732ab806e34e4df7c61b66114b Mon Sep 17 00:00:00 2001 From: Killian-Aidalinfo Date: Wed, 3 Jun 2026 22:04:46 +0200 Subject: [PATCH] feat(core): resilient structured output for providers without responseFormat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Structured-output agents on non-OpenAI providers (e.g. Scaleway) fail intermittently with AI_NoObjectGeneratedError when the model returns valid JSON but with drifted keys (casing/separator like documentType↔document_type, or semantic/translated aliases like type↔document_type). Nothing constrains the model since these providers don't enforce responseFormat. Adds two agnostic resilience layers to the structured pipeline, on by default and backward-compatible (a conformant object passes through unchanged, no extra LLM call): 1. Deterministic key normalization — recursively remaps keys onto the schema's property names, insensitive to case/separator. Fixes casing drift with zero LLM cost. 2. Error-driven repair retry — when the object still misses required keys or violates enums, re-queries the model (default 2 attempts) reinjecting the issues and the exact expected keys. Generic safety net for semantic aliases. Also fixes shouldUseStructuredPipeline, which gated on structuredOutput.type but the AI SDK's Output.object() identifies via .name === "object" — so the pipeline (and thus any resilience) was never selected for real usage. Now accepts type ?? name, consistent with getJsonSchemaFromStructuredOutput. New opt-out options on Agent config and per call: normalizeStructuredKeys, structuredOutputRepair. Validated live against Scaleway gpt-oss-120b: drifted keys converge to a zod-conformant object in one repair; disabling resilience reproduces the original failure. Co-Authored-By: Claude Opus 4.8 (1M context) --- packages/core/package.json | 2 +- packages/core/src/agents/index.ts | 33 +- .../core/src/agents/structurePipeline.test.ts | 133 ++++++ packages/core/src/agents/structurePipeline.ts | 199 +++++++-- .../agents/structuredOutputResilience.test.ts | 382 ++++++++++++++++++ .../src/agents/structuredOutputResilience.ts | 337 +++++++++++++++ packages/core/src/agents/types.ts | 3 +- 7 files changed, 1062 insertions(+), 27 deletions(-) create mode 100644 packages/core/src/agents/structuredOutputResilience.test.ts create mode 100644 packages/core/src/agents/structuredOutputResilience.ts diff --git a/packages/core/package.json b/packages/core/package.json index fb14cb9..0cf581d 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -1,6 +1,6 @@ { "name": "@ai_kit/core", - "version": "1.8.0", + "version": "1.9.0", "description": "", "type": "module", "main": "./dist/index.js", diff --git a/packages/core/src/agents/index.ts b/packages/core/src/agents/index.ts index 9b73147..9d0fb62 100644 --- a/packages/core/src/agents/index.ts +++ b/packages/core/src/agents/index.ts @@ -43,6 +43,10 @@ import { } from "./toolLoop.js"; import { buildToonSystemPrompt, parseToonStructuredOutput } from "./toon.js"; import { getJsonSchemaFromStructuredOutput } from "./structuredOutputSchema.js"; +import { + resolveResilienceConfig, + type StructuredOutputResilienceOptions, +} from "./structuredOutputResilience.js"; const Output = BaseOutput as typeof BaseOutput & { object: >(options: { @@ -61,8 +65,9 @@ export type { AgentStructuredOutput, } from "./types.js"; export { DEFAULT_MAX_STEP_TOOLS } from "./toolLoop.js"; +export type { StructuredOutputResilienceOptions } from "./structuredOutputResilience.js"; -export interface AgentConfig { +export interface AgentConfig extends StructuredOutputResilienceOptions { name: string; instructions?: string; model: LanguageModel; @@ -84,6 +89,8 @@ export class Agent { private loopToolsEnabled: boolean; private maxStepTools: number; private toonEnabled: boolean; + private normalizeStructuredKeysDefault?: boolean; + private structuredOutputRepairDefault?: boolean | { maxAttempts?: number }; readonly memory?: Memory; constructor({ @@ -96,6 +103,8 @@ export class Agent { maxStepTools, toon, memory, + normalizeStructuredKeys, + structuredOutputRepair, }: AgentConfig) { this.name = name; this.instructions = instructions; @@ -107,11 +116,22 @@ export class Agent { this.loopToolsEnabled = loopTools ?? false; this.maxStepTools = maxStepTools ?? DEFAULT_MAX_STEP_TOOLS; this.toonEnabled = toon ?? false; + this.normalizeStructuredKeysDefault = normalizeStructuredKeys; + this.structuredOutputRepairDefault = structuredOutputRepair; if (memory) { this.memory = new Memory(memory); } } + private resolveResilience(options: StructuredOutputResilienceOptions) { + return resolveResilienceConfig({ + normalizeStructuredKeys: + options.normalizeStructuredKeys ?? this.normalizeStructuredKeysDefault, + structuredOutputRepair: + options.structuredOutputRepair ?? this.structuredOutputRepairDefault, + }); + } + withTelemetry(enabled: boolean = true) { this.telemetryEnabled = enabled; return this; @@ -234,6 +254,7 @@ export class Agent { telemetryEnabled: this.telemetryEnabled, telemetryDefaults: this.telemetryDefaults, agentName: this.name, + resilienceConfig: this.resolveResilience(options), }), ); } @@ -261,6 +282,7 @@ export class Agent { telemetryDefaults: this.telemetryDefaults, agentName: this.name, loopToolsEnabled: loopSettings.enabled, + resilienceConfig: this.resolveResilience(options), }); return result; @@ -275,6 +297,8 @@ export class Agent { structuredOutput: _structured, runtime: _runtime, toon: _toon, + normalizeStructuredKeys: _normalizeStructuredKeys, + structuredOutputRepair: _structuredOutputRepair, ...rest } = options; const { @@ -355,6 +379,8 @@ export class Agent { structuredOutput: _structured, runtime: _runtime, toon: _toon, + normalizeStructuredKeys: _normalizeStructuredKeys, + structuredOutputRepair: _structuredOutputRepair, ...rest } = options; const { @@ -501,6 +527,7 @@ export class Agent { telemetryDefaults: this.telemetryDefaults, agentName: this.name, loopToolsEnabled: loopSettings.enabled, + resilienceConfig: this.resolveResilience(options), }); }, }); @@ -515,6 +542,8 @@ export class Agent { structuredOutput: _structured, runtime: _runtime, toon: _toon, + normalizeStructuredKeys: _normalizeStructuredKeys, + structuredOutputRepair: _structuredOutputRepair, ...rest } = options; const { @@ -605,6 +634,8 @@ export class Agent { structuredOutput: _structured, runtime: _runtime, toon: _toon, + normalizeStructuredKeys: _normalizeStructuredKeys, + structuredOutputRepair: _structuredOutputRepair, ...rest } = options; const { diff --git a/packages/core/src/agents/structurePipeline.test.ts b/packages/core/src/agents/structurePipeline.test.ts index 8ab949a..c01f055 100644 --- a/packages/core/src/agents/structurePipeline.test.ts +++ b/packages/core/src/agents/structurePipeline.test.ts @@ -33,10 +33,12 @@ vi.mock("./structuredOutputSchema.js", () => ({ })); import { + generateWithDirectStructuredObject, generateWithStructuredPipeline, shouldUseStructuredPipeline, streamWithStructuredPipeline, } from "./structurePipeline.js"; +import { resolveResilienceConfig } from "./structuredOutputResilience.js"; class MockStreamResult { text: Promise; @@ -81,6 +83,25 @@ describe("structurePipeline", () => { outputObjectMock.mockClear(); }); + it("shouldUseStructuredPipeline reconnaît Output.object via .name (sans .type)", () => { + // The AI SDK's Output.object() result carries { name: "object" } and no + // `type`. The gate must treat it as a structured object all the same. + const outputObjectShape = { name: "object" } as any; + + expect( + shouldUseStructuredPipeline({ provider: "scaleway.chat" } as any, {}, outputObjectShape), + ).toBe(true); + expect( + shouldUseStructuredPipeline({ provider: "anthropic" } as any, {}, outputObjectShape), + ).toBe(true); + expect( + shouldUseStructuredPipeline({ provider: "openai" } as any, {}, outputObjectShape), + ).toBe(false); + expect( + shouldUseStructuredPipeline({ provider: "anthropic" } as any, {}, { name: "text" } as any), + ).toBe(false); + }); + it("shouldUseStructuredPipeline respecte toon, type et provider", () => { const structuredOutput = { type: "object" } as any; @@ -205,4 +226,116 @@ describe("structurePipeline", () => { ); expect(() => (result as any).experimental_output).toThrow("structured stream failed"); }); + + it("normalise une clé dérivée de la sortie structurée sans reprise", async () => { + generateTextMock.mockResolvedValueOnce({ text: "Réponse libre" }); + generateTextMock.mockResolvedValueOnce({ output: { Summary: "OK" } }); + + const result = await generateWithStructuredPipeline({ + model: { provider: "anthropic" } as any, + structuredOutput: { type: "object" } as any, + options: { prompt: "Fais un résumé" } as any, + telemetryEnabled: false, + loopToolsEnabled: false, + }); + + expect((result as any).experimental_output).toEqual({ summary: "OK" }); + expect(generateTextMock).toHaveBeenCalledTimes(2); + }); + + it("relance le modèle pour réparer un alias sémantique (clé manquante)", async () => { + generateTextMock.mockResolvedValueOnce({ text: "Réponse libre" }); + generateTextMock.mockResolvedValueOnce({ output: { resume: "mauvaise clé" } }); + generateTextMock.mockResolvedValueOnce({ output: { summary: "réparé" } }); + + const result = await generateWithStructuredPipeline({ + model: { provider: "anthropic" } as any, + structuredOutput: { type: "object" } as any, + options: { prompt: "Fais un résumé" } as any, + telemetryEnabled: false, + loopToolsEnabled: false, + }); + + expect((result as any).experimental_output).toEqual({ summary: "réparé" }); + expect(generateTextMock).toHaveBeenCalledTimes(3); + + const repairPayload = generateTextMock.mock.calls[2]?.[0]; + const lastMessage = repairPayload.messages[repairPayload.messages.length - 1]; + expect(lastMessage.role).toBe("user"); + expect(lastMessage.content).toContain("summary"); + }); + + it("ne relance pas quand la réparation est désactivée", async () => { + generateTextMock.mockResolvedValueOnce({ text: "Réponse libre" }); + generateTextMock.mockResolvedValueOnce({ output: { resume: "mauvaise clé" } }); + + const result = await generateWithStructuredPipeline({ + model: { provider: "anthropic" } as any, + structuredOutput: { type: "object" } as any, + options: { prompt: "Fais un résumé" } as any, + telemetryEnabled: false, + loopToolsEnabled: false, + resilienceConfig: resolveResilienceConfig({ structuredOutputRepair: false }), + }); + + expect((result as any).experimental_output).toEqual({ resume: "mauvaise clé" }); + expect(generateTextMock).toHaveBeenCalledTimes(2); + }); + + it("normalise la clé dérivée du dernier objet en streaming", async () => { + const baseStreamResult = new MockStreamResult("Texte stream") as any; + const objectStream = { + output: Promise.resolve({ Summary: "stream-ok" }), + partialOutputStream: asAsyncIterable([{ summary: "partiel" }]), + }; + + streamTextMock.mockReturnValueOnce(baseStreamResult); + streamTextMock.mockResolvedValueOnce(objectStream); + + const result = await streamWithStructuredPipeline({ + model: { provider: "anthropic" } as any, + structuredOutput: { type: "object" } as any, + options: { prompt: "stream" } as any, + telemetryEnabled: false, + loopToolsEnabled: false, + }); + + await Promise.resolve(); + await Promise.resolve(); + expect((result as any).experimental_output).toEqual({ summary: "stream-ok" }); + }); + + it("normalise la clé dérivée sur le chemin direct (sans outils)", async () => { + generateTextMock.mockResolvedValueOnce({ output: { Summary: "OK" } }); + + const result = await generateWithDirectStructuredObject({ + model: { provider: "anthropic" } as any, + structuredOutput: { type: "object" } as any, + options: { prompt: "Fais un résumé" } as any, + telemetryEnabled: false, + }); + + expect((result as any).experimental_output).toEqual({ summary: "OK" }); + expect(generateTextMock).toHaveBeenCalledTimes(1); + }); + + it("répare un alias sémantique sur le chemin direct", async () => { + generateTextMock.mockResolvedValueOnce({ output: { resume: "mauvaise clé" } }); + generateTextMock.mockResolvedValueOnce({ output: { summary: "réparé" } }); + + const result = await generateWithDirectStructuredObject({ + model: { provider: "anthropic" } as any, + structuredOutput: { type: "object" } as any, + options: { prompt: "Fais un résumé" } as any, + telemetryEnabled: false, + }); + + expect((result as any).experimental_output).toEqual({ summary: "réparé" }); + expect(generateTextMock).toHaveBeenCalledTimes(2); + + const repairPayload = generateTextMock.mock.calls[1]?.[0]; + const lastMessage = repairPayload.messages[repairPayload.messages.length - 1]; + expect(lastMessage.role).toBe("user"); + expect(lastMessage.content).toContain("summary"); + }); }); diff --git a/packages/core/src/agents/structurePipeline.ts b/packages/core/src/agents/structurePipeline.ts index 0b594fc..ba90d69 100644 --- a/packages/core/src/agents/structurePipeline.ts +++ b/packages/core/src/agents/structurePipeline.ts @@ -27,6 +27,12 @@ import { } from "./types.js"; import { setExperimentalOutput } from "./experimentalOutput.js"; import { getJsonSchemaFromStructuredOutput } from "./structuredOutputSchema.js"; +import { + normalizeKeysToSchema, + resolveResilienceConfig, + resolveResilientObject, + type ResilienceConfig, +} from "./structuredOutputResilience.js"; const OPENAI_PROVIDER_ID = "openai"; @@ -46,6 +52,7 @@ interface StructuredGeneratePipelineParams< telemetryDefaults?: AgentTelemetryOverrides; agentName?: string; loopToolsEnabled: boolean; + resilienceConfig?: ResilienceConfig; } interface StructuredStreamPipelineParams< @@ -62,6 +69,7 @@ interface StructuredStreamPipelineParams< telemetryDefaults?: AgentTelemetryOverrides; agentName?: string; loopToolsEnabled: boolean; + resilienceConfig?: ResilienceConfig; } interface StructuredDirectGenerateParams< @@ -76,6 +84,7 @@ interface StructuredDirectGenerateParams< telemetryEnabled: boolean; telemetryDefaults?: AgentTelemetryOverrides; agentName?: string; + resilienceConfig?: ResilienceConfig; } export function shouldUseStructuredPipeline( @@ -88,7 +97,11 @@ export function shouldUseStructuredPipeline( return false; } - if (!structuredOutput || structuredOutput.type !== "object") { + // The AI SDK's `Output.object()` result identifies itself via `name` + // ("object"), while hand-built structured outputs may use `type`. Accept + // either, consistent with getJsonSchemaFromStructuredOutput. + const kind = structuredOutput?.type ?? structuredOutput?.name; + if (!structuredOutput || kind !== "object") { return false; } @@ -112,6 +125,7 @@ export async function generateWithStructuredPipeline< telemetryDefaults, agentName, loopToolsEnabled, + resilienceConfig, } = params; const originalPrompt = "prompt" in options ? options.prompt : undefined; @@ -128,26 +142,35 @@ export async function generateWithStructuredPipeline< loopToolsEnabled, }); - const schema = jsonSchema( - await getJsonSchemaFromStructuredOutput(structuredOutput), + const jsonSchemaObject = await getJsonSchemaFromStructuredOutput( + structuredOutput, ); + const schema = jsonSchema(jsonSchemaObject); const structuringMessages = buildStructuringMessages({ text: textResult.text, originalPrompt, originalMessages, }); + const objectCallSettings = extractObjectCallSettings( + options as unknown as Partial, + ); - const objectResult = await generateText({ - ...extractObjectCallSettings( - options as unknown as Partial, - ), - model, - system, - messages: structuringMessages, - output: Output.object({ schema }), + const runStructuredObjectCall = (messages: ModelMessages) => + generateStructuredObject({ model, system, schema, messages, objectCallSettings }); + + const initialObject = await runStructuredObjectCall(structuringMessages); + + const resilient = await resolveResilientObject({ + initialObject, + schema: jsonSchemaObject, + config: resilienceConfig ?? resolveResilienceConfig({}), + requery: ({ instruction, previousObject }) => + runStructuredObjectCall( + appendRepairTurn(structuringMessages, instruction, previousObject), + ), }); - setExperimentalOutput(textResult, objectResult.output as OUTPUT); + setExperimentalOutput(textResult, resilient as OUTPUT); return textResult; } @@ -167,11 +190,13 @@ export async function generateWithDirectStructuredObject< telemetryEnabled, telemetryDefaults, agentName, + resilienceConfig, } = params; - const schema = jsonSchema( - await getJsonSchemaFromStructuredOutput(structuredOutput), + const jsonSchemaObject = await getJsonSchemaFromStructuredOutput( + structuredOutput, ); + const schema = jsonSchema(jsonSchemaObject); const textResult = await callGenerateTextDirect({ model, @@ -183,7 +208,26 @@ export async function generateWithDirectStructuredObject< schema, }); - setExperimentalOutput(textResult, textResult.output as OUTPUT); + const objectCallSettings = extractObjectCallSettings( + options as unknown as Partial, + ); + const baseMessages = deriveBaseMessages(options); + + const resilient = await resolveResilientObject({ + initialObject: textResult.output as unknown, + schema: jsonSchemaObject, + config: resilienceConfig ?? resolveResilienceConfig({}), + requery: ({ instruction, previousObject }) => + generateStructuredObject({ + model, + system, + schema, + messages: appendRepairTurn(baseMessages, instruction, previousObject), + objectCallSettings, + }), + }); + + setExperimentalOutput(textResult, resilient as OUTPUT); return textResult; } @@ -205,6 +249,7 @@ export async function streamWithStructuredPipeline< telemetryDefaults, agentName, loopToolsEnabled, + resilienceConfig, } = params; const originalPrompt = "prompt" in options ? options.prompt : undefined; @@ -221,9 +266,11 @@ export async function streamWithStructuredPipeline< loopToolsEnabled, }); - const schema = jsonSchema( - await getJsonSchemaFromStructuredOutput(structuredOutput), + const jsonSchemaObject = await getJsonSchemaFromStructuredOutput( + structuredOutput, ); + const schema = jsonSchema(jsonSchemaObject); + const config = resilienceConfig ?? resolveResilienceConfig({}); let pipelineError: unknown; const objectStreamPromise = (async () => { @@ -246,7 +293,10 @@ export async function streamWithStructuredPipeline< try { const finalObject = await objectStream.output; - setExperimentalOutput(streamResult, finalObject as OUTPUT); + const normalized = config.normalizeKeys + ? normalizeKeysToSchema(finalObject, jsonSchemaObject) + : finalObject; + setExperimentalOutput(streamResult, normalized as OUTPUT); } catch (error) { pipelineError = error; throw error; @@ -283,6 +333,69 @@ function getProvider(model: LanguageModel): string | undefined { return typeof model === "string" ? undefined : model.provider; } +/** + * Runs a single structuring object generation and returns the produced object + * (unvalidated — the schema passed to `Output.object` here is a bare JSON + * schema, so resilience layers downstream own validation). + */ +async function generateStructuredObject({ + model, + system, + schema, + messages, + objectCallSettings, +}: { + model: LanguageModel; + system?: string; + schema: ReturnType; + messages: ModelMessages; + objectCallSettings: ReturnType; +}): Promise { + const objectResult = await generateText({ + ...objectCallSettings, + model, + system, + messages, + output: Output.object({ schema }), + }); + + return objectResult.output as unknown; +} + +/** Appends the repair re-query (issues + expected keys + prior JSON) as a user turn. */ +function appendRepairTurn( + messages: ModelMessages, + instruction: string, + previousObject: unknown, +): ModelMessages { + return [ + ...messages, + { + role: "user", + content: `${instruction}\n\nPrevious response:\n${JSON.stringify( + previousObject, + )}`, + }, + ]; +} + +/** Conversation to re-send when repairing a direct (single-pass) structured call. */ +function deriveBaseMessages(options: { + prompt?: GenerateTextParams["prompt"]; + messages?: GenerateTextParams["messages"]; +}): ModelMessages { + if (options.messages !== undefined) { + return options.messages as ModelMessages; + } + + const content = flattenPrompt(options.prompt); + if (content !== undefined) { + return [{ role: "user", content }]; + } + + return []; +} + function buildStructuringMessages({ text, originalPrompt, @@ -421,12 +534,19 @@ async function callGenerateText< experimental_telemetry, loopTools: _loopTools, maxStepTools: _maxStepTools, + normalizeStructuredKeys: _normalizeStructuredKeys, + structuredOutputRepair: _structuredOutputRepair, ...restWithoutContext } = rest as { experimental_context?: unknown; telemetry?: AgentTelemetryOverrides; experimental_telemetry?: GenerateTextParams["experimental_telemetry"]; - } & typeof rest & { loopTools?: unknown; maxStepTools?: unknown }; + } & typeof rest & { + loopTools?: unknown; + maxStepTools?: unknown; + normalizeStructuredKeys?: unknown; + structuredOutputRepair?: unknown; + }; const toolSet = toToolSet(tools); const payload = { ...restWithoutContext, @@ -484,12 +604,19 @@ async function callGenerateText< experimental_telemetry, loopTools: _loopTools, maxStepTools: _maxStepTools, + normalizeStructuredKeys: _normalizeStructuredKeys, + structuredOutputRepair: _structuredOutputRepair, ...restWithoutContext } = rest as { experimental_context?: unknown; telemetry?: AgentTelemetryOverrides; experimental_telemetry?: GenerateTextParams["experimental_telemetry"]; - } & typeof rest & { loopTools?: unknown; maxStepTools?: unknown }; + } & typeof rest & { + loopTools?: unknown; + maxStepTools?: unknown; + normalizeStructuredKeys?: unknown; + structuredOutputRepair?: unknown; + }; const toolSet = toToolSet(tools); const payload = { ...restWithoutContext, @@ -569,12 +696,17 @@ async function callGenerateTextDirect< experimental_context, telemetry: telemetryOverrides, experimental_telemetry, + normalizeStructuredKeys: _normalizeStructuredKeys, + structuredOutputRepair: _structuredOutputRepair, ...restWithoutContext } = rest as { experimental_context?: unknown; telemetry?: AgentTelemetryOverrides; experimental_telemetry?: GenerateTextParams["experimental_telemetry"]; - } & typeof rest; + } & typeof rest & { + normalizeStructuredKeys?: unknown; + structuredOutputRepair?: unknown; + }; const payload = { ...restWithoutContext, @@ -627,12 +759,17 @@ async function callGenerateTextDirect< experimental_context, telemetry: telemetryOverrides, experimental_telemetry, + normalizeStructuredKeys: _normalizeStructuredKeys, + structuredOutputRepair: _structuredOutputRepair, ...restWithoutContext } = rest as { experimental_context?: unknown; telemetry?: AgentTelemetryOverrides; experimental_telemetry?: GenerateTextParams["experimental_telemetry"]; - } & typeof rest; + } & typeof rest & { + normalizeStructuredKeys?: unknown; + structuredOutputRepair?: unknown; + }; const payload = { ...restWithoutContext, @@ -713,12 +850,19 @@ function callStreamText< experimental_telemetry, loopTools: _loopTools, maxStepTools: _maxStepTools, + normalizeStructuredKeys: _normalizeStructuredKeys, + structuredOutputRepair: _structuredOutputRepair, ...restWithoutContext } = rest as { experimental_context?: unknown; telemetry?: AgentTelemetryOverrides; experimental_telemetry?: StreamTextParams["experimental_telemetry"]; - } & typeof rest & { loopTools?: unknown; maxStepTools?: unknown }; + } & typeof rest & { + loopTools?: unknown; + maxStepTools?: unknown; + normalizeStructuredKeys?: unknown; + structuredOutputRepair?: unknown; + }; const toolSet = toToolSet(tools); const payload = { ...restWithoutContext, @@ -775,12 +919,19 @@ function callStreamText< experimental_telemetry, loopTools: _loopTools, maxStepTools: _maxStepTools, + normalizeStructuredKeys: _normalizeStructuredKeys, + structuredOutputRepair: _structuredOutputRepair, ...restWithoutContext } = rest as { experimental_context?: unknown; telemetry?: AgentTelemetryOverrides; experimental_telemetry?: StreamTextParams["experimental_telemetry"]; - } & typeof rest & { loopTools?: unknown; maxStepTools?: unknown }; + } & typeof rest & { + loopTools?: unknown; + maxStepTools?: unknown; + normalizeStructuredKeys?: unknown; + structuredOutputRepair?: unknown; + }; const toolSet = toToolSet(tools); const payload = { ...restWithoutContext, diff --git a/packages/core/src/agents/structuredOutputResilience.test.ts b/packages/core/src/agents/structuredOutputResilience.test.ts new file mode 100644 index 0000000..962b904 --- /dev/null +++ b/packages/core/src/agents/structuredOutputResilience.test.ts @@ -0,0 +1,382 @@ +import { describe, expect, it, vi } from "vitest"; + +import { + buildRepairInstruction, + collectSchemaIssues, + listExpectedKeys, + normalizeKeysToSchema, + resolveResilienceConfig, + resolveResilientObject, +} from "./structuredOutputResilience.js"; + +describe("normalizeKeysToSchema", () => { + it("remaps a camelCase key to the schema's snake_case property", () => { + const schema = { + type: "object", + properties: { document_type: { type: "string" } }, + required: ["document_type"], + }; + + expect(normalizeKeysToSchema({ documentType: "Bilan" }, schema)).toEqual({ + document_type: "Bilan", + }); + }); + + it("remaps kebab-case and PascalCase variants to the schema property", () => { + const schema = { + type: "object", + properties: { document_type: { type: "string" } }, + }; + + expect(normalizeKeysToSchema({ "document-type": "a" }, schema)).toEqual({ + document_type: "a", + }); + expect(normalizeKeysToSchema({ DocumentType: "b" }, schema)).toEqual({ + document_type: "b", + }); + }); + + it("is a no-op when keys already match the schema", () => { + const schema = { + type: "object", + properties: { document_type: { type: "string" } }, + }; + const value = { document_type: "ok" }; + + expect(normalizeKeysToSchema(value, schema)).toEqual({ document_type: "ok" }); + }); + + it("preserves keys that have no matching schema property", () => { + const schema = { + type: "object", + properties: { document_type: { type: "string" } }, + }; + + expect( + normalizeKeysToSchema({ documentType: "a", extra: 1 }, schema), + ).toEqual({ document_type: "a", extra: 1 }); + }); + + it("does not overwrite an exact match with a drifted alias", () => { + const schema = { + type: "object", + properties: { document_type: { type: "string" } }, + }; + + expect( + normalizeKeysToSchema( + { document_type: "snake", documentType: "camel" }, + schema, + ), + ).toEqual({ document_type: "snake" }); + }); + + it("recurses into nested object properties", () => { + const schema = { + type: "object", + properties: { + patient_info: { + type: "object", + properties: { first_name: { type: "string" } }, + }, + }, + }; + + expect( + normalizeKeysToSchema({ patientInfo: { firstName: "Ada" } }, schema), + ).toEqual({ patient_info: { first_name: "Ada" } }); + }); + + it("recurses into arrays of objects via items schema", () => { + const schema = { + type: "object", + properties: { + questions: { + type: "array", + items: { + type: "object", + properties: { question_id: { type: "string" } }, + }, + }, + }, + }; + + expect( + normalizeKeysToSchema( + { questions: [{ questionId: "q1" }, { questionId: "q2" }] }, + schema, + ), + ).toEqual({ questions: [{ question_id: "q1" }, { question_id: "q2" }] }); + }); + + it("returns primitives and unschemaed values unchanged", () => { + expect(normalizeKeysToSchema("hello", { type: "string" })).toBe("hello"); + expect(normalizeKeysToSchema(42, undefined)).toBe(42); + expect(normalizeKeysToSchema({ a: 1 }, undefined)).toEqual({ a: 1 }); + }); +}); + +describe("collectSchemaIssues", () => { + it("reports no issues for a conformant object", () => { + const schema = { + type: "object", + properties: { questionId: { type: "string" } }, + required: ["questionId"], + }; + + expect(collectSchemaIssues({ questionId: "q1" }, schema)).toEqual([]); + }); + + it("flags a missing required property", () => { + const schema = { + type: "object", + properties: { questionId: { type: "string" } }, + required: ["questionId"], + }; + + const issues = collectSchemaIssues({ id: "q1" }, schema); + expect(issues).toHaveLength(1); + expect(issues[0]).toContain("questionId"); + }); + + it("flags a missing required property nested in an object with its path", () => { + const schema = { + type: "object", + properties: { + patient: { + type: "object", + properties: { lastName: { type: "string" } }, + required: ["lastName"], + }, + }, + required: ["patient"], + }; + + const issues = collectSchemaIssues({ patient: { firstName: "Ada" } }, schema); + expect(issues).toHaveLength(1); + expect(issues[0]).toContain("patient.lastName"); + }); + + it("flags an enum violation", () => { + const schema = { + type: "object", + properties: { status: { type: "string", enum: ["draft", "final"] } }, + required: ["status"], + }; + + const issues = collectSchemaIssues({ status: "done" }, schema); + expect(issues).toHaveLength(1); + expect(issues[0]).toContain("status"); + expect(issues[0]).toContain("draft"); + }); + + it("does not flag absent optional nested objects", () => { + const schema = { + type: "object", + properties: { + meta: { + type: "object", + properties: { note: { type: "string" } }, + required: ["note"], + }, + }, + }; + + expect(collectSchemaIssues({}, schema)).toEqual([]); + }); + + it("flags a missing required property inside array items", () => { + const schema = { + type: "object", + properties: { + questions: { + type: "array", + items: { + type: "object", + properties: { questionId: { type: "string" } }, + required: ["questionId"], + }, + }, + }, + required: ["questions"], + }; + + const issues = collectSchemaIssues( + { questions: [{ questionId: "q1" }, { id: "q2" }] }, + schema, + ); + expect(issues).toHaveLength(1); + expect(issues[0]).toContain("questions.1.questionId"); + }); +}); + +describe("listExpectedKeys", () => { + it("returns the top-level property names in declaration order", () => { + const schema = { + type: "object", + properties: { + questionId: { type: "string" }, + answer: { type: "string" }, + }, + }; + + expect(listExpectedKeys(schema)).toEqual(["questionId", "answer"]); + }); + + it("returns an empty list when the schema has no properties", () => { + expect(listExpectedKeys({ type: "string" })).toEqual([]); + expect(listExpectedKeys(undefined)).toEqual([]); + }); +}); + +describe("buildRepairInstruction", () => { + it("includes the schema issues and the exact expected keys", () => { + const instruction = buildRepairInstruction({ + issues: ["missing required property 'questionId'"], + expectedKeys: ["questionId", "answer"], + }); + + expect(instruction).toContain("missing required property 'questionId'"); + expect(instruction).toContain("questionId"); + expect(instruction).toContain("answer"); + }); + + it("instructs the model to return only corrected JSON", () => { + const instruction = buildRepairInstruction({ + issues: ["missing required property 'questionId'"], + expectedKeys: ["questionId"], + }); + + expect(instruction.toLowerCase()).toContain("json"); + }); +}); + +describe("resolveResilienceConfig", () => { + it("defaults to normalization on and repair with 2 attempts", () => { + expect(resolveResilienceConfig({})).toEqual({ + normalizeKeys: true, + repair: { enabled: true, maxAttempts: 2 }, + }); + }); + + it("disables repair when structuredOutputRepair is false", () => { + expect( + resolveResilienceConfig({ structuredOutputRepair: false }).repair, + ).toEqual({ enabled: false, maxAttempts: 0 }); + }); + + it("honours a custom maxAttempts", () => { + expect( + resolveResilienceConfig({ structuredOutputRepair: { maxAttempts: 5 } }) + .repair, + ).toEqual({ enabled: true, maxAttempts: 5 }); + }); + + it("can disable key normalization", () => { + expect( + resolveResilienceConfig({ normalizeStructuredKeys: false }).normalizeKeys, + ).toBe(false); + }); +}); + +describe("resolveResilientObject", () => { + const schema = { + type: "object", + properties: { document_type: { type: "string" } }, + required: ["document_type"], + }; + const aliasSchema = { + type: "object", + properties: { questionId: { type: "string" } }, + required: ["questionId"], + }; + const config = resolveResilienceConfig({}); + + it("fixes a casing drift via normalization without any re-query", async () => { + const requery = vi.fn(); + const result = await resolveResilientObject({ + initialObject: { documentType: "Bilan" }, + schema, + config, + requery, + }); + + expect(result).toEqual({ document_type: "Bilan" }); + expect(requery).not.toHaveBeenCalled(); + }); + + it("leaves a conformant object untouched and never re-queries", async () => { + const requery = vi.fn(); + const result = await resolveResilientObject({ + initialObject: { document_type: "ok" }, + schema, + config, + requery, + }); + + expect(result).toEqual({ document_type: "ok" }); + expect(requery).not.toHaveBeenCalled(); + }); + + it("repairs a semantic alias with a single re-query", async () => { + const requery = vi.fn( + async (_params: { instruction: string; previousObject: unknown }) => ({ + questionId: "q1", + }), + ); + const result = await resolveResilientObject({ + initialObject: { id: "q1" }, + schema: aliasSchema, + config, + requery, + }); + + expect(result).toEqual({ questionId: "q1" }); + expect(requery).toHaveBeenCalledTimes(1); + const instruction = requery.mock.calls[0]?.[0]?.instruction as string; + expect(instruction).toContain("questionId"); + }); + + it("stops after maxAttempts and returns the best-effort object", async () => { + const requery = vi.fn(async () => ({ id: "still-wrong" })); + const result = await resolveResilientObject({ + initialObject: { id: "q1" }, + schema: aliasSchema, + config: resolveResilienceConfig({ + structuredOutputRepair: { maxAttempts: 3 }, + }), + requery, + }); + + expect(requery).toHaveBeenCalledTimes(3); + expect(result).toEqual({ id: "still-wrong" }); + }); + + it("does not re-query when repair is disabled", async () => { + const requery = vi.fn(); + const result = await resolveResilientObject({ + initialObject: { id: "q1" }, + schema: aliasSchema, + config: resolveResilienceConfig({ structuredOutputRepair: false }), + requery, + }); + + expect(result).toEqual({ id: "q1" }); + expect(requery).not.toHaveBeenCalled(); + }); + + it("delivers the best-effort object if a re-query throws", async () => { + const requery = vi.fn(async () => { + throw new Error("network"); + }); + const result = await resolveResilientObject({ + initialObject: { id: "q1" }, + schema: aliasSchema, + config, + requery, + }); + + expect(result).toEqual({ id: "q1" }); + expect(requery).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/core/src/agents/structuredOutputResilience.ts b/packages/core/src/agents/structuredOutputResilience.ts new file mode 100644 index 0000000..3c24ea8 --- /dev/null +++ b/packages/core/src/agents/structuredOutputResilience.ts @@ -0,0 +1,337 @@ +import type { JSONSchema7 } from "ai"; + +/** A JSON Schema node, or the boolean shorthand JSON Schema allows. */ +type JSONSchema7Definition = JSONSchema7 | boolean; + +/** + * Resilience helpers for structured output produced by providers that do not + * enforce a JSON schema (everything except OpenAI in this codebase). + * + * Models frequently return the right shape but with a drifted key name — a + * different casing/separator (`documentType` vs `document_type`) or a semantic + * alias (`id` vs `questionId`). These helpers make the structuring pass tolerant + * to that drift without hard-coding any application-specific key. + */ + +type JsonRecord = Record; + +/** Per-call configuration of structured-output resilience. */ +export interface StructuredOutputResilienceOptions { + /** Layer 1 — deterministic key normalization. Default: `true`. */ + normalizeStructuredKeys?: boolean; + /** + * Layer 2 — error-driven repair retries. `true` (default) enables repair with + * `maxAttempts: 2`; `false` disables it; the object form tunes the attempts. + */ + structuredOutputRepair?: boolean | { maxAttempts?: number }; +} + +export interface ResilienceConfig { + normalizeKeys: boolean; + repair: { enabled: boolean; maxAttempts: number }; +} + +const DEFAULT_REPAIR_ATTEMPTS = 2; + +/** Resolves user-facing options into a fully-defaulted resilience config. */ +export function resolveResilienceConfig( + options: StructuredOutputResilienceOptions, +): ResilienceConfig { + const normalizeKeys = options.normalizeStructuredKeys ?? true; + const repairOption = options.structuredOutputRepair; + + if (repairOption === false) { + return { normalizeKeys, repair: { enabled: false, maxAttempts: 0 } }; + } + + if (repairOption && typeof repairOption === "object") { + return { + normalizeKeys, + repair: { + enabled: true, + maxAttempts: repairOption.maxAttempts ?? DEFAULT_REPAIR_ATTEMPTS, + }, + }; + } + + return { + normalizeKeys, + repair: { enabled: true, maxAttempts: DEFAULT_REPAIR_ATTEMPTS }, + }; +} + +function isPlainObject(value: unknown): value is JsonRecord { + return ( + typeof value === "object" && + value !== null && + !Array.isArray(value) + ); +} + +function asObjectSchema( + schema: JSONSchema7Definition | undefined, +): JSONSchema7 | undefined { + return typeof schema === "object" && schema !== null ? schema : undefined; +} + +/** + * Canonical comparison form of a property name: lowercased with every + * non-alphanumeric character (separators like `_`, `-`, spaces) stripped. So + * `document_type`, `documentType`, `document-type` and `Document Type` all + * collapse to `documenttype`. + */ +function canonicalizeKey(key: string): string { + return key.toLowerCase().replace(/[^a-z0-9]/g, ""); +} + +/** + * Recursively remaps the keys of `value` onto the property names declared by + * `schema`, matching insensitively to case and separator. Keys that already + * match exactly, and keys with no schema counterpart, are left untouched. An + * exact match always wins over a drifted alias (the alias is dropped). + */ +export function normalizeKeysToSchema( + value: unknown, + schema: JSONSchema7Definition | undefined, +): unknown { + const objectSchema = asObjectSchema(schema); + + if (Array.isArray(value)) { + const itemSchema = objectSchema?.items; + const resolvedItemSchema = Array.isArray(itemSchema) + ? undefined + : itemSchema; + return value.map((item) => + normalizeKeysToSchema(item, resolvedItemSchema), + ); + } + + if (!isPlainObject(value) || !objectSchema?.properties) { + return value; + } + + const properties = objectSchema.properties; + const canonicalToProp = new Map(); + for (const propName of Object.keys(properties)) { + canonicalToProp.set(canonicalizeKey(propName), propName); + } + + const result: JsonRecord = {}; + for (const [key, child] of Object.entries(value)) { + const isExact = Object.prototype.hasOwnProperty.call(properties, key); + const canonicalName = isExact + ? key + : canonicalToProp.get(canonicalizeKey(key)) ?? key; + + const propSchema = Object.prototype.hasOwnProperty.call( + properties, + canonicalName, + ) + ? properties[canonicalName] + : undefined; + + const normalizedChild = normalizeKeysToSchema(child, propSchema); + + const alreadySet = Object.prototype.hasOwnProperty.call( + result, + canonicalName, + ); + if (alreadySet && !isExact) { + // An exact match already claimed this canonical name — drop the alias. + continue; + } + + result[canonicalName] = normalizedChild; + } + + return result; +} + +/** + * Conservatively reports the ways `value` fails to satisfy `schema`. It only + * surfaces problems a model can realistically fix on a retry — missing required + * properties (recursively) and enum violations — and never flags a conformant + * object, so it is safe to use as the "should we repair?" signal: zero issues + * means no extra LLM round-trip. + */ +export function collectSchemaIssues( + value: unknown, + schema: JSONSchema7Definition | undefined, + path = "", +): string[] { + const objectSchema = asObjectSchema(schema); + if (!objectSchema) { + return []; + } + + const issues: string[] = []; + + if (Array.isArray(objectSchema.enum) && value !== undefined) { + const allowed: unknown[] = objectSchema.enum; + if (!allowed.some((candidate) => candidate === value)) { + issues.push( + `property '${path || "value"}' must be one of [${allowed + .map((candidate) => JSON.stringify(candidate)) + .join(", ")}] (received ${JSON.stringify(value)})`, + ); + } + } + + if (Array.isArray(value)) { + const itemSchema = objectSchema.items; + const resolvedItemSchema = Array.isArray(itemSchema) + ? undefined + : itemSchema; + value.forEach((item, index) => { + issues.push( + ...collectSchemaIssues(item, resolvedItemSchema, joinPath(path, index)), + ); + }); + return issues; + } + + if (objectSchema.properties && isPlainObject(value)) { + const required = Array.isArray(objectSchema.required) + ? objectSchema.required + : []; + for (const requiredKey of required) { + if (!Object.prototype.hasOwnProperty.call(value, requiredKey)) { + issues.push( + `missing required property '${joinPath(path, requiredKey)}'`, + ); + } + } + + for (const [propName, propSchema] of Object.entries( + objectSchema.properties, + )) { + if (Object.prototype.hasOwnProperty.call(value, propName)) { + issues.push( + ...collectSchemaIssues( + value[propName], + propSchema, + joinPath(path, propName), + ), + ); + } + } + } + + return issues; +} + +function joinPath(path: string, segment: string | number): string { + return path ? `${path}.${segment}` : String(segment); +} + +/** Top-level property names declared by the schema, in declaration order. */ +export function listExpectedKeys( + schema: JSONSchema7Definition | undefined, +): string[] { + const objectSchema = asObjectSchema(schema); + if (!objectSchema?.properties) { + return []; + } + return Object.keys(objectSchema.properties); +} + +/** + * Builds the user turn used to re-query a model whose JSON drifted from the + * schema. It restates the validation issues and the exact keys the schema + * expects, then asks for corrected JSON only — the generic safety net for + * semantic aliases that deterministic key normalization cannot guess. + */ +export function buildRepairInstruction({ + issues, + expectedKeys, +}: { + issues: string[]; + expectedKeys: string[]; +}): string { + const lines = [ + "Your previous JSON response did not match the required schema.", + ]; + + if (issues.length > 0) { + lines.push( + "Issues:", + ...issues.map((issue) => `- ${issue}`), + ); + } + + if (expectedKeys.length > 0) { + lines.push( + `Return ONLY the corrected JSON object, using EXACTLY these top-level keys: ${expectedKeys + .map((key) => `"${key}"`) + .join(", ")}.`, + ); + } else { + lines.push("Return ONLY the corrected JSON object."); + } + + return lines.join("\n"); +} + +/** + * Orchestrates the two resilience layers over a structured object: + * + * 1. normalize drifted keys onto the schema's property names (cheap, no LLM); + * 2. if the object still misses required keys / violates enums, re-query the + * model up to `maxAttempts` times, re-normalizing each reply. + * + * A conformant (or normalization-fixable) object returns immediately with no + * re-query, keeping the nominal-case cost unchanged. If a re-query throws, the + * best-effort object obtained so far is returned rather than propagating the + * failure — the pipeline never regresses a success into an error. + */ +export async function resolveResilientObject({ + initialObject, + schema, + config, + requery, +}: { + initialObject: unknown; + schema: JSONSchema7Definition | undefined; + config: ResilienceConfig; + requery: (params: { + instruction: string; + previousObject: unknown; + }) => Promise; +}): Promise { + const normalize = (object: unknown) => + config.normalizeKeys ? normalizeKeysToSchema(object, schema) : object; + + let current = normalize(initialObject); + + if (!config.repair.enabled) { + return current; + } + + let issues = collectSchemaIssues(current, schema); + if (issues.length === 0) { + return current; + } + + const expectedKeys = listExpectedKeys(schema); + + for (let attempt = 0; attempt < config.repair.maxAttempts; attempt += 1) { + const instruction = buildRepairInstruction({ issues, expectedKeys }); + + let retried: unknown; + try { + retried = await requery({ instruction, previousObject: current }); + } catch { + // A failed re-query must not turn an already-obtained object into an + // error; deliver the best-effort result we have. + return current; + } + + current = normalize(retried); + issues = collectSchemaIssues(current, schema); + if (issues.length === 0) { + return current; + } + } + + return current; +} diff --git a/packages/core/src/agents/types.ts b/packages/core/src/agents/types.ts index 76497f9..ae86b1b 100644 --- a/packages/core/src/agents/types.ts +++ b/packages/core/src/agents/types.ts @@ -10,6 +10,7 @@ import { } from "ai"; import type { RuntimeState, RuntimeStore } from "../runtime/store.js"; +import type { StructuredOutputResilienceOptions } from "./structuredOutputResilience.js"; export interface AgentTelemetryOverrides { functionId?: string; @@ -68,7 +69,7 @@ export type BaseAgentOptions< loopTools?: boolean; maxStepTools?: number; memory?: MemoryOptions; -}; +} & StructuredOutputResilienceOptions; export type AgentGenerateOptions< OUTPUT = never,