Skip to content

Commit 2ef0c13

Browse files
nirganirga
authored
fix(openai): enrich token metrics on streaming requests (#183)
1 parent 83737ee commit 2ef0c13

File tree

8 files changed

+138
-4
lines changed

8 files changed

+138
-4
lines changed

‎package-lock.json

Lines changed: 7 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎packages/instrumentation-openai/package.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,7 +39,8 @@
3939
"@opentelemetry/core": "^1.22.0",
4040
"@opentelemetry/instrumentation": "^0.49.0",
4141
"@opentelemetry/semantic-conventions": "^1.22.0",
42-
"@traceloop/ai-semantic-conventions": "^0.5.27"
42+
"@traceloop/ai-semantic-conventions": "^0.5.27",
43+
"tiktoken": "^1.0.13"
4344
},
4445
"devDependencies": {
4546
"@pollyjs/adapter-node-http": "^6.0.6",

‎packages/instrumentation-openai/src/instrumentation.ts

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@ import type {
4545
} from "openai/resources";
4646
import type { Stream } from "openai/streaming";
4747
import { version } from "../package.json";
48+
import { encoding_for_model, TiktokenModel, Tiktoken } from "tiktoken";
4849

4950
export class OpenAIInstrumentation extends InstrumentationBase<any> {
5051
protected declare _config: OpenAIInstrumentationConfig;
@@ -198,6 +199,7 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
198199
plugin._streamingWrapPromise({
199200
span,
200201
type,
202+
params: args[0] as any,
201203
promise: execPromise,
202204
}),
203205
);
@@ -296,15 +298,18 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
296298
private async *_streamingWrapPromise({
297299
span,
298300
type,
301+
params,
299302
promise,
300303
}:
301304
| {
302305
span: Span;
303306
type: "chat";
307+
params: ChatCompletionCreateParamsStreaming;
304308
promise: Promise<Stream<ChatCompletionChunk>>;
305309
}
306310
| {
307311
span: Span;
312+
params: CompletionCreateParamsStreaming;
308313
type: "completion";
309314
promise: Promise<Stream<Completion>>;
310315
}) {
@@ -356,6 +361,29 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
356361
this._addLogProbsEvent(span, result.choices[0].logprobs);
357362
}
358363

364+
if (this._config.enrichTokens) {
365+
let promptTokens = 0;
366+
for (const message of params.messages) {
367+
promptTokens +=
368+
this.tokenCountFromString(
369+
message.content as string,
370+
result.model,
371+
) ?? 0;
372+
}
373+
374+
const completionTokens = this.tokenCountFromString(
375+
result.choices[0].message.content ?? "",
376+
result.model,
377+
);
378+
if (completionTokens) {
379+
result.usage = {
380+
prompt_tokens: promptTokens,
381+
completion_tokens: completionTokens,
382+
total_tokens: promptTokens + completionTokens,
383+
};
384+
}
385+
}
386+
359387
this._endSpan({ span, type, result });
360388
} else {
361389
const result: Completion = {
@@ -394,6 +422,23 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
394422
this._addLogProbsEvent(span, result.choices[0].logprobs);
395423
}
396424

425+
if (this._config.enrichTokens) {
426+
const promptTokens =
427+
this.tokenCountFromString(params.prompt as string, result.model) ?? 0;
428+
429+
const completionTokens = this.tokenCountFromString(
430+
result.choices[0].text ?? "",
431+
result.model,
432+
);
433+
if (completionTokens) {
434+
result.usage = {
435+
prompt_tokens: promptTokens,
436+
completion_tokens: completionTokens,
437+
total_tokens: promptTokens + completionTokens,
438+
};
439+
}
440+
}
441+
397442
this._endSpan({ span, type, result });
398443
}
399444
}
@@ -588,4 +633,23 @@ export class OpenAIInstrumentation extends InstrumentationBase<any> {
588633

589634
span.addEvent("logprobs", { logprobs: JSON.stringify(result) });
590635
}
636+
637+
private _encodingCache = new Map<string, Tiktoken>();
638+
639+
private tokenCountFromString(text: string, model: string) {
640+
if (!this._encodingCache.has(model)) {
641+
try {
642+
const encoding = encoding_for_model(model as TiktokenModel);
643+
this._encodingCache.set(model, encoding);
644+
} catch (e) {
645+
this._diag.warn(
646+
`Failed to get tiktoken encoding for model_name: ${model}, error: ${e}`,
647+
);
648+
return;
649+
}
650+
}
651+
652+
const encoding = this._encodingCache.get(model);
653+
return encoding!.encode(text).length;
654+
}
591655
}

‎packages/instrumentation-openai/src/types.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,10 @@ export interface OpenAIInstrumentationConfig extends InstrumentationConfig {
66
* @default true
77
*/
88
traceContent?: boolean;
9+
10+
/**
11+
* Whether to enrich token information if missing from the trace.
12+
* @default false
13+
*/
14+
enrichTokens?: boolean;
915
}

