Add token counts to OpenTel traces (#238)

This makes it easier to know how different types of content (e.g. system messages, docsQA, etc) are contributing to the final token count. This will not exhaustively cover all usecases. This also adds a unit test.
fixie-ai · Aug 15, 2023 · 0d2e6d8 · 0d2e6d8 · vercel · Aug 15, 2023
1 parent 04c3de5
commit 0d2e6d8
Show file tree

Hide file tree

Showing 6 changed files with 218 additions and 3 deletions.
diff --git a/packages/ai-jsx/package.json b/packages/ai-jsx/package.json
@@ -4,7 +4,7 @@
   "repository": "fixie-ai/ai-jsx",
   "bugs": "https://github.com/fixie-ai/ai-jsx/issues",
   "homepage": "https://ai-jsx.com",
-  "version": "0.9.0",
+  "version": "0.9.1",
   "volta": {
     "extends": "../../package.json"
   },

diff --git a/packages/ai-jsx/src/core/opentelemetry.ts b/packages/ai-jsx/src/core/opentelemetry.ts
@@ -3,6 +3,8 @@ import { PartiallyRendered, StreamRenderer } from './render.js';
 import { Element, isElement } from './node.js';
 import { memoizedIdSymbol } from './memoize.js';
 import { debug } from './debug.js';
+import { getEncoding } from 'js-tiktoken';
+import _ from 'lodash';
 
 function bindAsyncGeneratorToActiveContext<T = unknown, TReturn = any, TNext = unknown>(
   generator: AsyncGenerator<T, TReturn, TNext>
@@ -36,6 +38,8 @@ function spanAttributesForElement(element: Element<any>): opentelemetry.Attribut
   return { 'ai.jsx.tag': element.tag.name, 'ai.jsx.tree': debug(element, true) };
 }
 
+const getEncoder = _.once(() => getEncoding('cl100k_base'));
+
 export function openTelemetryStreamRenderer(streamRenderer: StreamRenderer): StreamRenderer {
   const tracer = opentelemetry.trace.getTracer('ai.jsx');
 
@@ -59,8 +63,12 @@ export function openTelemetryStreamRenderer(streamRenderer: StreamRenderer): Str
           throw ex;
         } finally {
           if (result) {
+            const resultIsPartial = result.find(isElement);
             // Record the rendered value.
-            span.setAttribute('ai.jsx.result', result.find(isElement) ? debug(result, true) : result.join(''));
+            span.setAttribute('ai.jsx.result', resultIsPartial ? debug(result, true) : result.join(''));
+            if (!resultIsPartial) {
+              span.setAttribute('ai.jsx.result.tokenCount', getEncoder().encode(result.join('')).length);
+            }
           }
           cleanup(span);
           span.end();

diff --git a/packages/docs/docs/changelog.md b/packages/docs/docs/changelog.md
@@ -1,6 +1,10 @@
 # Changelog
 
-## 0.9.0
+## 0.9.1
+
+- Add `tokenCount` field to [OpenTelemetry-emitted spans](./guides/observability.md#opentelemetry-integration). Now, if you're emitting via OpenTelemetry (e.g. to DataDog), the spans will tell you how many tokens each component resolved to. This is helpful for answering quetsions like "how big is my system message?".
+
+## [0.9.0](https://github.com/fixie-ai/ai-jsx/commit/94624bedc27defc96f7cfead96094c8a577c8e27)
 
 - **Breaking:** Remove prompt-engineered `UseTools`. Previously, if you called `UseTools` with a model that doesn't support native function calling (e.g. Anthropic), `UseTools` would use a polyfilled version that uses prompt engineering to simulate function calling. However, this wasn't reliable enough in practice, so we've dropped it.
 - Fix issue where `gpt-4-32k` didn't accept functions.

diff --git a/packages/examples/package.json b/packages/examples/package.json
@@ -78,6 +78,8 @@
     "@opentelemetry/instrumentation-fetch": "^0.41.1",
     "@opentelemetry/sdk-logs": "^0.41.1",
     "@opentelemetry/sdk-node": "^0.41.1",
+    "@opentelemetry/sdk-trace-base": "^1.15.2",
+    "@opentelemetry/sdk-trace-node": "^1.15.2",
     "@wandb/sdk": "^0.5.1",
     "ai-jsx": "workspace:*",
     "axios": "^1.4.0",

diff --git a/packages/examples/test/core/completion.tsx b/packages/examples/test/core/completion.tsx
@@ -42,6 +42,133 @@ import { Anthropic } from 'ai-jsx/lib/anthropic';
 import { CompletionCreateParams } from '@anthropic-ai/sdk/resources/completions';
 import { Jsonifiable } from 'type-fest';
 
+import { trace } from '@opentelemetry/api';
+import { SimpleSpanProcessor, InMemorySpanExporter } from '@opentelemetry/sdk-trace-base';
+import { NodeTracerProvider } from '@opentelemetry/sdk-trace-node';
+import _ from 'lodash';
+
+afterEach(() => {
+  fetchMock.resetMocks();
+});
+
+describe('OpenTelemetry', () => {
+  const memoryExporter = new InMemorySpanExporter();
+  const tracerProvider = new NodeTracerProvider();
+  tracerProvider.addSpanProcessor(new SimpleSpanProcessor(memoryExporter));
+  trace.setGlobalTracerProvider(tracerProvider);
+
+  beforeEach(() => {
+    memoryExporter.reset();
+  });
+
+  it('should emit a span with the correct attributes', async () => {
+    mockOpenAIResponse('opentel response from OpenAI');
+
+    // Call your function that creates spans
+    await AI.createRenderContext({ enableOpenTelemetry: true }).render(
+      <ChatCompletion>
+        <UserMessage>hello</UserMessage>
+      </ChatCompletion>
+    );
+
+    const spans = memoryExporter.getFinishedSpans();
+    const minimalSpans = _.map(spans, 'attributes');
+    // Unfortunately, the @memoizedId will be sensitive how many tests in this file ran before it.
+    // To avoid issues with that, we put this test first.
+    expect(minimalSpans).toMatchInlineSnapshot(`
+      [
+        {
+          "ai.jsx.result": "[<UserMessage @memoizedId=1>
+        {"hello"}
+      </UserMessage>]",
+          "ai.jsx.tag": "UserMessage",
+          "ai.jsx.tree": "<UserMessage @memoizedId=1>
+        {"hello"}
+      </UserMessage>",
+        },
+        {
+          "ai.jsx.result": "[<UserMessage @memoizedId=1>
+        {"hello"}
+      </UserMessage>]",
+          "ai.jsx.tag": "UserMessage",
+          "ai.jsx.tree": "<UserMessage @memoizedId=1>
+        {"hello"}
+      </UserMessage>",
+        },
+        {
+          "ai.jsx.result": "[<UserMessage @memoizedId=1>
+        {"hello"}
+      </UserMessage>]",
+          "ai.jsx.tag": "ShrinkConversation",
+          "ai.jsx.tree": "<ShrinkConversation cost={tokenCountForConversationMessage} budget={4093}>
+        <UserMessage>
+          {"hello"}
+        </UserMessage>
+      </ShrinkConversation>",
+        },
+        {
+          "ai.jsx.result": "hello",
+          "ai.jsx.result.tokenCount": 1,
+          "ai.jsx.tag": "UserMessage",
+          "ai.jsx.tree": "<UserMessage @memoizedId=1>
+        {"hello"}
+      </UserMessage>",
+        },
+        {
+          "ai.jsx.result": "opentel response from OpenAI",
+          "ai.jsx.result.tokenCount": 7,
+          "ai.jsx.tag": "Stream",
+          "ai.jsx.tree": ""▮"",
+        },
+        {
+          "ai.jsx.result": "opentel response from OpenAI",
+          "ai.jsx.result.tokenCount": 7,
+          "ai.jsx.tag": "AssistantMessage",
+          "ai.jsx.tree": "<AssistantMessage>
+        {"▮"}
+      </AssistantMessage>",
+        },
+        {
+          "ai.jsx.result": "opentel response from OpenAI",
+          "ai.jsx.result.tokenCount": 7,
+          "ai.jsx.tag": "Stream",
+          "ai.jsx.tree": ""opentel response from OpenAI"",
+        },
+        {
+          "ai.jsx.result": "opentel response from OpenAI",
+          "ai.jsx.result.tokenCount": 7,
+          "ai.jsx.tag": "OpenAIChatModel",
+          "ai.jsx.tree": "<OpenAIChatModel model="gpt-3.5-turbo">
+        <UserMessage>
+          {"hello"}
+        </UserMessage>
+      </OpenAIChatModel>",
+        },
+        {
+          "ai.jsx.result": "opentel response from OpenAI",
+          "ai.jsx.result.tokenCount": 7,
+          "ai.jsx.tag": "AutomaticChatModel",
+          "ai.jsx.tree": "<AutomaticChatModel>
+        <UserMessage>
+          {"hello"}
+        </UserMessage>
+      </AutomaticChatModel>",
+        },
+        {
+          "ai.jsx.result": "opentel response from OpenAI",
+          "ai.jsx.result.tokenCount": 7,
+          "ai.jsx.tag": "ChatCompletion",
+          "ai.jsx.tree": "<ChatCompletion>
+        <UserMessage>
+          {"hello"}
+        </UserMessage>
+      </ChatCompletion>",
+        },
+      ]
+    `);
+  });
+});
+
 it('passes creates a chat completion', async () => {
   mockOpenAIResponse('response from OpenAI');
 

diff --git a/yarn.lock b/yarn.lock
@@ -4307,6 +4307,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@opentelemetry/context-async-hooks@npm:1.15.2":
+  version: 1.15.2
+  resolution: "@opentelemetry/context-async-hooks@npm:1.15.2"
+  peerDependencies:
+    "@opentelemetry/api": ">=1.0.0 <1.5.0"
+  checksum: f5e00a9920e14f12a4abd4f855927d7f2fcae3e57048b9853860529608d680e43473216643089c95de95519cf3642f635b178849acafec034cd7174198297295
+  languageName: node
+  linkType: hard
+
 "@opentelemetry/core@npm:1.15.1":
   version: 1.15.1
   resolution: "@opentelemetry/core@npm:1.15.1"
@@ -4556,6 +4565,17 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@opentelemetry/propagator-b3@npm:1.15.2":
+  version: 1.15.2
+  resolution: "@opentelemetry/propagator-b3@npm:1.15.2"
+  dependencies:
+    "@opentelemetry/core": 1.15.2
+  peerDependencies:
+    "@opentelemetry/api": ">=1.0.0 <1.5.0"
+  checksum: 830b34f678b1b6c12f51135ada6555f2981d1d9c44f59fa21c4a550b0edb483b644791368ffc0b09b4ec7aed77e9ff6e437df91513c86c4dc35a25841f6248c1
+  languageName: node
+  linkType: hard
+
 "@opentelemetry/propagator-jaeger@npm:1.15.1":
   version: 1.15.1
   resolution: "@opentelemetry/propagator-jaeger@npm:1.15.1"
@@ -4567,6 +4587,17 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@opentelemetry/propagator-jaeger@npm:1.15.2":
+  version: 1.15.2
+  resolution: "@opentelemetry/propagator-jaeger@npm:1.15.2"
+  dependencies:
+    "@opentelemetry/core": 1.15.2
+  peerDependencies:
+    "@opentelemetry/api": ">=1.0.0 <1.5.0"
+  checksum: 319e4c5e197d338d016ad3d5eddc7307c34d2217e81c182d17868c3b5a44442d59de5c0eb4b88c9690d253c42cf8151201d69169641349d250503d9f36d9fe7b
+  languageName: node
+  linkType: hard
+
 "@opentelemetry/resources@npm:1.15.1":
   version: 1.15.1
   resolution: "@opentelemetry/resources@npm:1.15.1"
@@ -4579,6 +4610,18 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@opentelemetry/resources@npm:1.15.2":
+  version: 1.15.2
+  resolution: "@opentelemetry/resources@npm:1.15.2"
+  dependencies:
+    "@opentelemetry/core": 1.15.2
+    "@opentelemetry/semantic-conventions": 1.15.2
+  peerDependencies:
+    "@opentelemetry/api": ">=1.0.0 <1.5.0"
+  checksum: 072d64bee2a073ac3f1218ba9d24bd0a20fc69988d206688c2f53e813c9d84ae325093c3c0e8909c2208cfa583fe5bad71d16af7d6b087a348d866c1d0857396
+  languageName: node
+  linkType: hard
+
 "@opentelemetry/sdk-logs@npm:0.41.1, @opentelemetry/sdk-logs@npm:^0.41.1":
   version: 0.41.1
   resolution: "@opentelemetry/sdk-logs@npm:0.41.1"
@@ -4642,6 +4685,19 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@opentelemetry/sdk-trace-base@npm:1.15.2, @opentelemetry/sdk-trace-base@npm:^1.15.2":
+  version: 1.15.2
+  resolution: "@opentelemetry/sdk-trace-base@npm:1.15.2"
+  dependencies:
+    "@opentelemetry/core": 1.15.2
+    "@opentelemetry/resources": 1.15.2
+    "@opentelemetry/semantic-conventions": 1.15.2
+  peerDependencies:
+    "@opentelemetry/api": ">=1.0.0 <1.5.0"
+  checksum: 3ca9d71919451f8e4a0644434981c7fa6dc29d002da03784578fbe47689625d106c00ab5b35539a8d348e40676ea57d44ba6845a9a604dd6be670911f83835bf
+  languageName: node
+  linkType: hard
+
 "@opentelemetry/sdk-trace-node@npm:1.15.1":
   version: 1.15.1
   resolution: "@opentelemetry/sdk-trace-node@npm:1.15.1"
@@ -4658,6 +4714,22 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@opentelemetry/sdk-trace-node@npm:^1.15.2":
+  version: 1.15.2
+  resolution: "@opentelemetry/sdk-trace-node@npm:1.15.2"
+  dependencies:
+    "@opentelemetry/context-async-hooks": 1.15.2
+    "@opentelemetry/core": 1.15.2
+    "@opentelemetry/propagator-b3": 1.15.2
+    "@opentelemetry/propagator-jaeger": 1.15.2
+    "@opentelemetry/sdk-trace-base": 1.15.2
+    semver: ^7.5.1
+  peerDependencies:
+    "@opentelemetry/api": ">=1.0.0 <1.5.0"
+  checksum: 67ffa3c9c40a4571ee1aadeec3623bb9bbf981593b5d68b157250b9a4a6a544605b4c68107b5cffb8bb81775868b9d5fee23ccec643f79d773d800ac8d4c155b
+  languageName: node
+  linkType: hard
+
 "@opentelemetry/sdk-trace-web@npm:1.15.1":
   version: 1.15.1
   resolution: "@opentelemetry/sdk-trace-web@npm:1.15.1"
@@ -11500,6 +11572,8 @@ __metadata:
     "@opentelemetry/instrumentation-fetch": ^0.41.1
     "@opentelemetry/sdk-logs": ^0.41.1
     "@opentelemetry/sdk-node": ^0.41.1
+    "@opentelemetry/sdk-trace-base": ^1.15.2
+    "@opentelemetry/sdk-trace-node": ^1.15.2
     "@tsconfig/node18": ^2.0.1
     "@types/jest": ^29.5.2
     "@types/lodash": ^4.14.195