Skip to content

Commit ed7af2a

Browse files
lenatracelooplenatraceloop
committed
feat(instrumentation-llamaindex): migrate to OTel 1.40 GenAI semantic conventions.
1 parent d255ea4 commit ed7af2a

File tree

10 files changed

+972
-123
lines changed

10 files changed

+972
-123
lines changed

‎packages/instrumentation-llamaindex/package.json‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@
4242
"@opentelemetry/instrumentation": "^0.203.0",
4343
"@opentelemetry/semantic-conventions": "^1.38.0",
4444
"@traceloop/ai-semantic-conventions": "workspace:*",
45+
"@traceloop/instrumentation-utils": "workspace:*",
4546
"lodash": "^4.17.21",
4647
"tslib": "^2.8.1"
4748
},

‎packages/instrumentation-llamaindex/src/custom-llm-instrumentation.ts‎

Lines changed: 138 additions & 64 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
import * as lodash from "lodash";
21
import type * as llamaindex from "llamaindex";
32

43
import {
@@ -13,15 +12,29 @@ import {
1312
} from "@opentelemetry/api";
1413
import { safeExecuteInTheMiddle } from "@opentelemetry/instrumentation";
1514

16-
import { SpanAttributes } from "@traceloop/ai-semantic-conventions";
1715
import {
18-
ATTR_GEN_AI_COMPLETION,
19-
ATTR_GEN_AI_PROMPT,
16+
SpanAttributes,
17+
FinishReasons,
18+
} from "@traceloop/ai-semantic-conventions";
19+
import {
20+
ATTR_GEN_AI_INPUT_MESSAGES,
21+
ATTR_GEN_AI_OPERATION_NAME,
22+
ATTR_GEN_AI_OUTPUT_MESSAGES,
23+
ATTR_GEN_AI_PROVIDER_NAME,
2024
ATTR_GEN_AI_REQUEST_MODEL,
2125
ATTR_GEN_AI_REQUEST_TOP_P,
26+
ATTR_GEN_AI_RESPONSE_FINISH_REASONS,
2227
ATTR_GEN_AI_RESPONSE_MODEL,
23-
ATTR_GEN_AI_SYSTEM,
28+
ATTR_GEN_AI_USAGE_INPUT_TOKENS,
29+
ATTR_GEN_AI_USAGE_OUTPUT_TOKENS,
30+
GEN_AI_OPERATION_NAME_VALUE_CHAT,
31+
GEN_AI_PROVIDER_NAME_VALUE_OPENAI,
2432
} from "@opentelemetry/semantic-conventions/incubating";
33+
import {
34+
formatInputMessages,
35+
formatOutputMessage,
36+
mapOpenAIContentBlock,
37+
} from "@traceloop/instrumentation-utils";
2538

2639
import { LlamaIndexInstrumentationConfig } from "./types";
2740
import { shouldSendPrompts, llmGeneratorWrapper } from "./utils";
@@ -33,9 +46,21 @@ type AsyncResponseType =
3346
| AsyncIterable<llamaindex.ChatResponseChunk>
3447
| AsyncIterable<llamaindex.CompletionResponse>;
3548

49+
const classNameToProviderName: Record<string, string> = {
50+
OpenAI: GEN_AI_PROVIDER_NAME_VALUE_OPENAI,
51+
};
52+
53+
export const openAIFinishReasonMap: Record<string, string> = {
54+
stop: FinishReasons.STOP,
55+
length: FinishReasons.LENGTH,
56+
tool_calls: FinishReasons.TOOL_CALL,
57+
content_filter: FinishReasons.CONTENT_FILTER,
58+
function_call: FinishReasons.TOOL_CALL,
59+
};
60+
3661
export class CustomLLMInstrumentation {
3762
constructor(
38-
private config: LlamaIndexInstrumentationConfig,
63+
private config: () => LlamaIndexInstrumentationConfig,
3964
private diag: DiagLogger,
4065
private tracer: () => Tracer,
4166
) {}
@@ -50,44 +75,30 @@ export class CustomLLMInstrumentation {
5075
const messages = params?.messages;
5176
const streaming = params?.stream;
5277

53-
const span = plugin
54-
.tracer()
55-
.startSpan(`llamaindex.${lodash.snakeCase(className)}.chat`, {
56-
kind: SpanKind.CLIENT,
57-
});
78+
const span = plugin.tracer().startSpan(`chat ${this.metadata.model}`, {
79+
kind: SpanKind.CLIENT,
80+
});
5881

5982
try {
60-
span.setAttribute(ATTR_GEN_AI_SYSTEM, className);
83+
span.setAttribute(
84+
ATTR_GEN_AI_PROVIDER_NAME,
85+
classNameToProviderName[className] ?? className.toLowerCase(),
86+
);
6187
span.setAttribute(ATTR_GEN_AI_REQUEST_MODEL, this.metadata.model);
62-
span.setAttribute(SpanAttributes.LLM_REQUEST_TYPE, "chat");
88+
span.setAttribute(
89+
ATTR_GEN_AI_OPERATION_NAME,
90+
GEN_AI_OPERATION_NAME_VALUE_CHAT,
91+
);
6392
span.setAttribute(ATTR_GEN_AI_REQUEST_TOP_P, this.metadata.topP);
64-
if (shouldSendPrompts(plugin.config)) {
65-
for (const messageIdx in messages) {
66-
const content = messages[messageIdx].content;
67-
if (typeof content === "string") {
68-
span.setAttribute(
69-
`${ATTR_GEN_AI_PROMPT}.${messageIdx}.content`,
70-
content as string,
71-
);
72-
} else if (
73-
(content as llamaindex.MessageContentDetail[])[0].type ===
74-
"text"
75-
) {
76-
span.setAttribute(
77-
`${ATTR_GEN_AI_PROMPT}.${messageIdx}.content`,
78-
(content as llamaindex.MessageContentTextDetail[])[0].text,
79-
);
80-
}
81-
82-
span.setAttribute(
83-
`${ATTR_GEN_AI_PROMPT}.${messageIdx}.role`,
84-
messages[messageIdx].role,
85-
);
86-
}
93+
if (shouldSendPrompts(plugin.config()) && messages) {
94+
span.setAttribute(
95+
ATTR_GEN_AI_INPUT_MESSAGES,
96+
formatInputMessages(messages, mapOpenAIContentBlock),
97+
);
8798
}
8899
} catch (e) {
89100
plugin.diag.warn(e);
90-
plugin.config.exceptionLogger?.(e);
101+
plugin.config().exceptionLogger?.(e);
91102
}
92103

93104
const execContext = trace.setSpan(context.active(), span);
@@ -138,36 +149,59 @@ export class CustomLLMInstrumentation {
138149
): T {
139150
span.setAttribute(ATTR_GEN_AI_RESPONSE_MODEL, metadata.model);
140151

141-
if (!shouldSendPrompts(this.config)) {
142-
span.setStatus({ code: SpanStatusCode.OK });
143-
span.end();
144-
return result;
145-
}
146-
147152
try {
148-
if ((result as llamaindex.ChatResponse).message) {
153+
const raw = (result as any).raw;
154+
const finishReason: string | null =
155+
raw?.choices?.[0]?.finish_reason ?? null;
156+
157+
// finish_reasons: metadata, not content — always set outside shouldSendPrompts
158+
if (finishReason != null) {
159+
span.setAttribute(ATTR_GEN_AI_RESPONSE_FINISH_REASONS, [
160+
openAIFinishReasonMap[finishReason] ?? finishReason,
161+
]);
162+
}
163+
164+
// Token usage: always set when available
165+
const usage = raw?.usage;
166+
if (usage) {
167+
span.setAttribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, usage.prompt_tokens);
168+
span.setAttribute(
169+
ATTR_GEN_AI_USAGE_OUTPUT_TOKENS,
170+
usage.completion_tokens,
171+
);
149172
span.setAttribute(
150-
`${ATTR_GEN_AI_COMPLETION}.0.role`,
151-
(result as llamaindex.ChatResponse).message.role,
173+
SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS,
174+
usage.total_tokens,
152175
);
176+
}
177+
178+
// output messages: content — always set inside shouldSendPrompts
179+
if (
180+
shouldSendPrompts(this.config()) &&
181+
(result as llamaindex.ChatResponse).message
182+
) {
153183
const content = (result as llamaindex.ChatResponse).message.content;
154-
if (typeof content === "string") {
155-
span.setAttribute(`${ATTR_GEN_AI_COMPLETION}.0.content`, content);
156-
} else if (content[0].type === "text") {
157-
span.setAttribute(
158-
`${ATTR_GEN_AI_COMPLETION}.0.content`,
159-
content[0].text,
160-
);
161-
}
162-
span.setStatus({ code: SpanStatusCode.OK });
184+
// Normalize to array so mapOpenAIContentBlock handles both string and block array
185+
const contentArray = typeof content === "string" ? [content] : content;
186+
span.setAttribute(
187+
ATTR_GEN_AI_OUTPUT_MESSAGES,
188+
formatOutputMessage(
189+
contentArray,
190+
finishReason,
191+
openAIFinishReasonMap,
192+
GEN_AI_OPERATION_NAME_VALUE_CHAT,
193+
mapOpenAIContentBlock,
194+
),
195+
);
163196
}
197+
198+
span.setStatus({ code: SpanStatusCode.OK });
164199
} catch (e) {
165200
this.diag.warn(e);
166-
this.config.exceptionLogger?.(e);
201+
this.config().exceptionLogger?.(e);
167202
}
168203

169204
span.end();
170-
171205
return result;
172206
}
173207

@@ -178,14 +212,54 @@ export class CustomLLMInstrumentation {
178212
metadata: llamaindex.LLMMetadata,
179213
): T {
180214
span.setAttribute(ATTR_GEN_AI_RESPONSE_MODEL, metadata.model);
181-
if (!shouldSendPrompts(this.config)) {
182-
span.setStatus({ code: SpanStatusCode.OK });
183-
span.end();
184-
return result;
185-
}
186215

187-
return llmGeneratorWrapper(result, execContext, (message) => {
188-
span.setAttribute(`${ATTR_GEN_AI_COMPLETION}.0.content`, message);
216+
return llmGeneratorWrapper(result, execContext, (message, lastChunk) => {
217+
try {
218+
// Extract finish_reason and usage from the last chunk's raw OpenAI
219+
// response — available when stream_options: { include_usage: true }
220+
// is set on the LLM (OpenAI sends usage in the final streaming chunk).
221+
const lastRaw = lastChunk?.raw as any;
222+
const finishReason: string | null =
223+
lastRaw?.choices?.[0]?.finish_reason ?? null;
224+
const usage = lastRaw?.usage ?? null;
225+
226+
if (finishReason != null) {
227+
span.setAttribute(ATTR_GEN_AI_RESPONSE_FINISH_REASONS, [
228+
openAIFinishReasonMap[finishReason] ?? finishReason,
229+
]);
230+
}
231+
232+
if (usage) {
233+
span.setAttribute(
234+
ATTR_GEN_AI_USAGE_INPUT_TOKENS,
235+
usage.prompt_tokens,
236+
);
237+
span.setAttribute(
238+
ATTR_GEN_AI_USAGE_OUTPUT_TOKENS,
239+
usage.completion_tokens,
240+
);
241+
span.setAttribute(
242+
SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS,
243+
usage.total_tokens,
244+
);
245+
}
246+
247+
if (shouldSendPrompts(this.config())) {
248+
span.setAttribute(
249+
ATTR_GEN_AI_OUTPUT_MESSAGES,
250+
formatOutputMessage(
251+
[message],
252+
finishReason,
253+
openAIFinishReasonMap,
254+
GEN_AI_OPERATION_NAME_VALUE_CHAT,
255+
mapOpenAIContentBlock,
256+
),
257+
);
258+
}
259+
} catch (e) {
260+
this.diag.warn(e);
261+
this.config().exceptionLogger?.(e);
262+
}
189263
span.setStatus({ code: SpanStatusCode.OK });
190264
span.end();
191265
}) as any;

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase {
4141
constructor(config: LlamaIndexInstrumentationConfig = {}) {
4242
super("@traceloop/instrumentation-llamaindex", version, config);
4343
this.customLLMInstrumentation = new CustomLLMInstrumentation(
44-
this._config,
44+
() => this._config,
4545
this._diag,
4646
() => this.tracer,
4747
);

‎packages/instrumentation-llamaindex/src/utils.ts‎

Lines changed: 31 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -58,23 +58,43 @@ export async function* llmGeneratorWrapper(
5858
| AsyncIterable<llamaindex.ChatResponseChunk>
5959
| AsyncIterable<llamaindex.CompletionResponse>,
6060
ctx: Context,
61-
fn: (message: string) => void,
61+
fn: (message: string, lastChunk?: any) => void,
6262
) {
6363
let message = "";
64+
// Track the last chunk so the callback can extract usage/finish_reason from
65+
// chunk.raw — OpenAI sends these in the final streaming chunk when
66+
// stream_options: { include_usage: true } is set on the LLM.
67+
let lastChunk: any;
6468

65-
for await (const messageChunk of bindAsyncGenerator(
66-
ctx,
67-
streamingResult as AsyncGenerator,
68-
)) {
69-
if ((messageChunk as llamaindex.ChatResponseChunk).delta) {
70-
message += (messageChunk as llamaindex.ChatResponseChunk).delta;
69+
let fnCalled = false;
70+
try {
71+
for await (const messageChunk of bindAsyncGenerator(
72+
ctx,
73+
streamingResult as AsyncGenerator,
74+
)) {
75+
if ((messageChunk as llamaindex.ChatResponseChunk).delta) {
76+
message += (messageChunk as llamaindex.ChatResponseChunk).delta;
77+
}
78+
if ((messageChunk as llamaindex.CompletionResponse).text) {
79+
message += (messageChunk as llamaindex.CompletionResponse).text;
80+
}
81+
lastChunk = messageChunk;
82+
yield messageChunk;
83+
}
84+
} catch (err) {
85+
// Ensure span is finalized even if the stream throws
86+
if (!fnCalled) {
87+
fnCalled = true;
88+
fn(message, lastChunk);
7189
}
72-
if ((messageChunk as llamaindex.CompletionResponse).text) {
73-
message += (messageChunk as llamaindex.CompletionResponse).text;
90+
throw err;
91+
} finally {
92+
// Covers normal completion and early consumer exit (break/return)
93+
if (!fnCalled) {
94+
fnCalled = true;
95+
fn(message, lastChunk);
7496
}
75-
yield messageChunk;
7697
}
77-
fn(message);
7898
}
7999

80100
export function genericWrapper(

‎packages/instrumentation-llamaindex/test/finish_reasons.test.ts‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
/**
2+
* Unit tests for openAIFinishReasonMap.
3+
*
4+
* Each OpenAI raw finish reason value is tested individually.
5+
* Verified values from OpenAI API documentation.
6+
*/
7+
8+
import * as assert from "assert";
9+
import { FinishReasons } from "@traceloop/ai-semantic-conventions";
10+
import { openAIFinishReasonMap } from "../src/custom-llm-instrumentation";
11+
12+
const VALID_OTEL_FINISH_REASONS = new Set([
13+
FinishReasons.STOP,
14+
FinishReasons.LENGTH,
15+
FinishReasons.TOOL_CALL,
16+
FinishReasons.CONTENT_FILTER,
17+
FinishReasons.ERROR,
18+
]);
19+
20+
describe("openAIFinishReasonMap", () => {
21+
it("all mapped values are valid OTel finish reason strings", () => {
22+
for (const [raw, otel] of Object.entries(openAIFinishReasonMap)) {
23+
assert.ok(
24+
VALID_OTEL_FINISH_REASONS.has(otel),
25+
`openAIFinishReasonMap["${raw}"] = "${otel}" is not a valid OTel finish reason`,
26+
);
27+
}
28+
});
29+
30+
it('maps "stop" to stop', () => {
31+
assert.strictEqual(openAIFinishReasonMap["stop"], FinishReasons.STOP);
32+
});
33+
34+
it('maps "length" to length', () => {
35+
assert.strictEqual(openAIFinishReasonMap["length"], FinishReasons.LENGTH);
36+
});
37+
38+
it('maps "tool_calls" to tool_call', () => {
39+
assert.strictEqual(
40+
openAIFinishReasonMap["tool_calls"],
41+
FinishReasons.TOOL_CALL,
42+
);
43+
});
44+
45+
it('maps "content_filter" to content_filter', () => {
46+
assert.strictEqual(
47+
openAIFinishReasonMap["content_filter"],
48+
FinishReasons.CONTENT_FILTER,
49+
);
50+
});
51+
52+
it('maps "function_call" to tool_call (deprecated alias)', () => {
53+
assert.strictEqual(
54+
openAIFinishReasonMap["function_call"],
55+
FinishReasons.TOOL_CALL,
56+
);
57+
});
58+
});

0 commit comments

Comments
 (0)