‎packages/instrumentation-openai/test/instrumentation.test.ts

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ describe("Test OpenAI instrumentation", async function () {
5858
process.env.OPENAI_API_KEY = "test";
5959
}
6060
provider.addSpanProcessor(new SimpleSpanProcessor(memoryExporter));
61-
instrumentation = new OpenAIInstrumentation();
61+
instrumentation = new OpenAIInstrumentation({ enrichTokens: true });
6262
instrumentation.setTracerProvider(provider);
6363

6464
const openAIModule: typeof OpenAIModule = await import("openai");
@@ -103,6 +103,18 @@ describe("Test OpenAI instrumentation", async function () {
103103
completionSpan.attributes[`${SpanAttributes.LLM_PROMPTS}.0.content`],
104104
"Tell me a joke about OpenTelemetry",
105105
);
106+
assert.ok(
107+
completionSpan.attributes[`${SpanAttributes.LLM_USAGE_TOTAL_TOKENS}`],
108+
);
109+
assert.equal(
110+
completionSpan.attributes[`${SpanAttributes.LLM_USAGE_PROMPT_TOKENS}`],
111+
"15",
112+
);
113+
assert.ok(
114+
+completionSpan.attributes[
115+
`${SpanAttributes.LLM_USAGE_COMPLETION_TOKENS}`
116+
]! > 0,
117+
);
106118
});
107119

108120
it("should set attributes in span for streaming chat", async () => {
@@ -136,6 +148,18 @@ describe("Test OpenAI instrumentation", async function () {
136148
completionSpan.attributes[`${SpanAttributes.LLM_COMPLETIONS}.0.content`],
137149
result,
138150
);
151+
assert.ok(
152+
completionSpan.attributes[`${SpanAttributes.LLM_USAGE_TOTAL_TOKENS}`],
153+
);
154+
assert.equal(
155+
completionSpan.attributes[`${SpanAttributes.LLM_USAGE_PROMPT_TOKENS}`],
156+
"8",
157+
);
158+
assert.ok(
159+
+completionSpan.attributes[
160+
`${SpanAttributes.LLM_USAGE_COMPLETION_TOKENS}`
161+
]! > 0,
162+
);
139163
});
140164

141165
it("should set attributes in span for streaming chat with new API", async () => {
@@ -169,6 +193,26 @@ describe("Test OpenAI instrumentation", async function () {
169193
completionSpan.attributes[`${SpanAttributes.LLM_COMPLETIONS}.0.content`],
170194
result,
171195
);
196+
assert.ok(
197+
completionSpan.attributes[`${SpanAttributes.LLM_USAGE_PROMPT_TOKENS}`],
198+
);
199+
assert.ok(
200+
completionSpan.attributes[
201+
`${SpanAttributes.LLM_USAGE_COMPLETION_TOKENS}`
202+
],
203+
);
204+
assert.ok(
205+
completionSpan.attributes[`${SpanAttributes.LLM_USAGE_TOTAL_TOKENS}`],
206+
);
207+
assert.equal(
208+
completionSpan.attributes[`${SpanAttributes.LLM_USAGE_PROMPT_TOKENS}`],
209+
"8",
210+
);
211+
assert.ok(
212+
+completionSpan.attributes[
213+
`${SpanAttributes.LLM_USAGE_COMPLETION_TOKENS}`
214+
]! > 0,
215+
);
172216
});
173217

174218
it("should set attributes in span for completion", async () => {

‎packages/traceloop-sdk/src/lib/configuration/index.ts

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,10 @@ export const initialize = (options: InitializeOptions) => {
5454
options.traceloopSyncDevPollingInterval =
5555
Number(process.env.TRACELOOP_SYNC_DEV_POLLING_INTERVAL) || 5;
5656
}
57+
58+
if (options.shouldEnrichMetrics === undefined) {
59+
options.shouldEnrichMetrics = true;
60+
}
5761
}
5862

5963
validateConfiguration(options);

‎packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,12 @@ export interface InitializeOptions {
4646
*/
4747
logLevel?: "debug" | "info" | "warn" | "error";
4848

49+
/**
50+
* Whether to enrich metrics with additional data like OpenAI token usage for streaming requests. Optional.
51+
* Defaults to true.
52+
*/
53+
shouldEnrichMetrics?: boolean;
54+
4955
/**
5056
* Whether to log prompts, completions and embeddings on traces. Optional.
5157
* Defaults to true.

‎packages/traceloop-sdk/src/lib/tracing/index.ts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,9 @@ export const manuallyInitInstrumentations = (
7676
instrumentModules: InitializeOptions["instrumentModules"],
7777
) => {
7878
if (instrumentModules?.openAI) {
79-
openAIInstrumentation = new OpenAIInstrumentation();
79+
openAIInstrumentation = new OpenAIInstrumentation({
80+
enrichTokens: _configuration?.shouldEnrichMetrics,
81+
});
8082
instrumentations.push(openAIInstrumentation);
8183
openAIInstrumentation.manuallyInstrument(instrumentModules.openAI);
8284
}
@@ -149,6 +151,7 @@ export const startTracing = (options: InitializeOptions) => {
149151
if (!shouldSendTraces()) {
150152
openAIInstrumentation?.setConfig({
151153
traceContent: false,
154+
enrichTokens: _configuration?.shouldEnrichMetrics,
152155
});
153156
azureOpenAIInstrumentation?.setConfig({
154157
traceContent: false,

0 commit comments

Comments
 (0)