feat(tools): include structured output in exported .prompt files

genkit-tools/common/src/server/router.ts

@@ -34,7 +34,7 @@ import type { PromptFrontmatter } from '../types/prompt';
 import { logger } from '../utils';
 import { PageViewEvent, ToolsRequestEvent, record } from '../utils/analytics';
 import { toolsPackage } from '../utils/package';
-import { fromMessages } from '../utils/prompt';
+import { fromMessages, toFrontmatterOutput } from '../utils/prompt';
 
 const t = initTRPC.create({
   errorFormatter(opts) {
@@ -150,6 +150,7 @@ export const TOOLS_SERVER_ROUTER = (manager: BaseRuntimeManager) =>
           config: input.config,
           tools: input.tools?.map((toolDefinition) => toolDefinition.name),
           use: input.use,
+          output: toFrontmatterOutput(input.output),
         };
         return fromMessages(frontmatter, input.messages);
       }),

genkit-tools/common/src/types/apis.ts

@@ -172,6 +172,17 @@ export const CreatePromptRequestSchema = z.object({
   config: GenerationCommonConfigSchema.passthrough().optional(),
   tools: z.array(ToolDefinitionSchema).optional(),
   use: z.array(MiddlewareRefSchema).optional(),
+  output: z
+    .object({
+      format: z.string().optional(),
+      // Resolved JSON Schema lives under jsonSchema for a generate action; a
+      // model request carries it under schema. Either is accepted.
+      jsonSchema: z.unknown().optional(),
+      schema: z.unknown().optional(),
+      contentType: z.string().optional(),
+    })
+    .passthrough()
+    .optional(),
 });
 
 export type CreatePromptRequest = z.infer<typeof CreatePromptRequestSchema>;

genkit-tools/common/src/utils/prompt.ts

@@ -18,6 +18,119 @@ import { stringify } from 'yaml';
 import type { MessageData, Part } from '../types/model';
 import type { PromptFrontmatter } from '../types/prompt';
 
