diff --git a/packages/ai-jsx/package.json b/packages/ai-jsx/package.json index d3250ccd3..8bdb55abc 100644 --- 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 index b1228c2ce..d1b57ce43 100644 --- 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( generator: AsyncGenerator @@ -36,6 +38,8 @@ function spanAttributesForElement(element: Element): 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 index 49ea4d517..785cd4dc8 100644 --- 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 index 386b86f43..b580911ff 100644 --- 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 index 9fb91ae1e..f86260959 100644 --- 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( + + hello + + ); + + 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": "[ + {"hello"} + ]", + "ai.jsx.tag": "UserMessage", + "ai.jsx.tree": " + {"hello"} + ", + }, + { + "ai.jsx.result": "[ + {"hello"} + ]", + "ai.jsx.tag": "UserMessage", + "ai.jsx.tree": " + {"hello"} + ", + }, + { + "ai.jsx.result": "[ + {"hello"} + ]", + "ai.jsx.tag": "ShrinkConversation", + "ai.jsx.tree": " + + {"hello"} + + ", + }, + { + "ai.jsx.result": "hello", + "ai.jsx.result.tokenCount": 1, + "ai.jsx.tag": "UserMessage", + "ai.jsx.tree": " + {"hello"} + ", + }, + { + "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": " + {"▮"} + ", + }, + { + "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": " + + {"hello"} + + ", + }, + { + "ai.jsx.result": "opentel response from OpenAI", + "ai.jsx.result.tokenCount": 7, + "ai.jsx.tag": "AutomaticChatModel", + "ai.jsx.tree": " + + {"hello"} + + ", + }, + { + "ai.jsx.result": "opentel response from OpenAI", + "ai.jsx.result.tokenCount": 7, + "ai.jsx.tag": "ChatCompletion", + "ai.jsx.tree": " + + {"hello"} + + ", + }, + ] + `); + }); +}); + it('passes creates a chat completion', async () => { mockOpenAIResponse('response from OpenAI'); diff --git a/yarn.lock b/yarn.lock index 3586e9a9c..a787a49a5 100644 --- 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