diff --git a/js/ai/src/prompt.ts b/js/ai/src/prompt.ts index 7dc1160bd..d4b57d10c 100644 --- 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, renderOptions: PromptGenerateOptions | undefined ): Promise => { - 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) => {