fix: add span for prompt template rendering

firebase · Feb 25, 2025 · 64697ac · 64697ac
1 parent 443561e
commit 64697ac
Showing 1 changed file with 80 additions and 64 deletions.
diff --git a/js/ai/src/prompt.ts b/js/ai/src/prompt.ts
@@ -29,6 +29,7 @@ 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';
+import { runInNewSpan, SPAN_TYPE_ATTR } from '@genkit-ai/core/tracing';
 import { Message as DpMessage, PromptFunction } from 'dotprompt';
 import { existsSync, readdirSync, readFileSync } from 'fs';
 import { basename, join, resolve } from 'path';
@@ -249,72 +250,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(
+    return await runInNewSpan(
       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,
+      {
+        metadata: {
+          name,
+          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>) => {