weflayr 0.20.6 → 0.22.0

package/tests/instrumentation/google/get_embeddings_telemetry.test.js

@@ -0,0 +1,128 @@
+'use strict';
+
+// Tests for the bundled Google GenAI (Gemini) embeddings telemetry.
+//
+// The @google/genai `Models` class is stubbed with the prototype method
+// `embedContentInternal` (the stable patch point that the public `embedContent`
+// arrow delegates to). The Gemini Developer API returns no token usage, so the
+// billable unit is the input character count. `instrumentEmbeddings(Models,
+// patched)` wraps the method and pushes its restore tuple onto `patched`.
+
+const test = require('node:test');
+const assert = require('node:assert/strict');
+
+const { trace, SpanStatusCode } = require('@opentelemetry/api');
+const {
+  BasicTracerProvider,
+  InMemorySpanExporter,
+  SimpleSpanProcessor,
+} = require('@opentelemetry/sdk-trace-base');
+
+const {
+  instrumentEmbeddings,
+} = require('../../../src/instrumentation/google/get_embeddings_telemetry');
+
+const exporter = new InMemorySpanExporter();
+trace.setGlobalTracerProvider(
+  new BasicTracerProvider({ spanProcessors: [new SimpleSpanProcessor(exporter)] })
+);
+
+class StubModels {
+  embedContentInternal(params = {}) {
+    if (params.model === 'boom') return Promise.reject(new Error('api error'));
+    return Promise.resolve({ embeddings: [{ values: [0.1] }] });
+  }
+}
+
+function instrumented() {
+  const patched = [];
+  instrumentEmbeddings(StubModels, patched);
+  return () => {
+    for (const [owner, methodName, original] of patched) owner[methodName] = original;
+    exporter.reset();
+  };
+}
+
+test('embeddings emits genai span with character count and stripped model', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubModels();
+    await client.embedContentInternal({
+      model: 'models/text-embedding-004',
+      contents: 'hello world',
+    });
+
+    const [span] = exporter.getFinishedSpans();
+    assert.equal(span.name, 'embeddings');
+    // No apiClient on the stub → Gemini Developer API.
+    assert.equal(span.attributes['gen_ai.system'], 'gemini');
+    assert.equal(span.attributes['gen_ai.operation.name'], 'embeddings');
+    // The "models/" prefix is stripped to match the chat path.
+    assert.equal(span.attributes['gen_ai.request.model'], 'text-embedding-004');
+    assert.equal(span.attributes['gen_ai.usage.input_characters'], 'hello world'.length);
+    assert.deepEqual(JSON.parse(span.attributes['gen_ai.input.messages']), [
+      { role: 'user', parts: [{ type: 'text', content: 'hello world' }] },
+    ]);
+    assert.equal(span.status.code, SpanStatusCode.UNSET);
+  } finally {
+    teardown();
+  }
+});
+
+test('embeddings counts characters across list/part contents', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubModels();
+    await client.embedContentInternal({
+      model: 'text-embedding-004',
+      contents: ['ab', { text: 'cde' }, { parts: [{ text: 'fg' }] }],
+    });
+
+    const [span] = exporter.getFinishedSpans();
+    assert.equal(span.attributes['gen_ai.usage.input_characters'], 7);
+  } finally {
+    teardown();
+  }
+});
+
+test('embeddings system is vertex_ai when client targets vertex', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubModels();
+    // The Vertex backend is read off the api client.
+    client.apiClient = { isVertexAI: () => true };
+    await client.embedContentInternal({ model: 'text-embedding-004', contents: 'hi' });
+
+    const [span] = exporter.getFinishedSpans();
+    assert.equal(span.attributes['gen_ai.system'], 'vertex_ai');
+  } finally {
+    teardown();
+  }
+});
+
+test('embeddings records error status on failure', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubModels();
+    await assert.rejects(
+      () => client.embedContentInternal({ model: 'boom', contents: 'x' }),
+      /api error/
+    );
+
+    const [span] = exporter.getFinishedSpans();
+    assert.equal(span.status.code, SpanStatusCode.ERROR);
+  } finally {
+    teardown();
+  }
+});
+
+test('uninstrument restores the original method', async () => {
+  const teardown = instrumented();
+  const client = new StubModels();
+  await client.embedContentInternal({ model: 'text-embedding-004', contents: 'x' });
+  assert.equal(exporter.getFinishedSpans().length, 1);
+  teardown();
+  // After teardown the prototype method is restored (no new span is recorded).
+  await client.embedContentInternal({ model: 'text-embedding-004', contents: 'y' });
+  assert.equal(exporter.getFinishedSpans().length, 0);
+});

