fix(js/ai): default to JSON format when schema is set (#1728)

firebase · Jan 31, 2025 · 482d5c7 · 482d5c7
1 parent 4db4e9a
commit 482d5c7
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 4 deletions.
diff --git a/js/ai/src/generate.ts b/js/ai/src/generate.ts
@@ -358,7 +358,10 @@ export async function generate<
   });
 
   // If is schema is set but format is not explicitly set, default to `json` format.
-  if (resolvedOptions.output?.schema && !resolvedOptions.output?.format) {
+  if (
+    (resolvedOptions.output?.schema || resolvedOptions.output?.jsonSchema) &&
+    !resolvedOptions.output?.format
+  ) {
     resolvedOptions.output.format = 'json';
   }
   const resolvedFormat = await resolveFormat(registry, resolvedOptions.output);

diff --git a/js/genkit/tests/prompts/output.prompt b/js/genkit/tests/prompts/output.prompt
@@ -0,0 +1,11 @@
+---
+model: staticResponseModel
+input:
+  schema:
+    name: string
+output:
+  schema:
+    bar: string
+---
+
+Hi {{ name }}
diff --git a/js/genkit/tests/prompts_test.ts b/js/genkit/tests/prompts_test.ts
@@ -17,6 +17,7 @@
 import { ModelMiddleware, modelRef } from '@genkit-ai/ai/model';
 import * as assert from 'assert';
 import { beforeEach, describe, it } from 'node:test';
+import { stripUndefinedProps } from '../../core/src';
 import { GenkitBeta, genkit } from '../src/beta';
 import { PromptAction, z } from '../src/index';
 import {
@@ -61,7 +62,7 @@ const wrapResponse: ModelMiddleware = async (req, next) => {
   };
 };
 
-describe('definePrompt - functional', () => {
+describe('definePrompt', () => {
   let ai: GenkitBeta;
 
   beforeEach(() => {
@@ -172,7 +173,7 @@ describe('definePrompt - functional', () => {
   });
 });
 
-describe('definePrompt - dotprompt', () => {
+describe.only('definePrompt', () => {
   describe('default model', () => {
     let ai: GenkitBeta;
 
@@ -309,14 +310,15 @@ describe('definePrompt - dotprompt', () => {
     });
   });
 
-  describe('default model ref', () => {
+  describe.only('default model ref', () => {
     let ai: GenkitBeta;
 
     beforeEach(() => {
       ai = genkit({
         model: modelRef({
           name: 'echoModel',
         }),
+        promptDir: './tests/prompts',
       });
       defineEchoModel(ai);
     });
@@ -400,6 +402,52 @@ describe('definePrompt - dotprompt', () => {
       assert.deepStrictEqual(foo, { bar: 'baz' });
     });
 
+    it('defaults to json format from a loaded prompt', async () => {
+      defineStaticResponseModel(ai, {
+        role: 'model',
+        content: [
+          {
+            text: '```json\n{bar: "baz"}\n```',
+          },
+        ],
+      });
+      const hi = ai.prompt('output');
+
+      const response = await hi({ name: 'Genkit' });
+      const foo = response.output;
+      assert.deepStrictEqual(stripUndefinedProps(response.request), {
+        config: {},
+        messages: [
+          {
+            content: [
+              {
+                text: 'Hi Genkit',
+              },
+            ],
+            role: 'user',
+          },
+        ],
+        output: {
+          constrained: true,
+          contentType: 'application/json',
+          format: 'json',
+          schema: {
+            additionalProperties: false,
+            properties: {
+              bar: {
+                type: 'string',
+              },
+            },
+            required: ['bar'],
+            type: 'object',
+          },
+        },
+        tools: [],
+      });
+
+      assert.deepStrictEqual(foo, { bar: 'baz' });
+    });
+
     it('streams dotprompt with default model', async () => {
       const hi = ai.definePrompt({
         name: 'hi',