fix: add span for prompt template rendering (#2165)

Co-authored-by: Pavel Jbanov <[email protected]>

Author: MichaelDoyle

js/ai/src/prompt.ts

@@ -29,7 +29,8 @@ import { lazy } from '@genkit-ai/core/async';
 import { logger } from '@genkit-ai/core/logging';
 import type { Registry } from '@genkit-ai/core/registry';
 import { toJsonSchema } from '@genkit-ai/core/schema';
-import type { Message as DpMessage, PromptFunction } from 'dotprompt';
+import { SPAN_TYPE_ATTR, runInNewSpan } from '@genkit-ai/core/tracing';
+import { Message as DpMessage, PromptFunction } from 'dotprompt';
 import { existsSync, readFileSync, readdirSync } from 'fs';
 import { basename, join, resolve } from 'path';
 import type { DocumentData } from './document.js';
@@ -252,72 +253,87 @@ function definePromptAsync<
     input: z.infer<I>,
     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;
-
-    // order of these matters:
-    await renderSystemPrompt(
-      registry,
-      session,
-      input,
-      messages,
-      resolvedOptions,
-      promptCache,
-      renderOptions
-    );
-    await renderMessages(
-      registry,
-      session,
-      input,
-      messages,
-      resolvedOptions,
-      renderOptions,
-      promptCache
-    );
-    await renderUserPrompt(
+    return await runInNewSpan(
       registry,
-      session,
-      input,
-      messages,
-      resolvedOptions,
-      promptCache,
-      renderOptions
-    );
-
-    let docs: DocumentData[] | undefined;
-    if (typeof resolvedOptions.docs === 'function') {
-      docs = await resolvedOptions.docs(input, {
-        state: session?.state,
-        context: renderOptions?.context || getContext(registry) || {},
-      });
-    } else {
-      docs = resolvedOptions.docs;
-    }
-
-    const opts: GenerateOptions = stripUndefinedProps({
-      model: resolvedOptions.model,
-      maxTurns: resolvedOptions.maxTurns,
-      messages,
-      docs,
-      tools: resolvedOptions.tools,
-      returnToolRequests: resolvedOptions.returnToolRequests,
-      toolChoice: resolvedOptions.toolChoice,
-      context: resolvedOptions.context,
-      output: resolvedOptions.output,
-      use: resolvedOptions.use,
-      ...stripUndefinedProps(renderOptions),
-      config: {
-        ...resolvedOptions?.config,
-        ...renderOptions?.config,
+      {
+        metadata: {
+          name: 'render',
+          input,
+        },
+        labels: {
+          [SPAN_TYPE_ATTR]: 'promptTemplate',
+        },
       },
-    });
-    // if config is empty and it was not explicitly passed in, we delete it, don't want {}
-    if (Object.keys(opts.config).length === 0 && !renderOptions?.config) {
-      delete opts.config;
-    }
-    return opts;
+      async (metadata) => {
+        const messages: MessageData[] = [];
+        renderOptions = { ...renderOptions }; // make a copy, we will be trimming
+        const session = getCurrentSession(registry);
+        const resolvedOptions = await optionsPromise;
+
+        // order of these matters:
+        await renderSystemPrompt(
+          registry,
+          session,
+          input,
+          messages,
+          resolvedOptions,
+          promptCache,
+          renderOptions
+        );
+        await renderMessages(
+          registry,
+          session,
+          input,
+          messages,
+          resolvedOptions,
+          renderOptions,
+          promptCache
+        );
+        await renderUserPrompt(
+          registry,
+          session,
+          input,
+          messages,
+          resolvedOptions,
+          promptCache,
+          renderOptions
+        );
+
+        let docs: DocumentData[] | undefined;
+        if (typeof resolvedOptions.docs === 'function') {
+          docs = await resolvedOptions.docs(input, {
+            state: session?.state,
+            context: renderOptions?.context || getContext(registry) || {},
+          });
+        } else {
+          docs = resolvedOptions.docs;
+        }
+
+        const opts: GenerateOptions = stripUndefinedProps({
+          model: resolvedOptions.model,
+          maxTurns: resolvedOptions.maxTurns,
+          messages,
+          docs,
+          tools: resolvedOptions.tools,
+          returnToolRequests: resolvedOptions.returnToolRequests,
+          toolChoice: resolvedOptions.toolChoice,
+          context: resolvedOptions.context,
+          output: resolvedOptions.output,
+          use: resolvedOptions.use,
+          ...stripUndefinedProps(renderOptions),
+          config: {
+            ...resolvedOptions?.config,
+            ...renderOptions?.config,
+          },
+        });
+        // if config is empty and it was not explicitly passed in, we delete it, don't want {}
+        if (Object.keys(opts.config).length === 0 && !renderOptions?.config) {
+          delete opts.config;
+        }
+        metadata.output = opts;
+        return opts;
+      }
+    );
   };
   const rendererActionConfig = lazy(() =>
     optionsPromise.then((options: PromptConfig<I, O, CustomOptions>) => {
@@ -432,9 +448,25 @@ function wrapInExecutablePrompt<
     input?: I,
     opts?: PromptGenerateOptions<O, CustomOptions>
   ): Promise<GenerateResponse<z.infer<O>>> => {
-    return generate(registry, {
-      ...(await renderOptionsFn(input, opts)),
-    });
+    return await runInNewSpan(
+      registry,
+      {
+        metadata: {
+          name: (await rendererAction).__action.name,
+          input,
+        },
+        labels: {
+          [SPAN_TYPE_ATTR]: 'dotprompt',
+        },
+      },
+      async (metadata) => {
+        const output = await generate(registry, {
+          ...(await renderOptionsFn(input, opts)),
+        });
+        metadata.output = output;
+        return output;
+      }
+    );
   }) as ExecutablePrompt<z.infer<I>, O, CustomOptions>;
 
   executablePrompt.render = async (

js/genkit/src/genkit.ts

@@ -38,6 +38,7 @@ import {
   type GenerateOptions,
   type GenerateRequest,
   type GenerateResponse,
+  type GenerateResponseChunk,
   type GenerateResponseData,
   type GenerateStreamOptions,
   type GenerateStreamResponse,
@@ -110,11 +111,13 @@ import {
   type StreamingCallback,
   type z,
 } from '@genkit-ai/core';
+import { Channel } from '@genkit-ai/core/async';
 import type { HasRegistry } from '@genkit-ai/core/registry';
 import type { BaseEvalDataPointSchema } from './evaluator.js';
 import { logger } from './logging.js';
 import type { GenkitPlugin } from './plugin.js';
 import { Registry, type ActionType } from './registry.js';
+import { SPAN_TYPE_ATTR, runInNewSpan } from './tracing.js';
 
 /**
  * @deprecated use `ai.definePrompt({messages: fn})`
@@ -254,6 +257,7 @@ export class Genkit implements HasRegistry {
     options?: { variant?: string }
   ): ExecutablePrompt<z.infer<I>, O, CustomOptions> {
     return this.wrapExecutablePromptPromise(
+      `${name}${options?.variant ? `.${options?.variant}` : ''}`,
       prompt(this.registry, name, {
         ...options,
         dir: this.options.promptDir ?? './prompts',
@@ -265,7 +269,10 @@ export class Genkit implements HasRegistry {
     I extends z.ZodTypeAny = z.ZodTypeAny,
     O extends z.ZodTypeAny = z.ZodTypeAny,
     CustomOptions extends z.ZodTypeAny = z.ZodTypeAny,
-  >(promise: Promise<ExecutablePrompt<z.infer<I>, O, CustomOptions>>) {
+  >(
+    name: string,
+    promise: Promise<ExecutablePrompt<z.infer<I>, O, CustomOptions>>
+  ) {
     const executablePrompt = (async (
       input?: I,
       opts?: PromptGenerateOptions<O, CustomOptions>
@@ -286,13 +293,39 @@ export class Genkit implements HasRegistry {
       input?: I,
       opts?: PromptGenerateOptions<O, CustomOptions>
     ): GenerateStreamResponse<O> => {
-      return this.generateStream(
-        promise.then((action) =>
-          action.render(input, {
-            ...opts,
-          })
-        )
+      let channel = new Channel<GenerateResponseChunk>();
+
+      const generated = runInNewSpan(
+        this.registry,
+        {
+          metadata: {
+            name,
+            input,
+          },
+          labels: {
+            [SPAN_TYPE_ATTR]: 'dotprompt',
+          },
+        },
+        () =>
+          generate<O, CustomOptions>(
+            this.registry,
+            promise.then((action) =>
+              action.render(input, {
+                ...opts,
+                onChunk: (chunk) => channel.send(chunk),
+              })
+            )
+          )
+      );
+      generated.then(
+        () => channel.close(),
+        (err) => channel.error(err)
       );
+
+      return {
+        response: generated,
+        stream: channel,
+      };
     };
 
     executablePrompt.asTool = async (): Promise<ToolAction<I, O>> => {