package/tests/instrumentation/google/get_reasoning_telemetry.test.js

@@ -0,0 +1,41 @@
+'use strict';
+
+// Tests for the Gemini reasoning (thinking) token complement. It stamps
+// usageMetadata.thoughtsTokenCount onto the chat span as
+// gen_ai.usage.thoughts_token_count so the normalizer can fold the
+// (separately-reported) thinking tokens into the completion count.
+
+const test = require('node:test');
+const assert = require('node:assert/strict');
+
+const {
+  setReasoningTelemetry,
+} = require('../../../src/instrumentation/google/get_reasoning_telemetry');
+
+function fakeSpan() {
+  const attrs = {};
+  return {
+    attrs,
+    setAttribute(key, value) {
+      attrs[key] = value;
+    },
+  };
+}
+
+test('stamps thoughts token count', () => {
+  const span = fakeSpan();
+  setReasoningTelemetry(span, { usageMetadata: { thoughtsTokenCount: 369 } });
+  assert.equal(span.attrs['gen_ai.usage.thoughts_token_count'], 369);
+});
+
+test('sets nothing when thoughts is zero', () => {
+  const span = fakeSpan();
+  setReasoningTelemetry(span, { usageMetadata: { thoughtsTokenCount: 0 } });
+  assert.equal(span.attrs['gen_ai.usage.thoughts_token_count'], undefined);
+});
+
+test('is a no-op without usageMetadata', () => {
+  const span = fakeSpan();
+  setReasoningTelemetry(span, {});
+  assert.equal(Object.keys(span.attrs).length, 0);
+});

package/tests/instrumentation/google/set_telemetry_attributes.test.js

