Add schema input to LLM nodes for structured JSON output

bchapuis · claude · bchapuis · commit 1266c038fbaa · 2026-03-31T15:40:24.000+02:00
Enable structured output via JSON schema across Workers AI, OpenAI, and
Gemini LLM nodes. When a schema is selected, the executor converts it to
JSON Schema and passes it to each provider's native structured output API.
Output type changed from string to any to support JSON responses.

Co-Authored-By: Claude Opus 4.6 (1M context) &lt;noreply@anthropic.com&gt;
diff --git a/packages/runtime/src/nodes/gemini/execute-gemini-model.ts b/packages/runtime/src/nodes/gemini/execute-gemini-model.ts
@@ -1,7 +1,8 @@
 import type { NodeContext, ToolReference } from "@dafthunk/runtime";
-import type { NodeExecution } from "@dafthunk/types";
+import type { NodeExecution, Schema } from "@dafthunk/types";
 import { GoogleGenAI } from "@google/genai";
 import { getGoogleAIConfig } from "../../utils/ai-gateway";
+import { schemaToJsonSchema } from "../../utils/schema-to-json-schema";
 import { calculateTokenUsage, type TokenPricing } from "../../utils/usage";
 
 interface GeminiModelConfig {
@@ -41,6 +42,7 @@ export async function executeGeminiModel(
       thinking_budget,
       tools,
       googleSearch,
+      schema: schemaInput,
     } = context.inputs;
 
     if (!input) {
@@ -64,6 +66,16 @@ export async function executeGeminiModel(
       genConfig.thinkingConfig = { thinkingBudget: thinking_budget };
     }
 
+    // Add structured output when a schema is provided
+    if (
+      schemaInput &&
+      typeof schemaInput === "object" &&
+      "fields" in schemaInput
+    ) {
+      genConfig.responseMimeType = "application/json";
+      genConfig.responseSchema = schemaToJsonSchema(schemaInput as Schema);
+    }
+
     const functionDeclarations =
       await node.convertFunctionCallsToGeminiDeclarations(
         tools as ToolReference[],
diff --git a/packages/runtime/src/nodes/gemini/gemini-2-5-flash-node.ts b/packages/runtime/src/nodes/gemini/gemini-2-5-flash-node.ts
@@ -75,12 +75,18 @@ export class Gemini25FlashNode extends ExecutableNode {
         required: false,
         hidden: true,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "text",
-        type: "string",
-        description: "Generated text response from Gemini 2.5 Flash",
+        type: "any",
+        description: "Generated text or JSON response from Gemini 2.5 Flash",
       },
       {
         name: "candidates",
diff --git a/packages/runtime/src/nodes/gemini/gemini-2-5-pro-node.ts b/packages/runtime/src/nodes/gemini/gemini-2-5-pro-node.ts
@@ -75,12 +75,18 @@ export class Gemini25ProNode extends ExecutableNode {
         required: false,
         hidden: true,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "text",
-        type: "string",
-        description: "Generated text response from Gemini 2.5 Pro",
+        type: "any",
+        description: "Generated text or JSON response from Gemini 2.5 Pro",
       },
       {
         name: "candidates",
diff --git a/packages/runtime/src/nodes/gemini/gemini-3-1-pro-node.ts b/packages/runtime/src/nodes/gemini/gemini-3-1-pro-node.ts
@@ -75,12 +75,18 @@ export class Gemini31ProNode extends ExecutableNode {
         required: false,
         hidden: true,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "text",
-        type: "string",
-        description: "Generated text response from Gemini 3.1 Pro",
+        type: "any",
+        description: "Generated text or JSON response from Gemini 3.1 Pro",
       },
       {
         name: "candidates",
diff --git a/packages/runtime/src/nodes/gemini/gemini-3-flash-node.ts b/packages/runtime/src/nodes/gemini/gemini-3-flash-node.ts
@@ -75,12 +75,18 @@ export class Gemini3FlashNode extends ExecutableNode {
         required: false,
         hidden: true,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "text",
-        type: "string",
-        description: "Generated text response from Gemini 3 Flash",
+        type: "any",
+        description: "Generated text or JSON response from Gemini 3 Flash",
       },
       {
         name: "candidates",
diff --git a/packages/runtime/src/nodes/openai/execute-openai-model.ts b/packages/runtime/src/nodes/openai/execute-openai-model.ts
@@ -1,7 +1,8 @@
 import type { NodeContext } from "@dafthunk/runtime";
-import type { NodeExecution } from "@dafthunk/types";
+import type { NodeExecution, Schema } from "@dafthunk/types";
 import OpenAI from "openai";
 import { getOpenAIConfig } from "../../utils/ai-gateway";
+import { schemaToJsonSchema } from "../../utils/schema-to-json-schema";
 import { calculateTokenUsage, type TokenPricing } from "../../utils/usage";
 
 /**
@@ -23,7 +24,7 @@ export async function executeOpenAIModel(
   pricing: TokenPricing
 ): Promise<NodeExecution> {
   try {
-    const { instructions, input } = context.inputs;
+    const { instructions, input, schema: schemaInput } = context.inputs;
 
     if (!input) {
       return node.createErrorResult("Input is required");
@@ -35,6 +36,21 @@ export async function executeOpenAIModel(
       ...getOpenAIConfig(context.env),
     });
 
+    // Build response_format when a schema is provided
+    const responseFormat =
+      schemaInput &&
+      typeof schemaInput === "object" &&
+      "fields" in schemaInput
+        ? {
+            type: "json_schema" as const,
+            json_schema: {
+              name: "response",
+              schema: schemaToJsonSchema(schemaInput as Schema),
+              strict: true,
+            },
+          }
+        : undefined;
+
     const completion = await client.chat.completions.create({
       model: modelId,
       max_tokens: 1024,
@@ -44,6 +60,7 @@ export async function executeOpenAIModel(
           : []),
         { role: "user" as const, content: input },
       ],
+      ...(responseFormat && { response_format: responseFormat }),
     });
 
     const responseText = completion.choices[0]?.message?.content || "";
diff --git a/packages/runtime/src/nodes/openai/gpt-41-node.ts b/packages/runtime/src/nodes/openai/gpt-41-node.ts
@@ -34,12 +34,18 @@ export class Gpt41Node extends ExecutableNode {
         description: "The input text or question for GPT-4.1",
         required: true,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "text",
-        type: "string",
-        description: "Generated text response from GPT-4.1",
+        type: "any",
+        description: "Generated text or JSON response from GPT-4.1",
       },
     ],
   };
diff --git a/packages/runtime/src/nodes/openai/gpt-5-mini-node.ts b/packages/runtime/src/nodes/openai/gpt-5-mini-node.ts
@@ -34,12 +34,18 @@ export class Gpt5MiniNode extends ExecutableNode {
         description: "The input text or question for GPT-5 Mini",
         required: true,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "text",
-        type: "string",
-        description: "Generated text response from GPT-5 Mini",
+        type: "any",
+        description: "Generated text or JSON response from GPT-5 Mini",
       },
     ],
   };
diff --git a/packages/runtime/src/nodes/openai/gpt-5-nano-node.ts b/packages/runtime/src/nodes/openai/gpt-5-nano-node.ts
@@ -34,12 +34,18 @@ export class Gpt5NanoNode extends ExecutableNode {
         description: "The input text or question for GPT-5 Nano",
         required: true,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "text",
-        type: "string",
-        description: "Generated text response from GPT-5 Nano",
+        type: "any",
+        description: "Generated text or JSON response from GPT-5 Nano",
       },
     ],
   };
diff --git a/packages/runtime/src/nodes/openai/gpt-5-node.ts b/packages/runtime/src/nodes/openai/gpt-5-node.ts
@@ -34,12 +34,18 @@ export class Gpt5Node extends ExecutableNode {
         description: "The input text or question for GPT-5",
         required: true,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "text",
-        type: "string",
-        description: "Generated text response from GPT-5",
+        type: "any",
+        description: "Generated text or JSON response from GPT-5",
       },
     ],
   };
diff --git a/packages/runtime/src/nodes/text/deepseek-r1-distill-qwen-32b-node.ts b/packages/runtime/src/nodes/text/deepseek-r1-distill-qwen-32b-node.ts
@@ -101,12 +101,18 @@ export class DeepseekR1DistillQwen32BNode extends ExecutableNode {
         hidden: true,
         value: 0,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "response",
-        type: "string",
-        description: "Generated text response",
+        type: "any",
+        description: "Generated text or JSON response",
       },
     ],
   };
diff --git a/packages/runtime/src/nodes/text/execute-workers-ai-text-model.ts b/packages/runtime/src/nodes/text/execute-workers-ai-text-model.ts
@@ -1,7 +1,8 @@
 import { runWithTools } from "@cloudflare/ai-utils";
 import type { NodeContext, ToolReference } from "@dafthunk/runtime";
 import { type ExecutableNode, ToolCallTracker } from "@dafthunk/runtime";
-import type { NodeExecution } from "@dafthunk/types";
+import type { NodeExecution, Schema } from "@dafthunk/types";
+import { schemaToJsonSchema } from "../../utils/schema-to-json-schema";
 import { calculateTokenUsage, type TokenPricing } from "../../utils/usage";
 
 /**
@@ -29,7 +30,7 @@ export async function executeWorkersAiTextModel(
   config: WorkersAiTextConfig
 ): Promise<NodeExecution> {
   try {
-    const { prompt, messages, tools } = context.inputs;
+    const { prompt, messages, tools, schema: schemaInput } = context.inputs;
 
     if (!context.env?.AI) {
       return node.createErrorResult("AI service is not available");
@@ -54,7 +55,18 @@ export async function executeWorkersAiTextModel(
       context
     );
 
-    const extraParams = config.params ?? {};
+    // Build response_format when a schema is provided
+    const extraParams: Record<string, unknown> = { ...(config.params ?? {}) };
+    if (schemaInput && typeof schemaInput === "object" && "fields" in schemaInput) {
+      extraParams.response_format = {
+        type: "json_schema",
+        json_schema: {
+          name: "response",
+          schema: schemaToJsonSchema(schemaInput as Schema),
+        },
+      };
+    }
+
     let result: any;
     let executedToolCalls: any[] = [];
 
diff --git a/packages/runtime/src/nodes/text/glm-4-7-flash-node.ts b/packages/runtime/src/nodes/text/glm-4-7-flash-node.ts
@@ -90,12 +90,18 @@ export class Glm47FlashNode extends ExecutableNode {
         hidden: true,
         value: 0.0,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "response",
-        type: "string",
-        description: "Generated text response",
+        type: "any",
+        description: "Generated text or JSON response",
       },
       {
         name: "reasoning",
diff --git a/packages/runtime/src/nodes/text/hermes-2-pro-mistral-7b-node.ts b/packages/runtime/src/nodes/text/hermes-2-pro-mistral-7b-node.ts
@@ -104,12 +104,18 @@ export class Hermes2ProMistral7BNode extends ExecutableNode {
         hidden: true,
         value: 0.0,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "response",
-        type: "string",
-        description: "Generated text response",
+        type: "any",
+        description: "Generated text or JSON response",
       },
       {
         name: "tool_calls",
diff --git a/packages/runtime/src/nodes/text/kimi-k2-5-node.ts b/packages/runtime/src/nodes/text/kimi-k2-5-node.ts
@@ -90,12 +90,18 @@ export class KimiK25Node extends ExecutableNode {
         hidden: true,
         value: 0.0,
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "response",
-        type: "string",
-        description: "Generated text response",
+        type: "any",
+        description: "Generated text or JSON response",
       },
       {
         name: "reasoning",
diff --git a/packages/runtime/src/nodes/text/llama-3-1-8b-instruct-fast-node.ts b/packages/runtime/src/nodes/text/llama-3-1-8b-instruct-fast-node.ts
@@ -46,12 +46,18 @@ export class Llama318BInstructFastNode extends ExecutableNode {
         hidden: true,
         value: undefined, // Optional
       },
+      {
+        name: "schema",
+        type: "schema",
+        description: "JSON schema to constrain the output format",
+        hidden: true,
+      },
     ],
     outputs: [
       {
         name: "response",
-        type: "string",
-        description: "Generated text response",
+        type: "any",
+        description: "Generated text or JSON response",
       },
     ],
   };
diff --git a/packages/runtime/src/nodes/text/llama-3-3-70b-instruct-fp8-fast-node.ts b/packages/runtime/src/nodes/text/llama-3-3-70b-instruct-fp8-fast-node.ts
diff --git a/packages/runtime/src/nodes/text/llama-4-scout-17b-16e-instruct-node.ts b/packages/runtime/src/nodes/text/llama-4-scout-17b-16e-instruct-node.ts
diff --git a/packages/runtime/src/nodes/text/mistral-small-3-1-24b-instruct-node.ts b/packages/runtime/src/nodes/text/mistral-small-3-1-24b-instruct-node.ts
diff --git a/packages/runtime/src/nodes/text/qwen3-30b-a3b-fp8-node.ts b/packages/runtime/src/nodes/text/qwen3-30b-a3b-fp8-node.ts
diff --git a/packages/runtime/src/nodes/text/qwq-32b-node.ts b/packages/runtime/src/nodes/text/qwq-32b-node.ts
diff --git a/packages/runtime/src/utils/schema-to-json-schema.ts b/packages/runtime/src/utils/schema-to-json-schema.ts
