Fixing bugs in prompt rendering

js/ai/src/prompt.ts

@@ -25,9 +25,10 @@ import {
   ModelArgument,
 } from './model.js';
 
-export type PromptFn<I extends z.ZodTypeAny = z.ZodTypeAny> = (
-  input: z.infer<I>
-) => Promise<GenerateRequest>;
+export type PromptFn<
+  I extends z.ZodTypeAny = z.ZodTypeAny,
+  CustomOptionsSchema extends z.ZodTypeAny = z.ZodTypeAny,
+> = (input: z.infer<I>) => Promise<GenerateRequest<CustomOptionsSchema>>;
 
 export type PromptAction<I extends z.ZodTypeAny = z.ZodTypeAny> = Action<
   I,
@@ -87,26 +88,35 @@ export type PromptArgument<I extends z.ZodTypeAny = z.ZodTypeAny> =
 
 export async function renderPrompt<
   I extends z.ZodTypeAny = z.ZodTypeAny,
+  O extends z.ZodTypeAny = z.ZodTypeAny,
   CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
 >(params: {
   prompt: PromptArgument<I>;
   input: z.infer<I>;
   context?: DocumentData[];
   model: ModelArgument<CustomOptions>;
   config?: z.infer<CustomOptions>;
-}): Promise<GenerateOptions> {
+}): Promise<GenerateOptions<O, CustomOptions>> {
   let prompt: PromptAction<I>;
   if (typeof params.prompt === 'string') {
     prompt = await lookupAction(`/prompt/${params.prompt}`);
   } else {
     prompt = params.prompt as PromptAction<I>;
   }
-  const rendered = await prompt(params.input);
+  const rendered = (await prompt(
+    params.input
+  )) as GenerateRequest<CustomOptions>;
   return {
     model: params.model,
     config: { ...(rendered.config || {}), ...params.config },
     history: rendered.messages.slice(0, rendered.messages.length - 1),
     prompt: rendered.messages[rendered.messages.length - 1].content,
     context: params.context,
-  };
+    candidates: rendered.candidates || 1,
+    output: {
+      format: rendered.output?.format,
+      schema: rendered.output?.schema,
+    },
+    tools: rendered.tools || [],
+  } as GenerateOptions<O, CustomOptions>;
 }

js/ai/tests/prompt/prompt_test.ts