@@ -0,0 +1,219 @@
+'use strict';
+
+// Tests for the Gemini generate_content enrichment complement.
+//
+// It stamps `gen_ai.system` (vertex_ai/gemini) and audio token details onto the
+// ACTIVE span (traceloop's chat span) when `generateContent{,Stream}Internal`
+// runs. The @google/genai `Models` class is stubbed; the test plays traceloop by
+// starting a span and making it active around the call.
+
+const test = require('node:test');
+const assert = require('node:assert/strict');
+const Module = require('node:module');
+
+const { trace, context } = require('@opentelemetry/api');
+const {
+  BasicTracerProvider,
+  InMemorySpanExporter,
+  SimpleSpanProcessor,
+} = require('@opentelemetry/sdk-trace-base');
+const { AsyncLocalStorageContextManager } = require('@opentelemetry/context-async-hooks');
+
+const {
+  GoogleTelemetryInstrumentor,
+} = require('../../../src/instrumentation/google/set_telemetry_attributes');
+
+const GENAI_MODULE = '@google/genai';
+
+// The exact chunk objects the streaming stub yields. Shared so the passthrough
+// test can assert object identity — proof our wrapper re-yields the user's chunks
+// without copying or mutating them. The last chunk carries the streamed usage.
+const STREAM_CHUNKS = [
+  { text: 'hello ' },
+  { text: 'world', usageMetadata: { cachedContentTokenCount: 32 } },
+];
+
+const exporter = new InMemorySpanExporter();
+trace.setGlobalTracerProvider(
+  new BasicTracerProvider({ spanProcessors: [new SimpleSpanProcessor(exporter)] })
+);
+context.setGlobalContextManager(new AsyncLocalStorageContextManager().enable());
+
+class StubModels {
+  generateContentInternal(params = {}) {
+    return Promise.resolve({
+      text: 'ok',
+      usageMetadata: {
+        promptTokensDetails: params.audioInput
+          ? [{ modality: 'AUDIO', tokenCount: 9 }]
+          : [{ modality: 'TEXT', tokenCount: 2 }],
+        candidatesTokensDetails: params.audioOutput
+          ? [{ modality: 'AUDIO', tokenCount: 17 }]
+          : [{ modality: 'TEXT', tokenCount: 5 }],
+        cachedContentTokenCount: params.cached || 0,
+        cacheTokensDetails: params.cacheDetails,
+      },
+    });
+  }
+
+  generateContentStreamInternal() {
+    return Promise.resolve(
+      (async function* () {
+        for (const chunk of STREAM_CHUNKS) yield chunk;
+      })()
+    );
+  }
+}
+
+function instrumented() {
+  const stubs = { [GENAI_MODULE]: { Models: StubModels } };
+  const origResolve = Module._resolveFilename;
+  Module._resolveFilename = function patched(request, ...rest) {
+    if (request in stubs) return request;
+    return origResolve.call(this, request, ...rest);
+  };
+  for (const [moduleId, exports] of Object.entries(stubs)) {
+    require.cache[moduleId] = { id: moduleId, filename: moduleId, loaded: true, exports };
+  }
+  const instrumentor = new GoogleTelemetryInstrumentor();
+  instrumentor.instrument();
+  return () => {
+    instrumentor.uninstrument();
+    Module._resolveFilename = origResolve;
+    for (const moduleId of Object.keys(stubs)) delete require.cache[moduleId];
+    exporter.reset();
+  };
+}
+
+async function withActiveSpan(fn) {
+  const span = trace.getTracer('test').startSpan('generate_content');
+  await context.with(trace.setSpan(context.active(), span), fn);
+  span.end();
+  return exporter.getFinishedSpans().at(-1);
+}
+
+test('stamps provider gemini and audio output tokens on the active span', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubModels();
+    client.apiClient = { isVertexAI: () => false };
+    const span = await withActiveSpan(() =>
+      client.generateContentInternal({ model: 'tts', audioOutput: true })
+    );
+    assert.equal(span.attributes['gen_ai.system'], 'gemini');
+    assert.equal(span.attributes['gen_ai.usage.output_tokens_details.audio_tokens'], 17);
+    assert.equal(span.attributes['gen_ai.usage.input_tokens_details.audio_tokens'], undefined);
+  } finally {
+    teardown();
+  }
+});
+
+test('stamps provider vertex_ai and audio input tokens', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubModels();
+    client.apiClient = { isVertexAI: () => true };
+    const span = await withActiveSpan(() =>
+      client.generateContentInternal({ model: 'm', audioInput: true })
+    );
+    assert.equal(span.attributes['gen_ai.system'], 'vertex_ai');
+    assert.equal(span.attributes['gen_ai.usage.input_tokens_details.audio_tokens'], 9);
+  } finally {
+    teardown();
+  }
+});
+
+test('no audio token attributes for a text-only call', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubModels();
+    client.apiClient = { isVertexAI: () => false };
+    const span = await withActiveSpan(() => client.generateContentInternal({ model: 'm' }));
+    assert.equal(span.attributes['gen_ai.usage.output_tokens_details.audio_tokens'], undefined);
+    assert.equal(span.attributes['gen_ai.usage.cache_read_input_tokens'], undefined);
+    assert.equal(span.attributes['gen_ai.system'], 'gemini');
+  } finally {
+    teardown();
+  }
+});
+
+test('stamps cached-content tokens on the active span', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubModels();
+    client.apiClient = { isVertexAI: () => false };
+    const span = await withActiveSpan(() =>
+      client.generateContentInternal({ model: 'm', cached: 128 })
+    );
+    assert.equal(span.attributes['gen_ai.usage.cache_read_input_tokens'], 128);
+    assert.equal(
+      span.attributes['gen_ai.usage.input_tokens_details.cached_audio_tokens'],
+      undefined
+    );
+  } finally {
+    teardown();
+  }
+});
+
+test('splits cached tokens into audio and text remainder by modality', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubModels();
+    client.apiClient = { isVertexAI: () => false };
+    // 200 cached tokens, 50 of them audio: audio is billed separately, the rest
+    // stays as text cache-read so the two never double-count.
+    const span = await withActiveSpan(() =>
+      client.generateContentInternal({
+        model: 'm',
+        cached: 200,
+        cacheDetails: [
+          { modality: 'TEXT', tokenCount: 150 },
+          { modality: 'AUDIO', tokenCount: 50 },
+        ],
+      })
+    );
+    assert.equal(span.attributes['gen_ai.usage.input_tokens_details.cached_audio_tokens'], 50);
+    assert.equal(span.attributes['gen_ai.usage.cache_read_input_tokens'], 150);
+  } finally {
+    teardown();
+  }
+});
+
+test('streaming sets provider only and uninstrument restores methods', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubModels();
+    client.apiClient = { isVertexAI: () => true };
+    const span = await withActiveSpan(() => client.generateContentStreamInternal({ model: 'm' }));
+    assert.equal(span.attributes['gen_ai.system'], 'vertex_ai');
+  } finally {
+    teardown();
+  }
+  // After teardown the prototype method is restored.
+  const client = new StubModels();
+  assert.equal(typeof client.generateContentInternal, 'function');
+});
+
+test('streaming re-yields the user chunks unchanged and stamps cache off the stream', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubModels();
+    client.apiClient = { isVertexAI: () => false };
+    const collected = [];
+    const span = await withActiveSpan(async () => {
+      const stream = await client.generateContentStreamInternal({ model: 'm' });
+      for await (const chunk of stream) collected.push(chunk);
+    });
+    // Same count, order, and object identity: observeCache re-yields each chunk
+    // by reference — no telemetry collision with the caller's payload.
+    assert.equal(collected.length, STREAM_CHUNKS.length);
+    collected.forEach((chunk, index) => {
+      assert.equal(chunk, STREAM_CHUNKS[index]);
+      assert.equal(chunk.text, STREAM_CHUNKS[index].text);
+    });
+    // The wrapper still ran: cache tokens from the final chunk landed on the span.
+    assert.equal(span.attributes['gen_ai.usage.cache_read_input_tokens'], 32);
+  } finally {
+    teardown();
+  }
+});

