From 3eff1717ae3c61e3544fcc38cfd3da3500652289 Mon Sep 17 00:00:00 2001
From: Michael Doyle <michaeldoyle@google.com>
Date: Fri, 22 Nov 2024 11:37:23 -0500
Subject: [PATCH] Properly look up prompt variants (#1376)

---
 js/genkit/src/genkit.ts                     |  2 +-
 js/genkit/tests/prompts/test.variant.prompt |  6 ++++++
 js/genkit/tests/prompts_test.ts             | 11 +++++++++++
 3 files changed, 18 insertions(+), 1 deletion(-)
 create mode 100644 js/genkit/tests/prompts/test.variant.prompt

diff --git a/js/genkit/src/genkit.ts b/js/genkit/src/genkit.ts
index 23dd5f597..810fd8c3a 100644
--- a/js/genkit/src/genkit.ts
+++ b/js/genkit/src/genkit.ts
@@ -332,7 +332,7 @@ export class Genkit {
       // check the registry first as not all prompt types can be
       // loaded by dotprompt (e.g. functional)
       let action = (await this.registry.lookupAction(
-        `/prompt/${name}`
+        `/prompt/${name}${options?.variant ? `.${options?.variant}` : ''}`
       )) as PromptAction<I>;
       // nothing in registry - check for dotprompt file.
       if (!action) {
diff --git a/js/genkit/tests/prompts/test.variant.prompt b/js/genkit/tests/prompts/test.variant.prompt
new file mode 100644
index 000000000..b307d7e53
--- /dev/null
+++ b/js/genkit/tests/prompts/test.variant.prompt
@@ -0,0 +1,6 @@
+---
+description: a prompt variant in a file
+config: 
+  temperature: 13
+---
+Hello from a variant of the hello prompt
diff --git a/js/genkit/tests/prompts_test.ts b/js/genkit/tests/prompts_test.ts
index 7b4f32283..bb2588534 100644
--- a/js/genkit/tests/prompts_test.ts
+++ b/js/genkit/tests/prompts_test.ts
@@ -735,6 +735,17 @@ describe('prompt', () => {
     );
   });
 
+  it('loads a varaint from from the folder', async () => {
+    const testPrompt = ai.prompt('test', { variant: 'variant' }); // see tests/prompts folder
+
+    const { text } = await testPrompt();
+
+    assert.strictEqual(
+      text,
+      'Echo: Hello from a variant of the hello prompt\n; config: {"temperature":13}'
+    );
+  });
+
   it('returns a ref to functional prompts', async () => {
     ai.definePrompt(
       {