@@ -0,0 +1,56 @@
+/**
+ * Copyright 2024 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+import assert from 'node:assert';
+import { describe, it } from 'node:test';
+import * as z from 'zod';
+import { definePrompt, renderPrompt } from '../../src/prompt.ts';
+
+describe('prompt', () => {
+  describe('render()', () => {
+    it('respects output schema in the definition', async () => {
+      const schema1 = z.object({
+        puppyName: z.string({ description: 'A cute name for a puppy' }),
+      });
+      const prompt1 = definePrompt(
+        {
+          name: 'prompt1',
+          inputSchema: z.string({ description: 'Dog breed' }),
+        },
+        async (breed) => {
+          return {
+            messages: [
+              {
+                role: 'user',
+                content: [{ text: `Pick a name for a ${breed} puppy` }],
+              },
+            ],
+            output: {
+              format: 'json',
+              schema: schema1,
+            },
+          };
+        }
+      );
+      const generateRequest = await renderPrompt({
+        prompt: prompt1,
+        input: 'poodle',
+        model: 'geminiPro',
+      });
+      assert.equal(generateRequest.output?.schema, schema1);
+    });
+  });
+});

js/plugins/dotprompt/src/prompt.ts

@@ -24,7 +24,7 @@ import {
   PromptAction,
   toGenerateRequest,
 } from '@genkit-ai/ai';
-import { GenerationCommonConfigSchema, MessageData } from '@genkit-ai/ai/model';
+import { MessageData, ModelArgument } from '@genkit-ai/ai/model';
 import { DocumentData } from '@genkit-ai/ai/retriever';
 import { GenkitError } from '@genkit-ai/core';
 import { parseSchema } from '@genkit-ai/core/schema';
@@ -42,11 +42,11 @@ import { compile } from './template.js';
 
 export type PromptData = PromptFrontmatter & { template: string };
 
-export type PromptGenerateOptions<V = unknown> = Omit<
-  GenerateOptions<z.ZodTypeAny, typeof GenerationCommonConfigSchema>,
-  'prompt' | 'model'
-> & {
-  model?: string;
+export type PromptGenerateOptions<
+  V = unknown,
+  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+> = Omit<GenerateOptions<z.ZodTypeAny, CustomOptions>, 'prompt' | 'model'> & {
+  model?: ModelArgument<CustomOptions>;
   input?: V;
 };
 
@@ -55,7 +55,9 @@ interface RenderMetadata {
   history?: MessageData[];
 }
 
-export class Dotprompt<Variables = unknown> implements PromptMetadata {
+export class Dotprompt<Variables = unknown>
+  implements PromptMetadata<z.ZodTypeAny>
+{
   name: string;
   variant?: string;
   hash: string;
@@ -166,9 +168,12 @@ export class Dotprompt<Variables = unknown> implements PromptMetadata {
     );
   }
 
-  private _generateOptions<O extends z.ZodTypeAny = z.ZodTypeAny>(
+  private _generateOptions<
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+    CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+  >(
     options: PromptGenerateOptions<Variables>
-  ): GenerateOptions<z.ZodTypeAny, O> {
+  ): GenerateOptions<O, CustomOptions> {
     const messages = this.renderMessages(options.input, {
       history: options.history,
       context: options.context,
@@ -188,25 +193,31 @@ export class Dotprompt<Variables = unknown> implements PromptMetadata {
       tools: (options.tools || []).concat(this.tools || []),
       streamingCallback: options.streamingCallback,
       returnToolRequests: options.returnToolRequests,
-    } as GenerateOptions<z.ZodTypeAny, O>;
+    } as GenerateOptions<O, CustomOptions>;
   }
 
-  render<O extends z.ZodTypeAny = z.ZodTypeAny>(
-    opt: PromptGenerateOptions<Variables>
-  ): GenerateOptions<z.ZodTypeAny, O> {
-    return this._generateOptions<O>(opt);
+  render<
+    CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+  >(
+    opt: PromptGenerateOptions<Variables, CustomOptions>
+  ): GenerateOptions<CustomOptions, O> {
+    return this._generateOptions(opt);
   }
 
-  async generate<O extends z.ZodTypeAny = z.ZodTypeAny>(
-    opt: PromptGenerateOptions<Variables>
+  async generate<
+    CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+  >(
+    opt: PromptGenerateOptions<Variables, CustomOptions>
   ): Promise<GenerateResponse<z.infer<O>>> {
-    return generate<z.ZodTypeAny, O>(this.render<O>(opt));
+    return generate<CustomOptions, O>(this.render<CustomOptions, O>(opt));
   }
 
-  async generateStream(
-    opt: PromptGenerateOptions<Variables>
+  async generateStream<CustomOptions extends z.ZodTypeAny = z.ZodTypeAny>(
+    opt: PromptGenerateOptions<Variables, CustomOptions>
   ): Promise<GenerateStreamResponse> {
-    return generateStream(this.render(opt));
+    return generateStream(this.render<CustomOptions>(opt));
   }
 }
 
@@ -238,23 +249,32 @@ export class DotpromptRef<Variables = unknown> {
     return this._prompt;
   }
 
-  async generate<O extends z.ZodTypeAny = z.ZodTypeAny>(
-    opt: PromptGenerateOptions<Variables>
+  async generate<
+    CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+  >(
+    opt: PromptGenerateOptions<Variables, CustomOptions>
   ): Promise<GenerateResponse<z.infer<O>>> {
     const prompt = await this.loadPrompt();
-    return prompt.generate<O>(opt);
+    return prompt.generate<CustomOptions, O>(opt);
   }
 
-  async render<O extends z.ZodTypeAny = z.ZodTypeAny>(
-    opt: PromptGenerateOptions<Variables>
+  async render<
+    CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+  >(
+    opt: PromptGenerateOptions<Variables, CustomOptions>
   ): Promise<GenerateOptions<z.ZodTypeAny, O>> {
     const prompt = await this.loadPrompt();
-    return prompt.render<O>(opt);
+    return prompt.render<CustomOptions, O>(opt);
   }
 }
 
-export function defineDotprompt<V extends z.ZodTypeAny = z.ZodTypeAny>(
-  options: PromptMetadata<V>,
+export function defineDotprompt<
+  V extends z.ZodTypeAny = z.ZodTypeAny,
+  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+>(
+  options: PromptMetadata<V, CustomOptions>,
   template: string
 ): Dotprompt<z.infer<V>> {
   const prompt = new Dotprompt(options, template);

js/testapps/dev-ui-gallery/src/main/prompts.ts

@@ -46,26 +46,24 @@ export const codeDefinedPrompt = defineDotprompt(
       topK: 16,
       topP: 0.95,
       stopSequences: ['STAWP!'],
-      custom: {
-        safetySettings: [
-          {
-            category: 'HARM_CATEGORY_HATE_SPEECH',
-            threshold: 'BLOCK_ONLY_HIGH',
-          },
-          {
-            category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
-            threshold: 'BLOCK_ONLY_HIGH',
-          },
-          {
-            category: 'HARM_CATEGORY_HARASSMENT',
-            threshold: 'BLOCK_ONLY_HIGH',
-          },
-          {
-            category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
-            threshold: 'BLOCK_ONLY_HIGH',
-          },
-        ],
-      },
+      safetySettings: [
+        {
+          category: 'HARM_CATEGORY_HATE_SPEECH',
+          threshold: 'BLOCK_ONLY_HIGH',
+        },
+        {
+          category: 'HARM_CATEGORY_DANGEROUS_CONTENT',
+          threshold: 'BLOCK_ONLY_HIGH',
+        },
+        {
+          category: 'HARM_CATEGORY_HARASSMENT',
+          threshold: 'BLOCK_ONLY_HIGH',
+        },
+        {
+          category: 'HARM_CATEGORY_SEXUALLY_EXPLICIT',
+          threshold: 'BLOCK_ONLY_HIGH',
+        },
+      ],
     },
   },
   template