package/tests/instrumentation/openai/get_audio_telemetry.test.js

@@ -0,0 +1,170 @@
+'use strict';
+
+// Tests for the bundled OpenAI audio instrumentor.
+//
+// The real `openai` audio resource classes are stubbed with `create` methods
+// that mirror the SDK's shapes — Text-to-Speech resolves to the audio bytes and
+// Speech-to-Text to a transcription carrying `duration` — so the instrumentor
+// can be exercised without network access or the dependency installed.
+
+const test = require('node:test');
+const assert = require('node:assert/strict');
+const Module = require('node:module');
+
+const { trace, SpanStatusCode } = require('@opentelemetry/api');
+const {
+  BasicTracerProvider,
+  InMemorySpanExporter,
+  SimpleSpanProcessor,
+} = require('@opentelemetry/sdk-trace-base');
+
+const {
+  OpenAIAudioInstrumentor,
+} = require('../../../src/instrumentation/openai/get_audio_telemetry');
+
+const SPEECH_MODULE = 'openai/resources/audio/speech';
+const TRANSCRIPTIONS_MODULE = 'openai/resources/audio/transcriptions';
+
+const exporter = new InMemorySpanExporter();
+trace.setGlobalTracerProvider(
+  new BasicTracerProvider({ spanProcessors: [new SimpleSpanProcessor(exporter)] })
+);
+
+class StubSpeech {
+  create(body = {}) {
+    if (body.input === 'boom') return Promise.reject(new Error('api error'));
+    return Promise.resolve(Buffer.from('audio-bytes'));
+  }
+}
+
+class StubTranscriptions {
+  create(body = {}) {
+    if (body.model === 'boom') return Promise.reject(new Error('api error'));
+    return Promise.resolve({ duration: 12.4, text: 'hello' });
+  }
+}
+
+// Stub the openai audio resource modules, install the instrumentor + the
+// in-memory exporter, and return a teardown that restores everything.
+function instrumented() {
+  const stubs = {
+    [SPEECH_MODULE]: { Speech: StubSpeech },
+    [TRANSCRIPTIONS_MODULE]: { Transcriptions: StubTranscriptions },
+  };
+  const origResolve = Module._resolveFilename;
+  Module._resolveFilename = function patched(request, ...rest) {
+    if (request in stubs) return request;
+    return origResolve.call(this, request, ...rest);
+  };
+  for (const [moduleId, exports] of Object.entries(stubs)) {
+    require.cache[moduleId] = { id: moduleId, filename: moduleId, loaded: true, exports };
+  }
+
+  const instrumentor = new OpenAIAudioInstrumentor();
+  instrumentor.instrument();
+
+  return () => {
+    instrumentor.uninstrument();
+    Module._resolveFilename = origResolve;
+    for (const moduleId of Object.keys(stubs)) delete require.cache[moduleId];
+    exporter.reset();
+  };
+}
+
+test('speech emits genai span with char count', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubSpeech();
+    assert.deepEqual(
+      await client.create({ input: 'hello world', model: 'tts-1', voice: 'alloy' }),
+      Buffer.from('audio-bytes')
+    );
+
+    const [span] = exporter.getFinishedSpans();
+    assert.equal(span.name, 'text_to_speech');
+    assert.equal(span.attributes['gen_ai.system'], 'openai');
+    assert.equal(span.attributes['gen_ai.operation.name'], 'text_to_speech');
+    assert.equal(span.attributes['gen_ai.request.model'], 'tts-1');
+    assert.equal(span.attributes['gen_ai.usage.input_characters'], 'hello world'.length);
+    assert.equal(span.attributes['gen_ai.usage.input_seconds'], undefined);
+    assert.equal(span.status.code, SpanStatusCode.UNSET);
+  } finally {
+    teardown();
+  }
+});
+
+test('speech records error status on failure', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubSpeech();
+    await assert.rejects(() => client.create({ input: 'boom', model: 'tts-1' }), /api error/);
+
+    const [span] = exporter.getFinishedSpans();
+    assert.equal(span.status.code, SpanStatusCode.ERROR);
+  } finally {
+    teardown();
+  }
+});
+
+test('transcriptions emits span with audio duration', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubTranscriptions();
+    const result = await client.create({ file: Buffer.from('audio'), model: 'whisper-1' });
+    assert.equal(result.duration, 12.4);
+
+    const [span] = exporter.getFinishedSpans();
+    assert.equal(span.name, 'speech_to_text');
+    assert.equal(span.attributes['gen_ai.system'], 'openai');
+    assert.equal(span.attributes['gen_ai.operation.name'], 'speech_to_text');
+    assert.equal(span.attributes['gen_ai.request.model'], 'whisper-1');
+    assert.equal(span.attributes['gen_ai.usage.input_seconds'], 12.4);
+    assert.equal(span.attributes['gen_ai.usage.input_characters'], undefined);
+    assert.equal(span.status.code, SpanStatusCode.UNSET);
+  } finally {
+    teardown();
+  }
+});
+
+test('transcriptions records error status on failure', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubTranscriptions();
+    await assert.rejects(
+      () => client.create({ file: Buffer.from('audio'), model: 'boom' }),
+      /api error/
+    );
+
+    const [span] = exporter.getFinishedSpans();
+    assert.equal(span.status.code, SpanStatusCode.ERROR);
+  } finally {
+    teardown();
+  }
+});
+
+test('uninstrument restores original methods', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubSpeech();
+    await client.create({ input: 'x', model: 'tts-1' });
+    assert.equal(exporter.getFinishedSpans().length, 1);
+  } finally {
+    teardown();
+  }
+});
+
+test('system is azure when client targets azure', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubSpeech();
+    // Azure is reached via the same openai SDK; the vendor is read off the
+    // client base URL so the span is attributed to azure.ai.openai, not openai.
+    client._client = { baseURL: 'https://acme.openai.azure.com/' };
+    await client.create({ input: 'hi', model: 'tts-1' });
+
+    const [span] = exporter.getFinishedSpans();
+    assert.equal(span.attributes['gen_ai.system'], 'azure.ai.openai');
+  } finally {
+    teardown();
+  }
+});

