feat(js/ai/prompt): added prepare fn option to definePrompt

js/ai/src/index.ts

@@ -81,10 +81,14 @@ export {
   isExecutablePrompt,
   loadPromptFolder,
   prompt,
+  type DocsResolver,
   type ExecutablePrompt,
+  type MessagesResolver,
+  type PartsResolver,
   type PromptAction,
   type PromptConfig,
   type PromptGenerateOptions,
+  type PromptPrepare,
 } from './prompt.js';
 export {
   rerank,

js/ai/src/prompt.ts

@@ -125,6 +125,7 @@ export interface PromptConfig<
   toolChoice?: ToolChoice;
   use?: ModelMiddleware[];
   context?: ActionContext;
+  prepare?: PromptPrepare<I, O, CustomOptions>;
 }
 
 /**
@@ -208,6 +209,34 @@ export type DocsResolver<I, S = any> = (
   }
 ) => DocumentData[] | Promise<DocumentData[]>;
 
+/**
+ * A function that can be passes to the prompt that can produce additional propmpt render options
+ * like context, docs, or any other propmpt option really:
+ *
+ * ```ts
+ * const ragPrompt = ai.definePrompt({
+ *   prepare: async (input, {context}) => ({
+ *     context: {...context, userInfo: await fetchUserInfo(context.auth.uid)},
+ *     docs: await myRetriever({query: input.query}),
+ *   })
+ * })
+ * ```
+ */
+export type PromptPrepare<
+  I extends z.ZodTypeAny = z.ZodTypeAny,
+  O extends z.ZodTypeAny = z.ZodTypeAny,
+  CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
+  S = any,
+> = (
+  input: I,
+  options: {
+    state?: S;
+    context: ActionContext;
+  }
+) =>
+  | PromptGenerateOptions<O, CustomOptions>
+  | PromiseLike<PromptGenerateOptions<O, CustomOptions>>;
+
 interface PromptCache {
   userPrompt?: PromptFunction;
   system?: PromptFunction;
@@ -250,10 +279,22 @@ function definePromptAsync<
     renderOptions: PromptGenerateOptions<O, CustomOptions> | undefined
   ): Promise<GenerateOptions> => {
     const messages: MessageData[] = [];
-    renderOptions = { ...renderOptions }; // make a copy, we will be trimming
     const session = getCurrentSession(registry);
     const resolvedOptions = await optionsPromise;
 
+    // if prepare option is set, invoke it and merge with renderOptions
+    const preparedOptions = resolvedOptions?.prepare
+      ? await resolvedOptions?.prepare(input, {
+          state: session?.state,
+          context: renderOptions?.context || getContext(registry) || {},
+        })
+      : {};
+    // make a copy of renderOptions, we will be trimming
+    renderOptions = {
+      ...renderOptions,
+      ...preparedOptions,
+    };
+
     // order of these matters:
     await renderSystemPrompt(
       registry,

js/ai/tests/prompt/prompt_test.ts

@@ -753,6 +753,48 @@ describe('prompt', () => {
         tools: ['toolA'],
       },
     },
+    {
+      name: 'invoked prepare fn and merges with render options',
+      prompt: {
+        model: 'echoModel',
+        name: 'prompt1',
+        config: { banana: 'ripe' },
+        input: { schema: z.object({ name: z.string() }) },
+        prompt: 'hello {{@foo}} {{@baz}} ({{@state.name}})',
+        tools: ['toolA'],
+        prepare: async (input, { context }) => ({
+          context: { ...context, foo: 'bar' },
+          docs: [Document.fromText('doc txt')],
+        }),
+      },
+      input: { name: 'foo' },
+      state: { name: 'bar' },
+      inputOptions: { config: { temperature: 11 }, context: { baz: 'aux' } },
+      wantTextOutput:
+        'Echo: hello bar aux (bar),\n' +
+        '\n' +
+        'Use the following information to complete your task:\n' +
+        '\n' +
+        '- [0]: doc txt\n' +
+        '\n' +
+        '; config: {"banana":"ripe","temperature":11}',
+      wantRendered: {
+        context: {
+          baz: 'aux',
+          foo: 'bar',
+        },
+        docs: [Document.fromText('doc txt')],
+        config: {
+          banana: 'ripe',
+          temperature: 11,
+        },
+        messages: [
+          { content: [{ text: 'hello bar aux (bar)' }], role: 'user' },
+        ],
+        model: 'echoModel',
+        tools: ['toolA'],
+      },
+    },
   ];
 
   basicTests = basicTests.find((t) => t.only)

js/genkit/src/index.ts

@@ -47,6 +47,7 @@ export {
   indexerRef,
   rerankerRef,
   retrieverRef,
+  type DocsResolver,
   type DocumentData,
   type EmbedderAction,
   type EmbedderArgument,
@@ -79,13 +80,16 @@ export {
   type LlmStats,
   type MediaPart,
   type MessageData,
+  type MessagesResolver,
   type ModelArgument,
   type ModelReference,
   type ModelRequest,
   type ModelResponseData,
   type Part,
+  type PartsResolver,
   type PromptAction,
   type PromptConfig,
+  type PromptPrepare,
   type RankedDocument,
   type RerankerAction,
   type RerankerArgument,