weflayr 0.20.6 → 0.22.0

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

@@ -0,0 +1,145 @@
+'use strict';
+
+// Tests for the bundled OpenAI embeddings instrumentor.
+//
+// The real `openai` Embeddings resource class is stubbed with a `create` method
+// that mirrors the SDK's shape — resolving to a response carrying `model` and a
+// `usage.prompt_tokens` count — 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 {
+  OpenAIEmbeddingsInstrumentor,
+} = require('../../../src/instrumentation/openai/get_embeddings_telemetry');
+
+const EMBEDDINGS_MODULE = 'openai/resources/embeddings';
+
+const exporter = new InMemorySpanExporter();
+trace.setGlobalTracerProvider(
+  new BasicTracerProvider({ spanProcessors: [new SimpleSpanProcessor(exporter)] })
+);
+
+class StubEmbeddings {
+  create(body = {}) {
+    if (body.model === 'boom') return Promise.reject(new Error('api error'));
+    return Promise.resolve({
+      object: 'list',
+      data: [{ embedding: [0.1, 0.2] }],
+      model: 'text-embedding-3-small',
+      usage: { prompt_tokens: 8, total_tokens: 8 },
+    });
+  }
+}
+
+// Stub the openai embeddings resource module, install the instrumentor + the
+// in-memory exporter, and return a teardown that restores everything.
+function instrumented() {
+  const stubs = { [EMBEDDINGS_MODULE]: { Embeddings: StubEmbeddings } };
+  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 OpenAIEmbeddingsInstrumentor();
+  instrumentor.instrument();
+
+  return () => {
+    instrumentor.uninstrument();
+    Module._resolveFilename = origResolve;
+    for (const moduleId of Object.keys(stubs)) delete require.cache[moduleId];
+    exporter.reset();
+  };
+}
+
+test('embeddings emits genai span with prompt tokens', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubEmbeddings();
+    const result = await client.create({ input: 'hello', model: 'text-embedding-3-small' });
+    assert.equal(result.usage.prompt_tokens, 8);
+
+    const [span] = exporter.getFinishedSpans();
+    assert.equal(span.name, 'embeddings');
+    assert.equal(span.attributes['gen_ai.system'], 'openai');
+    assert.equal(span.attributes['gen_ai.operation.name'], 'embeddings');
+    assert.equal(span.attributes['gen_ai.request.model'], 'text-embedding-3-small');
+    assert.equal(span.attributes['gen_ai.response.model'], 'text-embedding-3-small');
+    assert.equal(span.attributes['gen_ai.usage.input_tokens'], 8);
+    assert.deepEqual(JSON.parse(span.attributes['gen_ai.input.messages']), [
+      { role: 'user', parts: [{ type: 'text', content: 'hello' }] },
+    ]);
+    assert.equal(span.status.code, SpanStatusCode.UNSET);
+  } finally {
+    teardown();
+  }
+});
+
+test('embeddings emits one input message per element of a list input', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubEmbeddings();
+    await client.create({ input: ['a', 'b'], model: 'text-embedding-3-small' });
+
+    const [span] = exporter.getFinishedSpans();
+    assert.deepEqual(JSON.parse(span.attributes['gen_ai.input.messages']), [
+      { role: 'user', parts: [{ type: 'text', content: 'a' }] },
+      { role: 'user', parts: [{ type: 'text', content: 'b' }] },
+    ]);
+  } finally {
+    teardown();
+  }
+});
+
+test('embeddings records error status on failure', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubEmbeddings();
+    await assert.rejects(() => client.create({ input: 'x', 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 StubEmbeddings();
+    await client.create({ input: 'x', model: 'text-embedding-3-small' });
+    assert.equal(exporter.getFinishedSpans().length, 1);
+  } finally {
+    teardown();
+  }
+});
+
+test('system is azure when client targets azure', async () => {
+  const teardown = instrumented();
+  try {
+    const client = new StubEmbeddings();
+    // 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: 'text-embedding-3-small' });
+
+    const [span] = exporter.getFinishedSpans();
+    assert.equal(span.attributes['gen_ai.system'], 'azure.ai.openai');
+  } finally {
+    teardown();
+  }
+});

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

@@ -0,0 +1,105 @@
+'use strict';
+
+// Tests for the OpenAI streamed token-usage complement. It wraps the primary
+// instrumentor's `_streamingWrapPromise` so the streamed usage chunk's counts land
+// on the span as gen_ai.usage.* attributes.
+
+const test = require('node:test');
+const assert = require('node:assert/strict');
+
+const {
+  OpenAIStreamingUsageInstrumentor,
+} = require('../../../src/instrumentation/openai/get_streaming_usage');
+
+function fakeSpan() {
+  const attrs = {};
+  return {
+    attrs,
+    isRecording: () => true,
+    setAttribute(key, value) {
+      attrs[key] = value;
+    },
+  };
+}
+
+test('backfills token usage from the streamed usage chunk', async () => {
+  const instrumentation = {
+    _streamingWrapPromise() {
+      return (async function* () {
+        yield { choices: [{ delta: { content: 'hi' } }] };
+        yield { usage: { prompt_tokens: 12, completion_tokens: 3, total_tokens: 15 } };
+      })();
+    },
+  };
+  new OpenAIStreamingUsageInstrumentor().instrument(instrumentation);
+
+  const span = fakeSpan();
+  const stream = instrumentation._streamingWrapPromise({ span });
+  for await (const _chunk of stream) {
+    void _chunk;
+  }
+
+  assert.equal(span.attrs['gen_ai.usage.input_tokens'], 12);
+  assert.equal(span.attrs['gen_ai.usage.output_tokens'], 3);
+  assert.equal(span.attrs['gen_ai.usage.total_tokens'], 15);
+});
+
+test('backfills cached-prompt tokens from the streamed usage chunk', async () => {
+  const instrumentation = {
+    _streamingWrapPromise() {
+      return (async function* () {
+        yield {
+          usage: {
+            prompt_tokens: 1600,
+            completion_tokens: 2,
+            total_tokens: 1602,
+            prompt_tokens_details: { cached_tokens: 1536 },
+          },
+        };
+      })();
+    },
+  };
+  new OpenAIStreamingUsageInstrumentor().instrument(instrumentation);
+
+  const span = fakeSpan();
+  for await (const _chunk of instrumentation._streamingWrapPromise({ span })) {
+    void _chunk;
+  }
+
+  assert.equal(span.attrs['gen_ai.usage.cache_read_input_tokens'], 1536);
+});
+
+test('re-yields every stream chunk unchanged — no telemetry collision with the payload', async () => {
+  const emitted = [
+    { choices: [{ delta: { content: 'hel' } }] },
+    { choices: [{ delta: { content: 'lo' } }] },
+    { usage: { prompt_tokens: 12, completion_tokens: 3, total_tokens: 15 } },
+  ];
+  const instrumentation = {
+    _streamingWrapPromise() {
+      return (async function* () {
+        for (const chunk of emitted) yield chunk;
+      })();
+    },
+  };
+  new OpenAIStreamingUsageInstrumentor().instrument(instrumentation);
+
+  const span = fakeSpan();
+  const collected = [];
+  for await (const chunk of instrumentation._streamingWrapPromise({ span })) {
+    collected.push(chunk);
+  }
+
+  // Same count, order, and object identity: backfillUsage re-yields each chunk by
+  // reference without copying or mutating it.
+  assert.equal(collected.length, emitted.length);
+  collected.forEach((chunk, index) => assert.equal(chunk, emitted[index]));
+  // …while still backfilling usage off the trailing chunk.
+  assert.equal(span.attrs['gen_ai.usage.output_tokens'], 3);
+});
+
+test('is a no-op when the instrumentor has no _streamingWrapPromise', () => {
+  const instrumentation = {};
+  new OpenAIStreamingUsageInstrumentor().instrument(instrumentation);
+  assert.equal(instrumentation._streamingWrapPromise, undefined);
+});

package/src/otel/openai-streaming-usage.js

@@ -1,60 +0,0 @@
-'use strict';
-
-// Backfill token usage onto streamed OpenAI spans.
-//
-// @traceloop/instrumentation-openai (the OpenAI OTel instrumentor we load for
-// the `openai` SDK) does not record token usage for streamed calls: its stream
-// handler accumulates `choices[].delta` but never reads the final usage chunk
-// that OpenAI sends when the request sets `stream_options: { include_usage:
-// true }`. Non-streaming calls work because OpenAI returns `usage` in the
-// response body, which the instrumentor reads directly. The Python OpenLLMetry
-// instrumentor does read the streamed usage, so this gap is JS-only.
-//
-// Read `usage` from the streamed chunks. The instrumentor's streaming handler hands us the span
-// and ends it only after its chunk loop, so when the usage chunk arrives the
-// span is still open and we set the same `gen_ai.usage.*` attributes the
-// instrumentor sets for non-streaming calls.
-
-const USAGE_INPUT_TOKENS_ATTR = 'gen_ai.usage.input_tokens';
-const USAGE_OUTPUT_TOKENS_ATTR = 'gen_ai.usage.output_tokens';
-const USAGE_TOTAL_TOKENS_ATTR = 'gen_ai.usage.total_tokens';
-
-function _setUsageAttributes(span, usage) {
-  if (typeof usage.prompt_tokens === 'number') {
-    span.setAttribute(USAGE_INPUT_TOKENS_ATTR, usage.prompt_tokens);
-  }
-  if (typeof usage.completion_tokens === 'number') {
-    span.setAttribute(USAGE_OUTPUT_TOKENS_ATTR, usage.completion_tokens);
-  }
-  if (typeof usage.total_tokens === 'number') {
-    span.setAttribute(USAGE_TOTAL_TOKENS_ATTR, usage.total_tokens);
-  }
-}
-
-/**
- * Wrap the OpenAI instrumentor's streaming handler so streamed spans carry
- * token usage. No-op if the instrumentor's internals have changed shape — the
- * worst case is streamed spans without token counts, never a crash.
- */
-function patchOpenAIStreamingUsage(instrumentation) {
-  const originalStreamingWrap = instrumentation._streamingWrapPromise;
-  if (typeof originalStreamingWrap !== 'function') return;
-
-  instrumentation._streamingWrapPromise = function streamingWrapWithUsage(streamArgs) {
-    const span = streamArgs && streamArgs.span;
-    const wrappedStream = originalStreamingWrap.call(this, streamArgs);
-    if (!span) return wrappedStream;
-
-    return (async function* backfillUsage() {
-      for await (const chunk of wrappedStream) {
-        const usage = chunk && chunk.usage;
-        if (usage && span.isRecording()) {
-          _setUsageAttributes(span, usage);
-        }
-        yield chunk;
-      }
-    })();
-  };
-}
-
-module.exports = { patchOpenAIStreamingUsage };