diff --git a/js/core/src/tracing.ts b/js/core/src/tracing.ts index dab2ffdf4..5b981bbe1 100644 --- a/js/core/src/tracing.ts +++ b/js/core/src/tracing.ts @@ -35,18 +35,27 @@ let nodeOtelConfig: TelemetryConfig | null = null; const instrumentationKey = '__GENKIT_TELEMETRY_INSTRUMENTED'; -export function ensureBasicTelemetryInstrumentation() { +export async function ensureBasicTelemetryInstrumentation() { if (global[instrumentationKey]) { - return; + return await global[instrumentationKey]; } - enableTelemetry({}); - global[instrumentationKey] = true; + await enableTelemetry({}); } /** * Enables tracing and metrics open telemetry configuration. */ -export function enableTelemetry(telemetryConfig: TelemetryConfig) { +export async function enableTelemetry( + telemetryConfig: TelemetryConfig | Promise +) { + global[instrumentationKey] = + telemetryConfig instanceof Promise ? telemetryConfig : Promise.resolve(); + + telemetryConfig = + telemetryConfig instanceof Promise + ? await telemetryConfig + : telemetryConfig; + nodeOtelConfig = telemetryConfig || {}; const processors: SpanProcessor[] = [createTelemetryServerProcessor()]; diff --git a/js/core/src/tracing/instrumentation.ts b/js/core/src/tracing/instrumentation.ts index f11cd675a..84f063f70 100644 --- a/js/core/src/tracing/instrumentation.ts +++ b/js/core/src/tracing/instrumentation.ts @@ -45,7 +45,7 @@ export async function newTrace( }, fn: (metadata: SpanMetadata, rootSpan: ApiSpan) => Promise ) { - ensureBasicTelemetryInstrumentation(); + await ensureBasicTelemetryInstrumentation(); const traceMetadata: TraceMetadata = traceMetadataAls.getStore() || { paths: new Set(), timestamp: performance.now(), @@ -78,7 +78,7 @@ export async function runInNewSpan( }, fn: (metadata: SpanMetadata, otSpan: ApiSpan, isRoot: boolean) => Promise ): Promise { - ensureBasicTelemetryInstrumentation(); + await ensureBasicTelemetryInstrumentation(); const tracer = trace.getTracer(TRACER_NAME, TRACER_VERSION); const parentStep = spanMetadataAls.getStore(); diff --git a/js/plugins/google-cloud/src/index.ts b/js/plugins/google-cloud/src/index.ts index 8112d6ff4..4c67643c2 100644 --- a/js/plugins/google-cloud/src/index.ts +++ b/js/plugins/google-cloud/src/index.ts @@ -23,13 +23,15 @@ import { GcpOpenTelemetry } from './gcpOpenTelemetry.js'; import { TelemetryConfigs } from './telemetry/defaults.js'; import { GcpTelemetryConfig, GcpTelemetryConfigOptions } from './types.js'; -export async function enableGoogleCloudTelemetry( +export function enableGoogleCloudTelemetry( options?: GcpTelemetryConfigOptions ) { - const pluginConfig = await configureGcpPlugin(options); - - enableTelemetry(await new GcpOpenTelemetry(pluginConfig).getConfig()); - logger.init(await new GcpLogger(pluginConfig).getLogger(getCurrentEnv())); + return enableTelemetry( + configureGcpPlugin(options).then(async (pluginConfig) => { + logger.init(await new GcpLogger(pluginConfig).getLogger(getCurrentEnv())); + return new GcpOpenTelemetry(pluginConfig).getConfig(); + }) + ); } /** diff --git a/js/plugins/google-cloud/tests/traces_test.ts b/js/plugins/google-cloud/tests/traces_test.ts index 470cb1251..654ce5f5f 100644 --- a/js/plugins/google-cloud/tests/traces_test.ts +++ b/js/plugins/google-cloud/tests/traces_test.ts @@ -188,7 +188,7 @@ describe('GoogleCloudTracing', () => { ); const testFlow = createFlow(ai, 'modelFlow', async () => { return run('runFlow', async () => { - ai.generate({ + await ai.generate({ model: echoModel, prompt: 'Testing model telemetry', }); @@ -199,10 +199,11 @@ describe('GoogleCloudTracing', () => { const spans = await getExportedSpans(); - assert.equal(spans[0].name, 'runFlow'); - assert.equal(spans[1].name, 'modelFlow'); - assert.equal(spans[2].name, 'echoModel'); - assert.equal(spans[2].attributes['genkit/model'], 'echoModel'); + assert.equal(spans[0].name, 'echoModel'); + assert.equal(spans[0].attributes['genkit/model'], 'echoModel'); + assert.equal(spans[1].name, 'generate'); + assert.equal(spans[2].name, 'runFlow'); + assert.equal(spans[3].name, 'modelFlow'); }); it('attaches additional span', async () => {