feat(js): added id and metadata to executable prompts (#3084)

Author: pavelgj

js/ai/src/prompt.ts

@@ -147,6 +147,9 @@ export interface ExecutablePrompt<
   O extends z.ZodTypeAny = z.ZodTypeAny,
   CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
 > {
+  /** Prompt reference. */
+  ref: { name: string; metadata?: Record<string, any> };
+
   /**
    * Generates a response by rendering the prompt template with given user input and then calling the model.
    *
@@ -234,7 +237,8 @@ export function definePrompt<
   return definePromptAsync(
     registry,
     `${options.name}${options.variant ? `.${options.variant}` : ''}`,
-    Promise.resolve(options)
+    Promise.resolve(options),
+    options.metadata
   );
 }
 
@@ -245,7 +249,8 @@ function definePromptAsync<
 >(
   registry: Registry,
   name: string,
-  optionsPromise: PromiseLike<PromptConfig<I, O, CustomOptions>>
+  optionsPromise: PromiseLike<PromptConfig<I, O, CustomOptions>>,
+  metadata?: Record<string, any>
 ): ExecutablePrompt<z.infer<I>, O, CustomOptions> {
   const promptCache = {} as PromptCache;
 
@@ -399,11 +404,13 @@ function definePromptAsync<
     }
   ) as Promise<ExecutablePromptAction<I>>;
 
-  const executablePrompt = wrapInExecutablePrompt(
+  const executablePrompt = wrapInExecutablePrompt({
     registry,
+    name,
     renderOptionsFn,
-    rendererAction
-  );
+    rendererAction,
+    metadata,
+  });
 
   return executablePrompt;
 }
@@ -436,57 +443,64 @@ function wrapInExecutablePrompt<
   I extends z.ZodTypeAny = z.ZodTypeAny,
   O extends z.ZodTypeAny = z.ZodTypeAny,
   CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
->(
-  registry: Registry,
+>(wrapOpts: {
+  registry: Registry;
+  name: string;
   renderOptionsFn: (
     input: z.infer<I>,
     renderOptions: PromptGenerateOptions<O, CustomOptions> | undefined
-  ) => Promise<GenerateOptions>,
-  rendererAction: Promise<PromptAction<I>>
-) {
+  ) => Promise<GenerateOptions>;
+  rendererAction: Promise<PromptAction<I>>;
+  metadata?: Record<string, any>;
+}) {
   const executablePrompt = (async (
     input?: I,
     opts?: PromptGenerateOptions<O, CustomOptions>
   ): Promise<GenerateResponse<z.infer<O>>> => {
     return await runInNewSpan(
-      registry,
+      wrapOpts.registry,
       {
         metadata: {
-          name: (await rendererAction).__action.name,
+          name: (await wrapOpts.rendererAction).__action.name,
           input,
         },
         labels: {
           [SPAN_TYPE_ATTR]: 'dotprompt',
         },
       },
       async (metadata) => {
-        const output = await generate(registry, {
-          ...(await renderOptionsFn(input, opts)),
+        const output = await generate(wrapOpts.registry, {
+          ...(await wrapOpts.renderOptionsFn(input, opts)),
         });
         metadata.output = output;
         return output;
       }
     );
   }) as ExecutablePrompt<z.infer<I>, O, CustomOptions>;
 
+  executablePrompt.ref = { name: wrapOpts.name, metadata: wrapOpts.metadata };
+
   executablePrompt.render = async (
     input?: I,
     opts?: PromptGenerateOptions<O, CustomOptions>
   ): Promise<GenerateOptions<O, CustomOptions>> => {
     return {
-      ...(await renderOptionsFn(input, opts)),
+      ...(await wrapOpts.renderOptionsFn(input, opts)),
     } as GenerateOptions<O, CustomOptions>;
   };
 
   executablePrompt.stream = (
     input?: I,
     opts?: PromptGenerateOptions<O, CustomOptions>
   ): GenerateStreamResponse<z.infer<O>> => {
-    return generateStream(registry, renderOptionsFn(input, opts));
+    return generateStream(
+      wrapOpts.registry,
+      wrapOpts.renderOptionsFn(input, opts)
+    );
   };
 
   executablePrompt.asTool = async (): Promise<ToolAction<I, O>> => {
-    return (await rendererAction) as unknown as ToolAction<I, O>;
+    return (await wrapOpts.rendererAction) as unknown as ToolAction<I, O>;
   };
   return executablePrompt;
 }

js/genkit/src/genkit.ts

@@ -280,6 +280,8 @@ export class Genkit implements HasRegistry {
       return (await promise)(input, opts);
     }) as ExecutablePrompt<z.infer<I>, O, CustomOptions>;
 
+    executablePrompt.ref = { name };
+
     executablePrompt.render = async (
       input?: I,
       opts?: PromptGenerateOptions<O, CustomOptions>

js/genkit/tests/prompts_test.ts

@@ -72,6 +72,36 @@ describe('definePrompt', () => {
     defineEchoModel(ai);
   });
 
+  it('should define the prompt', async () => {
+    const prompt = ai.definePrompt({
+      name: 'hi',
+      metadata: { foo: 'bar' },
+      input: {
+        schema: z.object({
+          name: z.string(),
+        }),
+      },
+      messages: async (input) => {
+        return [
+          {
+            role: 'user',
+            content: [{ text: `hi ${input.name}` }],
+          },
+        ];
+      },
+    });
+
+    assert.deepStrictEqual(prompt.ref, {
+      name: 'hi',
+      metadata: { foo: 'bar' },
+    });
+
+    const lookedUpPrompt = ai.prompt('hi');
+    // This is a known limitation -- prompt lookup is async under the hood,
+    // so we can't actually get the metadata...
+    assert.deepStrictEqual(lookedUpPrompt.ref, { name: 'hi' }); // ideally metadatashould be: { foo: 'bar' }
+  });
+
   it('should apply middleware to a prompt call', async () => {
     const prompt = ai.definePrompt({
       name: 'hi',

js/plugins/mcp/src/util/prompts.ts

@@ -96,12 +96,20 @@ function createExecutablePrompt<
     options?: PromptGenerateOptions<any, any>;
   }
 ): ExecutablePrompt<z.infer<I>, O, CustomOptions> {
-  const callPrompt = async (
+  const callPrompt = (async (
     input?: z.infer<I>,
     opts?: PromptGenerateOptions<O, CustomOptions>
   ): Promise<GenerateResponse<z.infer<O>>> => {
     logger.debug(`[MCP] Calling MCP prompt ${params.name}/${prompt.name}`);
     return params.ai.generate(callPrompt.render(input, opts));
+  }) as ExecutablePrompt<z.infer<I>, O, CustomOptions>;
+
+  callPrompt.ref = {
+    name: prompt.name,
+    metadata: {
+      description: prompt.description,
+      arguments: prompt.arguments,
+    },
   };
 
   callPrompt.stream = (
@@ -137,7 +145,7 @@ function createExecutablePrompt<
     });
   };
 
-  return callPrompt as ExecutablePrompt<z.infer<I>, O, CustomOptions>;
+  return callPrompt;
 }
 
 /**

js/plugins/mcp/tests/host_test.ts

@@ -300,6 +300,14 @@ describe('createMcpHost', () => {
       // Add a prompt to the first transport
       fakeTransport.prompts.push({
         name: 'testPrompt1',
+        arguments: [
+          {
+            name: 'foo',
+            description: 'foo arg',
+            required: false,
+          },
+        ],
+        description: 'descr',
       });
       let activePrompts = await clientHost.getActivePrompts(ai);
       assert.strictEqual(activePrompts.length, 1);
@@ -326,17 +334,36 @@ describe('createMcpHost', () => {
       });
 
       activePrompts = await clientHost.getActivePrompts(ai);
-      assert.strictEqual(activePrompts.length, 2);
+      assert.deepStrictEqual(activePrompts[0].ref.metadata, {
+        arguments: [
+          {
+            description: 'foo arg',
+            name: 'foo',
+            required: false,
+          },
+        ],
+        description: 'descr',
+      });
+      assert.deepStrictEqual(
+        activePrompts.map((p) => p.ref.name),
+        ['testPrompt1', 'testPrompt2']
+      );
 
       // Disable the first server
       await clientHost.disable('test-server');
       activePrompts = await clientHost.getActivePrompts(ai);
-      assert.strictEqual(activePrompts.length, 1);
+      assert.deepStrictEqual(
+        activePrompts.map((p) => p.ref.name),
+        ['testPrompt2']
+      );
 
       // Enable the first server again
       await clientHost.enable('test-server');
       activePrompts = await clientHost.getActivePrompts(ai);
-      assert.strictEqual(activePrompts.length, 2);
+      assert.deepStrictEqual(
+        activePrompts.map((p) => p.ref.name),
+        ['testPrompt1', 'testPrompt2']
+      );
     });
 
     it('should execute prompt', async () => {