From 4e53051b0e18aad185d7b5709262242bf50f0a3f Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Thu, 14 Nov 2024 14:25:01 -0500
Subject: [PATCH] fix: fixed dotprompt history (#1301)

---
 js/genkit/src/genkit.ts         | 23 ++++++++++---
 js/genkit/tests/prompts_test.ts | 59 +++++++++++++++++++++++++++++----
 2 files changed, 71 insertions(+), 11 deletions(-)

diff --git a/js/genkit/src/genkit.ts b/js/genkit/src/genkit.ts
index fae7dc87b..18af70822 100644
--- a/js/genkit/src/genkit.ts
+++ b/js/genkit/src/genkit.ts
@@ -122,6 +122,7 @@ import {
   defineDotprompt,
   defineHelper,
   definePartial,
+  Dotprompt,
   PromptMetadata as DotpromptPromptMetadata,
   loadPromptFolder,
   prompt,
@@ -398,7 +399,8 @@ export class Genkit {
       );
       return this.wrapPromptActionInExecutablePrompt(
         dotprompt.promptAction! as PromptAction<I>,
-        options
+        options,
+        dotprompt
       );
     } else {
       const p = definePrompt(
@@ -434,7 +436,8 @@ export class Genkit {
     promptAction: PromptAction<I> | Promise<PromptAction<I>>,
     options:
       | Partial<PromptMetadata<I, CustomOptions>>
-      | Promise<Partial<PromptMetadata<I, CustomOptions>>>
+      | Promise<Partial<PromptMetadata<I, CustomOptions>>>,
+    dotprompt?: Dotprompt<z.infer<I>>
   ): ExecutablePrompt<I, O, CustomOptions> {
     const executablePrompt = async (
       input?: z.infer<I>,
@@ -476,7 +479,11 @@ export class Genkit {
         }
       }
       const p = await promptAction;
-      const promptResult = await p(opt.input);
+      // If it's a dotprompt template, we invoke dotprompt template directly
+      // because it can take in more PromptGenerateOptions (not just inputs).
+      const promptResult = await (dotprompt
+        ? dotprompt.render(opt)
+        : p(opt.input));
       const resultOptions = {
         messages: promptResult.messages,
         docs: promptResult.docs,
@@ -485,7 +492,12 @@ export class Genkit {
           promptResult.output?.format || promptResult.output?.schema
             ? {
                 format: promptResult.output?.format,
-                jsonSchema: promptResult.output?.schema,
+                jsonSchema: dotprompt
+                  ? (promptResult as GenerateOptions).output?.jsonSchema
+                  : promptResult.output.schema,
+                contentType: promptResult.output?.contentType,
+                instructions: promptResult.output?.instructions,
+                schema: promptResult.output?.schema,
               }
             : options.output,
         config: {
@@ -496,6 +508,9 @@ export class Genkit {
         model,
       } as GenerateOptions<O, CustomOptions>;
       delete (resultOptions as any).input;
+      if ((promptResult as GenerateOptions).prompt) {
+        resultOptions.prompt = (promptResult as GenerateOptions).prompt;
+      }
       return resultOptions;
     };
     (executablePrompt as ExecutablePrompt<I, O, CustomOptions>).asTool =
diff --git a/js/genkit/tests/prompts_test.ts b/js/genkit/tests/prompts_test.ts
index a478385cd..f1c7f6d5e 100644
--- a/js/genkit/tests/prompts_test.ts
+++ b/js/genkit/tests/prompts_test.ts
@@ -148,6 +148,7 @@ describe('definePrompt - dotprompt', () => {
             }),
           },
           output: {
+            format: 'json',
             schema: Foo,
           },
         },
@@ -235,6 +236,54 @@ describe('definePrompt - dotprompt', () => {
       assert.strictEqual(response.text, 'Echo: hi Genkit; config: {}');
     });
 
+    it('calls dotprompt with history', async () => {
+      const hi = ai.definePrompt(
+        {
+          name: 'hi',
+          model: 'echoModel',
+          input: {
+            schema: z.object({
+              name: z.string(),
+            }),
+          },
+        },
+        '{{ history}} hi {{ name }}'
+      );
+
+      const response = await hi(
+        { name: 'Genkit' },
+        {
+          messages: [
+            { role: 'user', content: [{ text: 'hi' }] },
+            { role: 'model', content: [{ text: 'bye' }] },
+          ],
+        }
+      );
+      assert.deepStrictEqual(response.messages, [
+        {
+          role: 'user',
+          content: [{ text: 'hi' }],
+          metadata: { purpose: 'history' },
+        },
+        {
+          role: 'model',
+          content: [{ text: 'bye' }],
+          metadata: { purpose: 'history' },
+        },
+        {
+          role: 'model',
+          content: [{ text: ' hi Genkit' }],
+        },
+        {
+          role: 'model',
+          content: [
+            { text: 'Echo: hi,bye, hi Genkit' },
+            { text: '; config: {}' },
+          ],
+        },
+      ]);
+    });
+
     it('calls dotprompt with default model with config', async () => {
       const hi = ai.definePrompt(
         {
@@ -306,16 +355,12 @@ describe('definePrompt - dotprompt', () => {
       assert.deepStrictEqual(response, {
         config: {},
         docs: undefined,
-        messages: [
+        prompt: [
           {
-            content: [
-              {
-                text: 'hi Genkit',
-              },
-            ],
-            role: 'user',
+            text: 'hi Genkit',
           },
         ],
+        messages: [],
         output: undefined,
         tools: [],
       });