package/tests/instrumentation/openai/get_cache_and_reasoning_telemetry.test.js

@@ -0,0 +1,83 @@
+'use strict';
+
+// Tests for the OpenAI non-streaming cached-prompt token complement. It wraps the
+// primary instrumentor's `_endSpan({ span, type, result })` to stamp
+// `gen_ai.usage.cache_read_input_tokens` from `usage.prompt_tokens_details.cached_tokens`.
+
+const test = require('node:test');
+const assert = require('node:assert/strict');
+
+const {
+  OpenAIPromptCachingInstrumentor,
+} = require('../../../src/instrumentation/openai/get_cache_and_reasoning_telemetry');
+
+function fakeSpan() {
+  const attrs = {};
+  return {
+    attrs,
+    isRecording: () => true,
+    setAttribute(key, value) {
+      attrs[key] = value;
+    },
+  };
+}
+
+test('stamps cache_read_input_tokens and still calls the original _endSpan', () => {
+  const calls = [];
+  const instrumentation = {
+    _endSpan(args) {
+      calls.push(args);
+    },
+  };
+  new OpenAIPromptCachingInstrumentor().instrument(instrumentation);
+
+  const span = fakeSpan();
+  instrumentation._endSpan({
+    span,
+    type: 'chat',
+    result: { usage: { prompt_tokens: 1600, prompt_tokens_details: { cached_tokens: 1536 } } },
+  });
+
+  assert.equal(span.attrs['gen_ai.usage.cache_read_input_tokens'], 1536);
+  assert.equal(calls.length, 1);
+});
+
+test('sets nothing when there are no cached tokens', () => {
+  const instrumentation = { _endSpan() {} };
+  new OpenAIPromptCachingInstrumentor().instrument(instrumentation);
+
+  const span = fakeSpan();
+  instrumentation._endSpan({ span, type: 'chat', result: { usage: { prompt_tokens: 10 } } });
+
+  assert.equal(span.attrs['gen_ai.usage.cache_read_input_tokens'], undefined);
+});
+
+test('stamps reasoning_tokens from completion_tokens_details', () => {
+  const instrumentation = { _endSpan() {} };
+  new OpenAIPromptCachingInstrumentor().instrument(instrumentation);
+
+  const span = fakeSpan();
+  instrumentation._endSpan({
+    span,
+    type: 'chat',
+    result: { usage: { output_tokens: 115, completion_tokens_details: { reasoning_tokens: 64 } } },
+  });
+
+  assert.equal(span.attrs['gen_ai.usage.reasoning_tokens'], 64);
+});
+
+test('is a no-op when the instrumentor has no _endSpan', () => {
+  const instrumentation = {};
+  new OpenAIPromptCachingInstrumentor().instrument(instrumentation);
+  assert.equal(instrumentation._endSpan, undefined);
+});
+
+test('uninstrument restores the original _endSpan', () => {
+  const original = () => {};
+  const instrumentation = { _endSpan: original };
+  const instrumentor = new OpenAIPromptCachingInstrumentor();
+  instrumentor.instrument(instrumentation);
+  assert.notEqual(instrumentation._endSpan, original);
+  instrumentor.uninstrument();
+  assert.equal(instrumentation._endSpan, original);
+});