fix(js/ai/prompt): fix prompt loading in node 20, make sure-sure that…

… .prompt loading is lazy (#1700)

.github/workflows/builder.yml

@@ -27,10 +27,10 @@ jobs:
     steps:
     - uses: actions/checkout@v3
     - uses: pnpm/action-setup@v3
-    - name: Set up node v21
+    - name: Set up node v20
       uses: actions/setup-node@v4
       with:
-        node-version: 21.x
+        node-version: 20.x
         cache: 'pnpm'
     - name: Install dependencies
       run: pnpm install

js/ai/src/prompt.ts

@@ -25,7 +25,7 @@ import {
   stripUndefinedProps,
   z,
 } from '@genkit-ai/core';
-import { LazyPromise } from '@genkit-ai/core/async';
+import { lazy } from '@genkit-ai/core/async';
 import { logger } from '@genkit-ai/core/logging';
 import { Registry } from '@genkit-ai/core/registry';
 import { toJsonSchema } from '@genkit-ai/core/schema';
@@ -311,8 +311,8 @@ function definePromptAsync<
       },
     });
   };
-  const rendererActionConfig = optionsPromise.then(
-    (options: PromptConfig<I, O, CustomOptions>) => {
+  const rendererActionConfig = lazy(() =>
+    optionsPromise.then((options: PromptConfig<I, O, CustomOptions>) => {
       const metadata = promptMetadata(options);
       return {
         name: `${options.name}${options.variant ? `.${options.variant}` : ''}`,
@@ -330,17 +330,21 @@ function definePromptAsync<
           );
         },
       } as ActionAsyncParams<any, any, any>;
-    }
+    })
   );
   const rendererAction = defineActionAsync(
     registry,
     'prompt',
     name,
-    rendererActionConfig
+    rendererActionConfig,
+    (action) => {
+      (action as PromptAction<I>).__executablePrompt =
+        executablePrompt as never as ExecutablePrompt<z.infer<I>>;
+    }
   ) as Promise<PromptAction<I>>;
 
-  const executablePromptActionConfig = optionsPromise.then(
-    (options: PromptConfig<I, O, CustomOptions>) => {
+  const executablePromptActionConfig = lazy(() =>
+    optionsPromise.then((options: PromptConfig<I, O, CustomOptions>) => {
       const metadata = promptMetadata(options);
       return {
         name: `${options.name}${options.variant ? `.${options.variant}` : ''}`,
@@ -359,14 +363,18 @@ function definePromptAsync<
           });
         },
       } as ActionAsyncParams<any, any, any>;
-    }
+    })
   );
 
-  const executablePromptAction = defineActionAsync(
+  defineActionAsync(
     registry,
     'executable-prompt',
     name,
-    executablePromptActionConfig
+    executablePromptActionConfig,
+    (action) => {
+      (action as ExecutablePromptAction<I>).__executablePrompt =
+        executablePrompt as never as ExecutablePrompt<z.infer<I>>;
+    }
   ) as Promise<ExecutablePromptAction<I>>;
 
   const executablePrompt = wrapInExecutablePrompt(
@@ -375,17 +383,6 @@ function definePromptAsync<
     rendererAction
   );
 
-  executablePromptAction.then((action) => {
-    action.__executablePrompt = executablePrompt as never as ExecutablePrompt<
-      z.infer<I>
-    >;
-  });
-  rendererAction.then((action) => {
-    action.__executablePrompt = executablePrompt as never as ExecutablePrompt<
-      z.infer<I>
-    >;
-  });
-
   return executablePrompt;
 }
 
