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,
@@ -47,6 +48,15 @@ export function isPrompt(arg: any): boolean {
   );
 }
 
+/**
+ * Defines and registers a prompt action. The action can be called to obtain
+ * a `GenerateRequest` which can be passed to a model action. The given
+ * `PromptFn` can perform any action needed to create the request such as rendering
+ * a template or fetching a prompt from a database.
+ *
+ * @returns The new `PromptAction`.
+ */
+
 export function definePrompt<I extends z.ZodTypeAny>(
   {
     name,
@@ -77,36 +87,46 @@ export function definePrompt<I extends z.ZodTypeAny>(
   return a as PromptAction<I>;
 }
 
-/**
- * A veneer for rendering a prompt action to GenerateOptions.
- */
-
 export type PromptArgument<I extends z.ZodTypeAny = z.ZodTypeAny> =
   | string
   | PromptAction<I>;
 
+/**
+ * This veneer renders a `PromptAction` into a `GenerateOptions` object.
+ *
+ * @returns A promise of an options object for use with the `generate()` function.
+ */
 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,7 @@ interface RenderMetadata {
   history?: MessageData[];
 }
 
-export class Dotprompt<Variables = unknown> implements PromptMetadata {
+export class Dotprompt<I = unknown> implements PromptMetadata<z.ZodTypeAny> {
   name: string;
   variant?: string;
   hash: string;
@@ -70,10 +70,7 @@ export class Dotprompt<Variables = unknown> implements PromptMetadata {
   config?: PromptMetadata['config'];
   candidates?: PromptMetadata['candidates'];
 
-  private _render: (
-    input: Variables,
-    options?: RenderMetadata
-  ) => MessageData[];
+  private _render: (input: I, options?: RenderMetadata) => MessageData[];
 
   static parse(name: string, source: string) {
     try {
@@ -123,7 +120,15 @@ export class Dotprompt<Variables = unknown> implements PromptMetadata {
     this._render = compile(this.template, options);
   }
 
-  renderText(input: Variables, options?: RenderMetadata): string {
+  /**
+   * Renders all of the prompt's text parts into a raw string.
+   *
+   * @param input User input to the prompt template.
+   * @param options Optional context and/or history for the prompt template.
+   * @returns all of the text parts concatenated into a string.
+   */
+
+  renderText(input: I, options?: RenderMetadata): string {
     const result = this.renderMessages(input, options);
     if (result.length !== 1) {
       throw new Error("Multi-message prompt can't be rendered as text.");
@@ -138,7 +143,15 @@ export class Dotprompt<Variables = unknown> implements PromptMetadata {
     return out;
   }
 
-  renderMessages(input?: Variables, options?: RenderMetadata): MessageData[] {
+  /**
+   * Renders the prompt template into an array of messages.
+   *
+   * @param input User input to the prompt template
+   * @param options optional context and/or history for the prompt template.
+   * @returns an array of messages representing an exchange between a user and a model.
+   */
+
+  renderMessages(input?: I, options?: RenderMetadata): MessageData[] {
     input = parseSchema(input, {
       schema: this.input?.schema,
       jsonSchema: this.input?.jsonSchema,
@@ -162,13 +175,14 @@ export class Dotprompt<Variables = unknown> implements PromptMetadata {
           prompt: this.toJSON(),
         },
       },
-      async (input?: Variables) => toGenerateRequest(this.render({ input }))
+      async (input?: I) => toGenerateRequest(this.render({ input }))
     );
   }
 
-  private _generateOptions<O extends z.ZodTypeAny = z.ZodTypeAny>(
-    options: PromptGenerateOptions<Variables>
-  ): GenerateOptions<z.ZodTypeAny, O> {
+  private _generateOptions<
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+    CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+  >(options: PromptGenerateOptions<I>): GenerateOptions<O, CustomOptions> {
     const messages = this.renderMessages(options.input, {
       history: options.history,
       context: options.context,
@@ -188,25 +202,49 @@ 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);
+  /**
+   * Renders the prompt template based on user input.
+   *
+   * @param opt Options for the prompt template, including user input variables and custom model configuration options.
+   * @returns a `GenerateOptions` object to be used with the `generate()` function from @genkit-ai/ai.
+   */
+  render<
+    CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+  >(
+    opt: PromptGenerateOptions<I, CustomOptions>
+  ): GenerateOptions<CustomOptions, O> {
+    return this._generateOptions(opt);
   }
 
-  async generate<O extends z.ZodTypeAny = z.ZodTypeAny>(
-    opt: PromptGenerateOptions<Variables>
+  /**
+   * Generates a response by rendering the prompt template with given user input and then calling the model.
+   *
+   * @param opt Options for the prompt template, including user input variables and custom model configuration options.
+   * @returns the model response as a promise of `GenerateResponse`.
+   */
+  async generate<
+    CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+    O extends z.ZodTypeAny = z.ZodTypeAny,
+  >(
+    opt: PromptGenerateOptions<I, 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>
+  /**
+   * Generates a streaming response by rendering the prompt template with given user input and then calling the model.
+   *
+   * @param opt Options for the prompt template, including user input variables and custom model configuration options.
+   * @returns the model response as a promise of `GenerateStreamResponse`.
+   */
+  async generateStream<CustomOptions extends z.ZodTypeAny = z.ZodTypeAny>(
+    opt: PromptGenerateOptions<I, CustomOptions>
   ): Promise<GenerateStreamResponse> {
-    return generateStream(this.render(opt));
+    return generateStream(this.render<CustomOptions>(opt));
   }
 }
 
@@ -228,6 +266,7 @@ export class DotpromptRef<Variables = unknown> {
     this.dir = options?.dir;
   }
 
+  /** Loads the prompt which is referenced. */
   async loadPrompt(): Promise<Dotprompt<Variables>> {
     if (this._prompt) return this._prompt;
     this._prompt = (await lookupPrompt(
@@ -238,25 +277,54 @@ export class DotpromptRef<Variables = unknown> {
     return this._prompt;
   }
 
-  async generate<O extends z.ZodTypeAny = z.ZodTypeAny>(
-    opt: PromptGenerateOptions<Variables>
+  /**
+   * Generates a response by rendering the prompt template with given user input and then calling the model.
+   *
+   * @param opt Options for the prompt template, including user input variables and custom model configuration options.
+   * @returns the model response as a promise of `GenerateResponse`.
+   */
+
+  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>
+  /**
+   * Renders the prompt template based on user input.
+   *
+   * @param opt Options for the prompt template, including user input variables and custom model configuration options.
+   * @returns a `GenerateOptions` object to be used with the `generate()` function from @genkit-ai/ai.
+   */
+  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>,
+/**
+ * Define a dotprompt in code. This function is offered as an alternative to definitions in .prompt files.
+ *
+ * @param options the prompt definition, including its name, variant and model. Any options from .prompt file front matter are accepted.
+ * @param template a string template, comparable to the main body of a prompt file.
+ * @returns the newly defined prompt.
+ */
+export function defineDotprompt<
+  I extends z.ZodTypeAny = z.ZodTypeAny,
+  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+>(
+  options: PromptMetadata<I, CustomOptions>,
   template: string
-): Dotprompt<z.infer<V>> {
+): Dotprompt<z.infer<I>> {
   const prompt = new Dotprompt(options, template);
   prompt.define();
   return prompt;