+/** A JSON Schema-ish object. */
+type JsonSchema = Record<string, any>;
+
+function scalarType(schema: JsonSchema | null | undefined): string {
+  switch (schema?.type) {
+    case 'string':
+    case 'integer':
+    case 'number':
+    case 'boolean':
+    case 'null':
+      return schema.type;
+    default:
+      return 'any';
+  }
+}
+
+/** Wraps a type kind with an optional description, e.g. `(array, the tags)`. */
+function wrap(kind: string, description?: string): string {
+  return description ? `(${kind}, ${description})` : `(${kind})`;
+}
+
+/**
+ * Encodes a single property as a Picoschema key suffix and value. Scalars carry
+ * their description after a comma in the value (`string, the title`); wrapped
+ * kinds (object, array, enum) carry it inside the parentheses on the key.
+ */
+function picoEntry(schema: JsonSchema): { suffix: string; value: any } {
+  const description =
+    typeof schema?.description === 'string' ? schema.description : undefined;
+
+  if (Array.isArray(schema?.enum)) {
+    return { suffix: wrap('enum', description), value: schema.enum };
+  }
+  if (schema?.type === 'array') {
+    const items: JsonSchema = schema.items ?? {};
+    const value =
+      items.type === 'object' || items.properties
+        ? picoObject(items)
+        : scalarType(items);
+    return { suffix: wrap('array', description), value };
+  }
+  if (schema?.type === 'object' || schema?.properties) {
+    return { suffix: wrap('object', description), value: picoObject(schema) };
+  }
+  const type = scalarType(schema);
+  return { suffix: '', value: description ? `${type}, ${description}` : type };
+}
+
+/** Converts an object JSON Schema into a Picoschema object structure. */
+function picoObject(schema: JsonSchema): Record<string, any> {
+  const required = new Set<string>(
+    Array.isArray(schema.required) ? schema.required : []
+  );
+  const out: Record<string, any> = {};
+  for (const [name, propSchema] of Object.entries<any>(
+    schema.properties ?? {}
+  )) {
+    const optional = required.has(name) ? '' : '?';
+    const { suffix, value } = picoEntry(propSchema);
+    out[`${name}${optional}${suffix}`] = value;
+  }
+  const additional = schema.additionalProperties;
+  if (additional && typeof additional === 'object') {
+    const { suffix, value } = picoEntry(additional);
+    out[`(*)${suffix}`] = value;
+  }
+  return out;
+}
+
+/**
+ * Converts a JSON Schema into the equivalent Picoschema for a `.prompt` file.
+ * Object schemas become the compact Picoschema form. Non-object top-level
+ * schemas (a bare array or scalar) have no Picoschema form, so the JSON Schema
+ * is returned unchanged; Dotprompt accepts raw JSON Schema there too.
+ */
+export function jsonSchemaToPicoschema(schema: unknown): any {
+  if (!schema || typeof schema !== 'object') {
+    return schema;
+  }
+  const s = schema as JsonSchema;
+  if (s.type === 'object' || s.properties) {
+    return picoObject(s);
+  }
+  return schema;
+}
+
+/**
+ * Maps a generate request's output config onto `.prompt` frontmatter, converting
+ * the JSON Schema to Picoschema. The frontmatter format is limited to
+ * json/text/media, so the JSON-producing formats (json, jsonl, array, enum) map
+ * onto `json`. Returns undefined when there is nothing to record.
+ */
+export function toFrontmatterOutput(output?: {
+  format?: string;
+  jsonSchema?: unknown;
+  schema?: unknown;
+}): PromptFrontmatter['output'] | undefined {
+  if (!output) return undefined;
+  const result: NonNullable<PromptFrontmatter['output']> = {};
+  if (output.format === 'text') {
+    result.format = 'text';
+  } else if (output.format === 'media') {
+    result.format = 'media';
+  } else if (output.format) {
+    result.format = 'json';
+  }
+  const schema = output.jsonSchema ?? output.schema;
+  if (schema && typeof schema === 'object') {
+    result.schema = jsonSchemaToPicoschema(schema);
+  }
+  return result.format || result.schema ? result : undefined;
+}
+
 export function fromMessages(
   frontmatter: PromptFrontmatter,
   messages: MessageData[]

genkit-tools/common/tests/utils/prompt_test.ts

@@ -17,7 +17,11 @@
 import { describe, expect, it } from '@jest/globals';
 import type { MessageData } from '../../src/types/model';
 import type { PromptFrontmatter } from '../../src/types/prompt';
-import { fromMessages } from '../../src/utils/prompt';
+import {
+  fromMessages,
+  jsonSchemaToPicoschema,
+  toFrontmatterOutput,
+} from '../../src/utils/prompt';
 
 describe('fromMessages', () => {
   it('builds a template from messages', () => {
@@ -182,3 +186,128 @@ describe('fromMessages', () => {
     expect(fromMessages(frontmatter, messages)).toStrictEqual(expected);
   });
 });
+
+describe('jsonSchemaToPicoschema', () => {
+  it('converts an object schema with required, optional, and described fields', () => {
+    const schema = {
+      type: 'object',
+      properties: {
+        title: { type: 'string' },
+        subtitle: { type: 'string', description: 'optional subtitle' },
+        servings: { type: 'integer' },
+      },
+      required: ['title', 'servings'],
+    };
+    expect(jsonSchemaToPicoschema(schema)).toEqual({
+      title: 'string',
+      'subtitle?': 'string, optional subtitle',
+      servings: 'integer',
+    });
+  });
+
+  it('encodes enums, arrays of scalars, arrays of objects, and nested objects', () => {
+    const schema = {
+      type: 'object',
+      properties: {
+        status: {
+          type: 'string',
+          enum: ['PENDING', 'APPROVED'],
+          description: 'approval status',
+        },
+        tags: {
+          type: 'array',
+          items: { type: 'string' },
+          description: 'relevant tags',
+        },
+        authors: {
+          type: 'array',
+          items: {
+            type: 'object',
+            properties: { name: { type: 'string' } },
+            required: ['name'],
+          },
+        },
+        metadata: {
+          type: 'object',
+          properties: { updatedAt: { type: 'string' } },
+        },
+      },
+      required: ['status', 'tags', 'authors'],
+    };
+    expect(jsonSchemaToPicoschema(schema)).toEqual({
+      'status(enum, approval status)': ['PENDING', 'APPROVED'],
+      'tags(array, relevant tags)': 'string',
+      'authors(array)': { name: 'string' },
+      'metadata?(object)': { 'updatedAt?': 'string' },
+    });
+  });
+
+  it('encodes additionalProperties as a wildcard field', () => {
+    const schema = {
+      type: 'object',
+      properties: { id: { type: 'string' } },
+      required: ['id'],
+      additionalProperties: { type: 'number' },
+    };
+    expect(jsonSchemaToPicoschema(schema)).toEqual({
+      id: 'string',
+      '(*)': 'number',
+    });
+  });
+
+  it('passes non-object top-level schemas through unchanged', () => {
+    const arraySchema = { type: 'array', items: { type: 'string' } };
+    expect(jsonSchemaToPicoschema(arraySchema)).toBe(arraySchema);
+  });
+
+  it('does not crash on a null or malformed property', () => {
+    const schema = {
+      type: 'object',
+      properties: { id: { type: 'string' }, broken: null, items: {} },
+      required: ['id'],
+    };
+    expect(jsonSchemaToPicoschema(schema)).toEqual({
+      id: 'string',
+      'broken?': 'any',
+      'items?': 'any',
+    });
+  });
+});
+
+describe('toFrontmatterOutput', () => {
+  const SCHEMA = {
+    type: 'object',
+    properties: { title: { type: 'string' } },
+    required: ['title'],
+  };
+
+  it('returns undefined when there is no output', () => {
+    expect(toFrontmatterOutput(undefined)).toBeUndefined();
+  });
+
+  it('reads the schema from jsonSchema and maps json formats', () => {
+    expect(toFrontmatterOutput({ format: 'json', jsonSchema: SCHEMA })).toEqual(
+      { format: 'json', schema: { title: 'string' } }
+    );
+  });
+
+  it('reads the schema from the schema field (model request shape)', () => {
+    expect(toFrontmatterOutput({ format: 'json', schema: SCHEMA })).toEqual({
+      format: 'json',
+      schema: { title: 'string' },
+    });
+  });
+
+  it('maps json-producing formats onto json', () => {
+    expect(
+      toFrontmatterOutput({ format: 'jsonl', jsonSchema: SCHEMA })?.format
+    ).toBe('json');
+  });
+
+  it('keeps text and media formats', () => {
+    expect(toFrontmatterOutput({ format: 'text' })).toEqual({ format: 'text' });
+    expect(toFrontmatterOutput({ format: 'media' })).toEqual({
+      format: 'media',
+    });
+  });
+});