@@ -670,26 +667,27 @@ export function loadPromptFolder(
       recursive: true,
     });
     for (const dirEnt of dirEnts) {
+      const parentPath = dirEnt.parentPath ?? dirEnt.path;
       if (dirEnt.isFile() && dirEnt.name.endsWith('.prompt')) {
         if (dirEnt.name.startsWith('_')) {
           const partialName = dirEnt.name.substring(1, dirEnt.name.length - 7);
           definePartial(
             registry,
             partialName,
-            readFileSync(join(dirEnt.path, dirEnt.name), {
+            readFileSync(join(parentPath, dirEnt.name), {
               encoding: 'utf8',
             })
           );
           logger.debug(
-            `Registered Dotprompt partial "${partialName}" from "${join(dirEnt.parentPath, dirEnt.name)}"`
+            `Registered Dotprompt partial "${partialName}" from "${join(parentPath, dirEnt.name)}"`
           );
         } else {
           // If this prompt is in a subdirectory, we need to include that
           // in the namespace to prevent naming conflicts.
           let prefix = '';
           let name = dirEnt.name;
-          if (promptsPath !== dirEnt.parentPath) {
-            prefix = dirEnt.parentPath.replace(`${promptsPath}/`, '') + '/';
+          if (promptsPath !== parentPath) {
+            prefix = parentPath.replace(`${promptsPath}/`, '') + '/';
           }
           loadPrompt(registry, promptsPath, name, prefix, ns);
         }
@@ -736,7 +734,7 @@ function loadPrompt(
     // We use a lazy promise here because we only want prompt loaded when it's first used.
     // This is important because otherwise the loading may happen before the user has configured
     // all the schemas, etc., which will result in dotprompt.renderMetadata errors.
-    new LazyPromise<PromptConfig>(async (resolvePromptConfig) => {
+    lazy(async () => {
       const promptMetadata =
         await registry.dotprompt.renderMetadata(parsedPrompt);
       if (variant) {
@@ -751,7 +749,7 @@ function loadPrompt(
         delete promptMetadata.input.schema.description;
       }
 
-      resolvePromptConfig({
+      return {
         name: registryDefinitionKey(name, variant ?? undefined, ns),
         model: promptMetadata.model,
         config: promptMetadata.config,
@@ -776,7 +774,7 @@ function loadPrompt(
         toolChoice: promptMetadata.raw?.['toolChoice'],
         returnToolRequests: promptMetadata.raw?.['returnToolRequests'],
         messages: parsedPrompt.template,
-      });
+      };
     })
   );
 }

js/core/src/action.ts

@@ -16,6 +16,7 @@
 
 import { JSONSchema7 } from 'json-schema';
 import * as z from 'zod';
+import { lazy } from './async.js';
 import { ActionContext, getContext, runWithContext } from './context.js';
 import { ActionType, Registry } from './registry.js';
 import { parseSchema } from './schema.js';
@@ -446,24 +447,28 @@ export function defineActionAsync<
         pluginId: string;
         actionId: string;
       },
-  config: PromiseLike<ActionAsyncParams<I, O, S>>
+  config: PromiseLike<ActionAsyncParams<I, O, S>>,
+  onInit?: (action: Action<I, O, S>) => void
 ): PromiseLike<Action<I, O, S>> {
   const actionName =
     typeof name === 'string' ? name : `${name.pluginId}/${name.actionId}`;
-  const actionPromise = config.then((resolvedConfig) => {
-    const act = action(
-      registry,
-      resolvedConfig,
-      async (i: I, options): Promise<z.infer<O>> => {
-        await registry.initializeAllPlugins();
-        return await runInActionRuntimeContext(registry, () =>
-          resolvedConfig.fn(i, options)
-        );
-      }
-    );
-    act.__action.actionType = actionType;
-    return act;
-  });
+  const actionPromise = lazy(() =>
+    config.then((resolvedConfig) => {
+      const act = action(
+        registry,
+        resolvedConfig,
+        async (i: I, options): Promise<z.infer<O>> => {
+          await registry.initializeAllPlugins();
+          return await runInActionRuntimeContext(registry, () =>
+            resolvedConfig.fn(i, options)
+          );
+        }
+      );
+      act.__action.actionType = actionType;
+      onInit?.(act);
+      return act;
+    })
+  );
   registry.registerActionAsync(actionType, actionName, actionPromise);
   return actionPromise;
 }

js/core/src/async.ts

@@ -113,3 +113,14 @@ export class LazyPromise<T> implements PromiseLike<T> {
     return this.promise.then(onfulfilled, onrejected);
   }
 }
+
+/** Lazily call the provided function to resolve the LazyPromise. */
+export function lazy<T>(fn: () => T | PromiseLike<T>): PromiseLike<T> {
+  return new LazyPromise<T>((resolve, reject) => {
+    try {
+      resolve(fn());
+    } catch (e) {
+      reject(e);
+    }
+  });
+}

js/genkit/tests/prompts/badSchemaRef.prompt

@@ -0,0 +1,9 @@
+---
+model: googleai/gemini-1.5-flash
+input:
+  schema: badSchemaRef1
+output:
+  schema: badSchemaRef2
+---
+
+doesn't matter

js/genkit/tests/prompts_test.ts

@@ -924,7 +924,7 @@ describe('prompt', () => {
     });
   });
 
-  it.only('loads from from the sub folder', async () => {
+  it('loads from from the sub folder', async () => {
     const testPrompt = ai.prompt('sub/test'); // see tests/prompts/sub folder
 
     const { text } = await testPrompt();
@@ -1043,6 +1043,14 @@ describe('prompt', () => {
     );
   });
 
+  it('lazily resolved schema refs', async () => {
+    const prompt = ai.prompt('badSchemaRef');
+
+    await assert.rejects(prompt.render({ foo: 'bar' }), (e: Error) =>
+      e.message.includes("NOT_FOUND: Schema 'badSchemaRef1' not found")
+    );
+  });
+
   it('loads a varaint from from the folder', async () => {
     const testPrompt = ai.prompt('test', { variant: 'variant' }); // see tests/prompts folder