diff --git a/packages/instrumentation-llamaindex/package.json b/packages/instrumentation-llamaindex/package.json index c9ded014f..dd1d27b9b 100644 --- a/packages/instrumentation-llamaindex/package.json +++ b/packages/instrumentation-llamaindex/package.json @@ -40,8 +40,9 @@ "@opentelemetry/api": "^1.9.0", "@opentelemetry/core": "^2.0.1", "@opentelemetry/instrumentation": "^0.203.0", - "@opentelemetry/semantic-conventions": "^1.38.0", + "@opentelemetry/semantic-conventions": "^1.40.0", "@traceloop/ai-semantic-conventions": "workspace:*", + "@traceloop/instrumentation-utils": "workspace:*", "lodash": "^4.17.21", "tslib": "^2.8.1" }, diff --git a/packages/instrumentation-llamaindex/src/custom-llm-instrumentation.ts b/packages/instrumentation-llamaindex/src/custom-llm-instrumentation.ts index fb9308e29..c3897e471 100644 --- a/packages/instrumentation-llamaindex/src/custom-llm-instrumentation.ts +++ b/packages/instrumentation-llamaindex/src/custom-llm-instrumentation.ts @@ -1,4 +1,3 @@ -import * as lodash from "lodash"; import type * as llamaindex from "llamaindex"; import { @@ -13,15 +12,31 @@ import { } from "@opentelemetry/api"; import { safeExecuteInTheMiddle } from "@opentelemetry/instrumentation"; -import { SpanAttributes } from "@traceloop/ai-semantic-conventions"; import { - ATTR_GEN_AI_COMPLETION, - ATTR_GEN_AI_PROMPT, + SpanAttributes, + FinishReasons, +} from "@traceloop/ai-semantic-conventions"; +import { + ATTR_GEN_AI_INPUT_MESSAGES, + ATTR_GEN_AI_OPERATION_NAME, + ATTR_GEN_AI_OUTPUT_MESSAGES, + ATTR_GEN_AI_PROVIDER_NAME, ATTR_GEN_AI_REQUEST_MODEL, + ATTR_GEN_AI_REQUEST_TEMPERATURE, ATTR_GEN_AI_REQUEST_TOP_P, + ATTR_GEN_AI_RESPONSE_FINISH_REASONS, + ATTR_GEN_AI_RESPONSE_ID, ATTR_GEN_AI_RESPONSE_MODEL, - ATTR_GEN_AI_SYSTEM, + ATTR_GEN_AI_USAGE_INPUT_TOKENS, + ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + GEN_AI_PROVIDER_NAME_VALUE_OPENAI, } from "@opentelemetry/semantic-conventions/incubating"; +import { + formatInputMessages, + formatOutputMessage, + mapOpenAIContentBlock, +} from "@traceloop/instrumentation-utils"; import { LlamaIndexInstrumentationConfig } from "./types"; import { shouldSendPrompts, llmGeneratorWrapper } from "./utils"; @@ -33,9 +48,23 @@ type AsyncResponseType = | AsyncIterable | AsyncIterable; +const classNameToProviderName: Record = { + OpenAI: GEN_AI_PROVIDER_NAME_VALUE_OPENAI, + // Future providers: Anthropic: "anthropic", Gemini: "gcp.gemini", etc. + // See well-known values: https://opentelemetry.io/docs/specs/semconv/registry/attributes/gen-ai/#gen-ai-provider-name +}; + +export const openAIFinishReasonMap: Record = { + stop: FinishReasons.STOP, + length: FinishReasons.LENGTH, + tool_calls: FinishReasons.TOOL_CALL, + content_filter: FinishReasons.CONTENT_FILTER, + function_call: FinishReasons.TOOL_CALL, +}; + export class CustomLLMInstrumentation { constructor( - private config: LlamaIndexInstrumentationConfig, + private config: () => LlamaIndexInstrumentationConfig, private diag: DiagLogger, private tracer: () => Tracer, ) {} @@ -50,44 +79,34 @@ export class CustomLLMInstrumentation { const messages = params?.messages; const streaming = params?.stream; - const span = plugin - .tracer() - .startSpan(`llamaindex.${lodash.snakeCase(className)}.chat`, { - kind: SpanKind.CLIENT, - }); + const span = plugin.tracer().startSpan(`chat ${this.metadata.model}`, { + kind: SpanKind.CLIENT, + }); try { - span.setAttribute(ATTR_GEN_AI_SYSTEM, className); + span.setAttribute( + ATTR_GEN_AI_PROVIDER_NAME, + classNameToProviderName[className] ?? className.toLowerCase(), + ); span.setAttribute(ATTR_GEN_AI_REQUEST_MODEL, this.metadata.model); - span.setAttribute(SpanAttributes.LLM_REQUEST_TYPE, "chat"); + span.setAttribute( + ATTR_GEN_AI_OPERATION_NAME, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + ); span.setAttribute(ATTR_GEN_AI_REQUEST_TOP_P, this.metadata.topP); - if (shouldSendPrompts(plugin.config)) { - for (const messageIdx in messages) { - const content = messages[messageIdx].content; - if (typeof content === "string") { - span.setAttribute( - `${ATTR_GEN_AI_PROMPT}.${messageIdx}.content`, - content as string, - ); - } else if ( - (content as llamaindex.MessageContentDetail[])[0].type === - "text" - ) { - span.setAttribute( - `${ATTR_GEN_AI_PROMPT}.${messageIdx}.content`, - (content as llamaindex.MessageContentTextDetail[])[0].text, - ); - } - - span.setAttribute( - `${ATTR_GEN_AI_PROMPT}.${messageIdx}.role`, - messages[messageIdx].role, - ); - } + span.setAttribute( + ATTR_GEN_AI_REQUEST_TEMPERATURE, + this.metadata.temperature, + ); + if (shouldSendPrompts(plugin.config()) && messages) { + span.setAttribute( + ATTR_GEN_AI_INPUT_MESSAGES, + formatInputMessages(messages, mapOpenAIContentBlock), + ); } } catch (e) { plugin.diag.warn(e); - plugin.config.exceptionLogger?.(e); + plugin.config().exceptionLogger?.(e); } const execContext = trace.setSpan(context.active(), span); @@ -138,36 +157,62 @@ export class CustomLLMInstrumentation { ): T { span.setAttribute(ATTR_GEN_AI_RESPONSE_MODEL, metadata.model); - if (!shouldSendPrompts(this.config)) { - span.setStatus({ code: SpanStatusCode.OK }); - span.end(); - return result; - } - try { - if ((result as llamaindex.ChatResponse).message) { + const raw = (result as any).raw; + if (raw?.id) { + span.setAttribute(ATTR_GEN_AI_RESPONSE_ID, raw.id); + } + const finishReason: string | null = + raw?.choices?.[0]?.finish_reason ?? null; + + // finish_reasons: metadata, not content — always set outside shouldSendPrompts + if (finishReason != null) { + span.setAttribute(ATTR_GEN_AI_RESPONSE_FINISH_REASONS, [ + openAIFinishReasonMap[finishReason] ?? finishReason, + ]); + } + + // Token usage: always set when available + const usage = raw?.usage; + if (usage) { + span.setAttribute(ATTR_GEN_AI_USAGE_INPUT_TOKENS, usage.prompt_tokens); + span.setAttribute( + ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, + usage.completion_tokens, + ); span.setAttribute( - `${ATTR_GEN_AI_COMPLETION}.0.role`, - (result as llamaindex.ChatResponse).message.role, + SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS, + usage.total_tokens, ); + } + + // output messages: content — always set inside shouldSendPrompts + if ( + shouldSendPrompts(this.config()) && + (result as llamaindex.ChatResponse).message + ) { const content = (result as llamaindex.ChatResponse).message.content; - if (typeof content === "string") { - span.setAttribute(`${ATTR_GEN_AI_COMPLETION}.0.content`, content); - } else if (content[0].type === "text") { - span.setAttribute( - `${ATTR_GEN_AI_COMPLETION}.0.content`, - content[0].text, - ); - } - span.setStatus({ code: SpanStatusCode.OK }); + // Normalize to array so mapOpenAIContentBlock handles both string and block array + const contentArray = typeof content === "string" ? [content] : content; + span.setAttribute( + ATTR_GEN_AI_OUTPUT_MESSAGES, + formatOutputMessage( + contentArray, + finishReason, + openAIFinishReasonMap, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + mapOpenAIContentBlock, + ), + ); } + + span.setStatus({ code: SpanStatusCode.OK }); } catch (e) { this.diag.warn(e); - this.config.exceptionLogger?.(e); + this.config().exceptionLogger?.(e); } span.end(); - return result; } @@ -178,14 +223,67 @@ export class CustomLLMInstrumentation { metadata: llamaindex.LLMMetadata, ): T { span.setAttribute(ATTR_GEN_AI_RESPONSE_MODEL, metadata.model); - if (!shouldSendPrompts(this.config)) { - span.setStatus({ code: SpanStatusCode.OK }); - span.end(); - return result; - } - return llmGeneratorWrapper(result, execContext, (message) => { - span.setAttribute(`${ATTR_GEN_AI_COMPLETION}.0.content`, message); + return llmGeneratorWrapper(result, execContext, (message, lastChunk) => { + try { + // Extract finish_reason and usage from the last chunk's raw OpenAI + // response — available when stream_options: { include_usage: true } + // is set on the LLM (OpenAI sends usage in the final streaming chunk). + const lastRaw = lastChunk?.raw as any; + if (lastRaw?.id) { + span.setAttribute(ATTR_GEN_AI_RESPONSE_ID, lastRaw.id); + } + const finishReason: string | null = + lastRaw?.choices?.[0]?.finish_reason ?? null; + const usage = lastRaw?.usage ?? null; + + if (finishReason != null) { + span.setAttribute(ATTR_GEN_AI_RESPONSE_FINISH_REASONS, [ + openAIFinishReasonMap[finishReason] ?? finishReason, + ]); + } + + if (usage) { + span.setAttribute( + ATTR_GEN_AI_USAGE_INPUT_TOKENS, + usage.prompt_tokens, + ); + span.setAttribute( + ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, + usage.completion_tokens, + ); + span.setAttribute( + SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS, + usage.total_tokens, + ); + } + + if (!finishReason && !usage) { + this.diag.debug( + "LlamaIndex streaming: no finish_reason or usage in last chunk. " + + "Set stream_options: { include_usage: true } on the LLM to capture token usage.", + ); + } + + // Note: streaming only produces text parts — LlamaIndex's streaming interface + // yields text deltas only, not full content blocks. Tool calls or multi-modal + // content are collapsed into a single text string by llmGeneratorWrapper. + if (shouldSendPrompts(this.config())) { + span.setAttribute( + ATTR_GEN_AI_OUTPUT_MESSAGES, + formatOutputMessage( + [message], + finishReason, + openAIFinishReasonMap, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + mapOpenAIContentBlock, + ), + ); + } + } catch (e) { + this.diag.warn(e); + this.config().exceptionLogger?.(e); + } span.setStatus({ code: SpanStatusCode.OK }); span.end(); }) as any; diff --git a/packages/instrumentation-llamaindex/src/instrumentation.ts b/packages/instrumentation-llamaindex/src/instrumentation.ts index 4ed49f0d8..65cb0b930 100644 --- a/packages/instrumentation-llamaindex/src/instrumentation.ts +++ b/packages/instrumentation-llamaindex/src/instrumentation.ts @@ -36,19 +36,29 @@ import { version } from "../package.json"; export class LlamaIndexInstrumentation extends InstrumentationBase { declare protected _config: LlamaIndexInstrumentationConfig; + private customLLMInstrumentation!: CustomLLMInstrumentation; constructor(config: LlamaIndexInstrumentationConfig = {}) { super("@traceloop/instrumentation-llamaindex", version, config); + this.customLLMInstrumentation = new CustomLLMInstrumentation( + () => this._config, + this._diag, + () => this.tracer, + ); } public override setConfig(config: LlamaIndexInstrumentationConfig = {}) { super.setConfig(config); } - public manuallyInstrument(module: typeof llamaindex) { + public manuallyInstrument(module: typeof llamaindex, openaiModule?: any) { this._diag.debug("Manually instrumenting llamaindex"); this.patch(module); + + if (openaiModule) { + this.patchOpenAI(openaiModule); + } } protected init(): InstrumentationModuleDefinition[] { @@ -94,12 +104,6 @@ export class LlamaIndexInstrumentation extends InstrumentationBase { private patch(moduleExports: typeof llamaindex, moduleVersion?: string) { this._diag.debug(`Patching llamaindex@${moduleVersion}`); - const customLLMInstrumentation = new CustomLLMInstrumentation( - this._config, - this._diag, - () => this.tracer, // this is on purpose. Tracer may change - ); - this._wrap( moduleExports.RetrieverQueryEngine.prototype, "query", @@ -133,7 +137,7 @@ export class LlamaIndexInstrumentation extends InstrumentationBase { this._wrap( cls.prototype, "chat", - customLLMInstrumentation.chatWrapper({ className: cls.name }), + this.customLLMInstrumentation.chatWrapper({ className: cls.name }), ); } else if (this.isEmbedding(cls.prototype)) { this._wrap( @@ -202,7 +206,16 @@ export class LlamaIndexInstrumentation extends InstrumentationBase { private patchOpenAI(moduleExports: any, moduleVersion?: string) { this._diag.debug(`Patching @llamaindex/openai@${moduleVersion}`); - // Instrument OpenAIAgent if it exists + if (moduleExports.OpenAI && this.isLLM(moduleExports.OpenAI.prototype)) { + this._wrap( + moduleExports.OpenAI.prototype, + "chat", + this.customLLMInstrumentation.chatWrapper({ + className: moduleExports.OpenAI.name, + }), + ); + } + if (moduleExports.OpenAIAgent && moduleExports.OpenAIAgent.prototype) { this._wrap( moduleExports.OpenAIAgent.prototype, @@ -223,7 +236,10 @@ export class LlamaIndexInstrumentation extends InstrumentationBase { private unpatchOpenAI(moduleExports: any, moduleVersion?: string) { this._diag.debug(`Unpatching @llamaindex/openai@${moduleVersion}`); - // Unwrap OpenAIAgent if it exists + if (moduleExports.OpenAI && moduleExports.OpenAI.prototype) { + this._unwrap(moduleExports.OpenAI.prototype, "chat"); + } + if (moduleExports.OpenAIAgent && moduleExports.OpenAIAgent.prototype) { this._unwrap(moduleExports.OpenAIAgent.prototype, "chat"); } diff --git a/packages/instrumentation-llamaindex/src/utils.ts b/packages/instrumentation-llamaindex/src/utils.ts index 5bdb8c485..22d8c7c67 100644 --- a/packages/instrumentation-llamaindex/src/utils.ts +++ b/packages/instrumentation-llamaindex/src/utils.ts @@ -58,23 +58,43 @@ export async function* llmGeneratorWrapper( | AsyncIterable | AsyncIterable, ctx: Context, - fn: (message: string) => void, + fn: (message: string, lastChunk?: any) => void, ) { let message = ""; + // Track the last chunk so the callback can extract usage/finish_reason from + // chunk.raw — OpenAI sends these in the final streaming chunk when + // stream_options: { include_usage: true } is set on the LLM. + let lastChunk: any; - for await (const messageChunk of bindAsyncGenerator( - ctx, - streamingResult as AsyncGenerator, - )) { - if ((messageChunk as llamaindex.ChatResponseChunk).delta) { - message += (messageChunk as llamaindex.ChatResponseChunk).delta; + let fnCalled = false; + try { + for await (const messageChunk of bindAsyncGenerator( + ctx, + streamingResult as AsyncGenerator, + )) { + if ((messageChunk as llamaindex.ChatResponseChunk).delta) { + message += (messageChunk as llamaindex.ChatResponseChunk).delta; + } + if ((messageChunk as llamaindex.CompletionResponse).text) { + message += (messageChunk as llamaindex.CompletionResponse).text; + } + lastChunk = messageChunk; + yield messageChunk; + } + } catch (err) { + // Ensure span is finalized even if the stream throws + if (!fnCalled) { + fnCalled = true; + fn(message, lastChunk); } - if ((messageChunk as llamaindex.CompletionResponse).text) { - message += (messageChunk as llamaindex.CompletionResponse).text; + throw err; + } finally { + // Covers normal completion and early consumer exit (break/return) + if (!fnCalled) { + fnCalled = true; + fn(message, lastChunk); } - yield messageChunk; } - fn(message); } export function genericWrapper( diff --git a/packages/instrumentation-llamaindex/test/finish_reasons.test.ts b/packages/instrumentation-llamaindex/test/finish_reasons.test.ts new file mode 100644 index 000000000..f4dd1ed31 --- /dev/null +++ b/packages/instrumentation-llamaindex/test/finish_reasons.test.ts @@ -0,0 +1,58 @@ +/** + * Unit tests for openAIFinishReasonMap. + * + * Each OpenAI raw finish reason value is tested individually. + * Verified values from OpenAI API documentation. + */ + +import * as assert from "assert"; +import { FinishReasons } from "@traceloop/ai-semantic-conventions"; +import { openAIFinishReasonMap } from "../src/custom-llm-instrumentation"; + +const VALID_OTEL_FINISH_REASONS = new Set([ + FinishReasons.STOP, + FinishReasons.LENGTH, + FinishReasons.TOOL_CALL, + FinishReasons.CONTENT_FILTER, + FinishReasons.ERROR, +]); + +describe("openAIFinishReasonMap", () => { + it("all mapped values are valid OTel finish reason strings", () => { + for (const [raw, otel] of Object.entries(openAIFinishReasonMap)) { + assert.ok( + VALID_OTEL_FINISH_REASONS.has(otel), + `openAIFinishReasonMap["${raw}"] = "${otel}" is not a valid OTel finish reason`, + ); + } + }); + + it('maps "stop" to stop', () => { + assert.strictEqual(openAIFinishReasonMap["stop"], FinishReasons.STOP); + }); + + it('maps "length" to length', () => { + assert.strictEqual(openAIFinishReasonMap["length"], FinishReasons.LENGTH); + }); + + it('maps "tool_calls" to tool_call', () => { + assert.strictEqual( + openAIFinishReasonMap["tool_calls"], + FinishReasons.TOOL_CALL, + ); + }); + + it('maps "content_filter" to content_filter', () => { + assert.strictEqual( + openAIFinishReasonMap["content_filter"], + FinishReasons.CONTENT_FILTER, + ); + }); + + it('maps "function_call" to tool_call (deprecated alias)', () => { + assert.strictEqual( + openAIFinishReasonMap["function_call"], + FinishReasons.TOOL_CALL, + ); + }); +}); diff --git a/packages/instrumentation-llamaindex/test/instrumentation.test.ts b/packages/instrumentation-llamaindex/test/instrumentation.test.ts index ab1c87227..6e4153411 100644 --- a/packages/instrumentation-llamaindex/test/instrumentation.test.ts +++ b/packages/instrumentation-llamaindex/test/instrumentation.test.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { context } from "@opentelemetry/api"; +import { context, diag, DiagLogger } from "@opentelemetry/api"; import { AsyncHooksContextManager } from "@opentelemetry/context-async-hooks"; import { LlamaIndexInstrumentation } from "../src/instrumentation"; import * as assert from "assert"; @@ -24,6 +24,26 @@ import { SimpleSpanProcessor, } from "@opentelemetry/sdk-trace-node"; import type * as llamaindexImport from "llamaindex"; +import { + ATTR_GEN_AI_INPUT_MESSAGES, + ATTR_GEN_AI_OPERATION_NAME, + ATTR_GEN_AI_OUTPUT_MESSAGES, + ATTR_GEN_AI_PROVIDER_NAME, + ATTR_GEN_AI_REQUEST_MODEL, + ATTR_GEN_AI_REQUEST_TEMPERATURE, + ATTR_GEN_AI_RESPONSE_FINISH_REASONS, + ATTR_GEN_AI_RESPONSE_ID, + ATTR_GEN_AI_RESPONSE_MODEL, + ATTR_GEN_AI_USAGE_INPUT_TOKENS, + ATTR_GEN_AI_USAGE_OUTPUT_TOKENS, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + GEN_AI_PROVIDER_NAME_VALUE_OPENAI, +} from "@opentelemetry/semantic-conventions/incubating"; +import { + SpanAttributes, + FinishReasons, +} from "@traceloop/ai-semantic-conventions"; +import { CustomLLMInstrumentation } from "../src/custom-llm-instrumentation"; import { Polly, setupMocha as setupPolly } from "@pollyjs/core"; import NodeHttpAdapter from "@pollyjs/adapter-node-http"; @@ -128,60 +148,523 @@ describe("Test LlamaIndex instrumentation", async function () { ); }).timeout(60000); - it.skip("should build proper trace on streaming query engine", async () => { - const directoryReader = new llamaindex.SimpleDirectoryReader(); - const documents = await directoryReader.loadData({ - directoryPath: "test/data", - }); + it.skip( + "should build proper trace on streaming query engine (legacy)", + async () => { + const directoryReader = new llamaindex.SimpleDirectoryReader(); + const documents = await directoryReader.loadData({ + directoryPath: "test/data", + }); - const index = await llamaindex.VectorStoreIndex.fromDocuments(documents); - const queryEngine = index.asQueryEngine(); + const index = await llamaindex.VectorStoreIndex.fromDocuments(documents); + const queryEngine = index.asQueryEngine(); - const result = await queryEngine.query({ - query: "Where was albert einstein born?", - stream: true, - }); + const result = await queryEngine.query({ + query: "Where was albert einstein born?", + stream: true, + }); + + for await (const res of result) { + assert.ok(res); + } + + const spans = memoryExporter.getFinishedSpans(); + + const retrieverQueryEngineQuerySpan = spans.find( + (span) => span.name === "retriever_query_engine.query", + ); + const synthesizeSpan = spans.find( + (span) => span.name === "base_synthesizer.synthesize", + ); + const retrieverQueryEngineRetrieveSpan = spans.find( + (span) => span.name === "retriever_query_engine.retrieve", + ); + const openAIEmbeddingSpan = spans.find( + (span) => span.name === "open_ai_embedding.get_query_embedding", + ); + const vectorIndexRetrieverSpan = spans.find( + (span) => span.name === "vector_index_retriever.retrieve", + ); + + assert.strictEqual( + synthesizeSpan?.parentSpanId, + retrieverQueryEngineQuerySpan?.spanContext().spanId, + ); - for await (const res of result) { - assert.ok(res); + assert.strictEqual( + retrieverQueryEngineRetrieveSpan?.parentSpanId, + retrieverQueryEngineQuerySpan?.spanContext().spanId, + ); + + assert.strictEqual( + vectorIndexRetrieverSpan?.parentSpanId, + retrieverQueryEngineRetrieveSpan?.spanContext().spanId, + ); + + assert.strictEqual( + openAIEmbeddingSpan?.parentSpanId, + vectorIndexRetrieverSpan?.spanContext().spanId, + ); + }, + ).timeout(60000); +}); + +// ───────────────────────────────────────────────────────────────────────────── +// OTel 1.40 migration tests — CustomLLMInstrumentation with mock LLM +// No HTTP, no Polly, no API keys needed. +// ───────────────────────────────────────────────────────────────────────────── + +const testDiag: DiagLogger = diag; + +function makeMockChat(options: { + responseContent?: string; + finishReason?: string; + promptTokens?: number; + completionTokens?: number; +}) { + const responseContent = options.responseContent ?? "Hello!"; + const finishReason = options.finishReason ?? "stop"; + const promptTokens = options.promptTokens ?? 10; + const completionTokens = options.completionTokens ?? 5; + + return async function chat({ stream }: any) { + if (stream) { + async function* generate() { + yield { delta: responseContent }; + } + return generate(); } + return { + message: { role: "assistant", content: responseContent }, + raw: { + id: "chatcmpl-test123", + choices: [{ finish_reason: finishReason }], + usage: { + prompt_tokens: promptTokens, + completion_tokens: completionTokens, + total_tokens: promptTokens + completionTokens, + }, + }, + }; + }; +} - const spans = memoryExporter.getFinishedSpans(); +function makeMockChatWithStreamUsage(options: { + responseContent?: string; + finishReason?: string; + promptTokens?: number; + completionTokens?: number; +}) { + const responseContent = options.responseContent ?? "Hello!"; + const finishReason = options.finishReason ?? "stop"; + const promptTokens = options.promptTokens ?? 10; + const completionTokens = options.completionTokens ?? 5; - const retrieverQueryEngineQuerySpan = spans.find( - (span) => span.name === "retriever_query_engine.query", - ); - const synthesizeSpan = spans.find( - (span) => span.name === "base_synthesizer.synthesize", - ); - const retrieverQueryEngineRetrieveSpan = spans.find( - (span) => span.name === "retriever_query_engine.retrieve", - ); - const openAIEmbeddingSpan = spans.find( - (span) => span.name === "open_ai_embedding.get_query_embedding", - ); - const vectorIndexRetrieverSpan = spans.find( - (span) => span.name === "vector_index_retriever.retrieve", - ); + return async function chat({ stream }: any) { + if (stream) { + async function* generate() { + yield { delta: "partial " }; + yield { + delta: responseContent, + raw: { + id: "chatcmpl-test123", + choices: [{ finish_reason: finishReason }], + usage: { + prompt_tokens: promptTokens, + completion_tokens: completionTokens, + total_tokens: promptTokens + completionTokens, + }, + }, + }; + } + return generate(); + } + return { + message: { role: "assistant", content: responseContent }, + raw: { + choices: [{ finish_reason: finishReason }], + usage: { + prompt_tokens: promptTokens, + completion_tokens: completionTokens, + total_tokens: promptTokens + completionTokens, + }, + }, + }; + }; +} - assert.strictEqual( - synthesizeSpan?.parentSpanId, - retrieverQueryEngineQuerySpan?.spanContext().spanId, - ); +describe("CustomLLMInstrumentation — OTel 1.40 attributes", () => { + const otelExporter = new InMemorySpanExporter(); + const otelProvider = new NodeTracerProvider({ + spanProcessors: [new SimpleSpanProcessor(otelExporter)], + }); + let contextManager: AsyncHooksContextManager; - assert.strictEqual( - retrieverQueryEngineRetrieveSpan?.parentSpanId, - retrieverQueryEngineQuerySpan?.spanContext().spanId, - ); + const mockLLMMeta = { model: "gpt-4o", topP: 1, temperature: 0.7 }; - assert.strictEqual( - vectorIndexRetrieverSpan?.parentSpanId, - retrieverQueryEngineRetrieveSpan?.spanContext().spanId, - ); + before(() => { + otelProvider.register(); + }); - assert.strictEqual( - openAIEmbeddingSpan?.parentSpanId, - vectorIndexRetrieverSpan?.spanContext().spanId, + beforeEach(() => { + contextManager = new AsyncHooksContextManager().enable(); + context.setGlobalContextManager(contextManager); + otelExporter.reset(); + }); + + afterEach(() => { + context.disable(); + }); + + function makeInstrumentation(traceContent = true) { + return new CustomLLMInstrumentation( + () => ({ traceContent }), + testDiag, + () => otelProvider.getTracer("test"), ); - }).timeout(60000); + } + + describe("traceContent: true", () => { + it("sets span name as 'chat {model}'", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({}); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const spans = otelExporter.getFinishedSpans(); + assert.strictEqual(spans.length, 1); + assert.strictEqual(spans[0].name, "chat gpt-4o"); + }); + + it("sets gen_ai.provider.name to openai", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({}); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.strictEqual( + span.attributes[ATTR_GEN_AI_PROVIDER_NAME], + GEN_AI_PROVIDER_NAME_VALUE_OPENAI, + ); + }); + + it("sets gen_ai.operation.name to chat", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({}); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.strictEqual( + span.attributes[ATTR_GEN_AI_OPERATION_NAME], + GEN_AI_OPERATION_NAME_VALUE_CHAT, + ); + }); + + it("sets request and response model", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({}); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.strictEqual(span.attributes[ATTR_GEN_AI_REQUEST_MODEL], "gpt-4o"); + assert.strictEqual(span.attributes[ATTR_GEN_AI_RESPONSE_MODEL], "gpt-4o"); + }); + + it("sets gen_ai.request.temperature", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({}); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.strictEqual(span.attributes[ATTR_GEN_AI_REQUEST_TEMPERATURE], 0.7); + }); + + it("sets gen_ai.input.messages with correct structure", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({}); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "What is 2+2?" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + const inputMessages = JSON.parse( + span.attributes[ATTR_GEN_AI_INPUT_MESSAGES] as string, + ); + assert.strictEqual(inputMessages[0].role, "user"); + assert.strictEqual(inputMessages[0].parts[0].type, "text"); + assert.strictEqual(inputMessages[0].parts[0].content, "What is 2+2?"); + }); + + it("sets gen_ai.output.messages with correct structure", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({ responseContent: "The answer is 4." }); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + const outputMessages = JSON.parse( + span.attributes[ATTR_GEN_AI_OUTPUT_MESSAGES] as string, + ); + assert.strictEqual(outputMessages[0].role, "assistant"); + assert.strictEqual(outputMessages[0].parts[0].type, "text"); + assert.strictEqual( + outputMessages[0].parts[0].content, + "The answer is 4.", + ); + }); + + it("sets gen_ai.response.finish_reasons (metadata — always set)", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({ finishReason: "stop" }); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.deepStrictEqual( + span.attributes[ATTR_GEN_AI_RESPONSE_FINISH_REASONS], + [FinishReasons.STOP], + ); + }); + + it("sets gen_ai.response.id", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({}); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.strictEqual( + span.attributes[ATTR_GEN_AI_RESPONSE_ID], + "chatcmpl-test123", + ); + }); + + it("unknown finish_reason passes through as-is to span attribute", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({ finishReason: "some_future_reason" }); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.deepStrictEqual( + span.attributes[ATTR_GEN_AI_RESPONSE_FINISH_REASONS], + ["some_future_reason"], + ); + }); + + it("sets token usage attributes", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({ promptTokens: 10, completionTokens: 5 }); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.strictEqual(span.attributes[ATTR_GEN_AI_USAGE_INPUT_TOKENS], 10); + assert.strictEqual(span.attributes[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS], 5); + assert.strictEqual( + span.attributes[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS], + 15, + ); + }); + }); + + describe("traceContent: false", () => { + it("does NOT set gen_ai.input.messages", async () => { + const instr = makeInstrumentation(false); + const chat = makeMockChat({}); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.strictEqual( + span.attributes[ATTR_GEN_AI_INPUT_MESSAGES], + undefined, + ); + }); + + it("does NOT set gen_ai.output.messages", async () => { + const instr = makeInstrumentation(false); + const chat = makeMockChat({}); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.strictEqual( + span.attributes[ATTR_GEN_AI_OUTPUT_MESSAGES], + undefined, + ); + }); + + it("still sets finish_reasons (metadata, not content)", async () => { + const instr = makeInstrumentation(false); + const chat = makeMockChat({ finishReason: "stop" }); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.deepStrictEqual( + span.attributes[ATTR_GEN_AI_RESPONSE_FINISH_REASONS], + [FinishReasons.STOP], + ); + }); + + it("still sets token usage (metadata, not content)", async () => { + const instr = makeInstrumentation(false); + const chat = makeMockChat({ promptTokens: 10, completionTokens: 5 }); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }] }, + ); + + const span = otelExporter.getFinishedSpans()[0]; + assert.strictEqual(span.attributes[ATTR_GEN_AI_USAGE_INPUT_TOKENS], 10); + assert.strictEqual(span.attributes[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS], 5); + }); + }); + + describe("streaming", () => { + it("sets gen_ai.output.messages after stream completes", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({ responseContent: "streamed response" }); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + const stream = await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }], stream: true }, + ); + + for await (const _chunk of stream) { + /* consume stream */ + } + + const span = otelExporter.getFinishedSpans()[0]; + const outputMessages = JSON.parse( + span.attributes[ATTR_GEN_AI_OUTPUT_MESSAGES] as string, + ); + assert.strictEqual(outputMessages[0].role, "assistant"); + assert.strictEqual( + outputMessages[0].parts[0].content, + "streamed response", + ); + }); + + it("does NOT set finish_reasons in streaming when mock omits them", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChat({}); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + const stream = await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }], stream: true }, + ); + + for await (const _chunk of stream) { + /* consume stream */ + } + + const span = otelExporter.getFinishedSpans()[0]; + assert.strictEqual( + span.attributes[ATTR_GEN_AI_RESPONSE_FINISH_REASONS], + undefined, + ); + }); + + it("streaming with traceContent=false: does NOT set gen_ai.output.messages", async () => { + const instr = makeInstrumentation(false); + const chat = makeMockChat({}); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + const stream = await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }], stream: true }, + ); + + for await (const _chunk of stream) { + /* consume stream */ + } + + const span = otelExporter.getFinishedSpans()[0]; + assert.strictEqual( + span.attributes[ATTR_GEN_AI_OUTPUT_MESSAGES], + undefined, + ); + }); + + it("sets finish_reasons and usage when streaming with raw data in last chunk", async () => { + const instr = makeInstrumentation(); + const chat = makeMockChatWithStreamUsage({ + finishReason: "stop", + promptTokens: 10, + completionTokens: 5, + }); + const wrapped = instr.chatWrapper({ className: "OpenAI" })(chat as any); + const stream = await wrapped.call( + { metadata: mockLLMMeta }, + { messages: [{ role: "user", content: "hi" }], stream: true }, + ); + + for await (const _chunk of stream) { + /* consume stream */ + } + + const span = otelExporter.getFinishedSpans()[0]; + assert.deepStrictEqual( + span.attributes[ATTR_GEN_AI_RESPONSE_FINISH_REASONS], + [FinishReasons.STOP], + ); + assert.strictEqual(span.attributes[ATTR_GEN_AI_USAGE_INPUT_TOKENS], 10); + assert.strictEqual(span.attributes[ATTR_GEN_AI_USAGE_OUTPUT_TOKENS], 5); + assert.strictEqual( + span.attributes[SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS], + 15, + ); + assert.strictEqual( + span.attributes[ATTR_GEN_AI_RESPONSE_ID], + "chatcmpl-test123", + ); + }); + }); }); diff --git a/packages/instrumentation-llamaindex/test/semconv.test.ts b/packages/instrumentation-llamaindex/test/semconv.test.ts new file mode 100644 index 000000000..fddd697ca --- /dev/null +++ b/packages/instrumentation-llamaindex/test/semconv.test.ts @@ -0,0 +1,252 @@ +/** + * OTel GenAI Semantic Convention compliance tests for LlamaIndex instrumentation. + * + * Pure unit tests — no HTTP, no Polly, no API keys needed. + * They validate: + * - openAIFinishReasonMap covers all known OpenAI values and produces valid OTel values + * - mapOpenAIContentBlock produces schema-compliant OTel parts + * - formatInputMessages / formatOutputMessage produce valid OTel JSON + * - traceContent: false omits content but keeps metadata + */ + +import * as assert from "assert"; +import { FinishReasons } from "@traceloop/ai-semantic-conventions"; +import { + formatInputMessages, + formatOutputMessage, + mapOpenAIContentBlock, +} from "@traceloop/instrumentation-utils"; +import { GEN_AI_OPERATION_NAME_VALUE_CHAT } from "@opentelemetry/semantic-conventions/incubating"; +import { openAIFinishReasonMap } from "../src/custom-llm-instrumentation"; + +// ───────────────────────────────────────────────────────────────────────────── +// Helpers +// ───────────────────────────────────────────────────────────────────────────── + +function assertValidOtelJsonArray(value: unknown, label: string): any[] { + assert.ok(typeof value === "string", `${label} must be a string`); + let parsed: any; + try { + parsed = JSON.parse(value as string); + } catch { + assert.fail(`${label} is not valid JSON: ${value}`); + } + assert.ok(Array.isArray(parsed), `${label} must be a JSON array`); + assert.ok(parsed.length > 0, `${label} must not be empty`); + return parsed; +} + +// ───────────────────────────────────────────────────────────────────────────── +// P1-1: Provider and operation name constants +// ───────────────────────────────────────────────────────────────────────────── + +describe("OTel provider and operation name constants", () => { + it("GEN_AI_OPERATION_NAME_VALUE_CHAT is chat", () => { + assert.strictEqual(GEN_AI_OPERATION_NAME_VALUE_CHAT, "chat"); + }); + + it("span name format is 'chat {model}'", () => { + const model = "gpt-4o"; + assert.strictEqual(`chat ${model}`, "chat gpt-4o"); + }); +}); + +// ───────────────────────────────────────────────────────────────────────────── +// P1-2: mapOpenAIContentBlock produces valid OTel parts +// ───────────────────────────────────────────────────────────────────────────── + +describe("mapOpenAIContentBlock produces valid OTel parts", () => { + it("string input → TextPart", () => { + const part = mapOpenAIContentBlock("hello world") as any; + assert.strictEqual(part.type, "text"); + assert.strictEqual(part.content, "hello world"); + }); + + it("text block → TextPart", () => { + const part = mapOpenAIContentBlock({ type: "text", text: "hello" }) as any; + assert.strictEqual(part.type, "text"); + assert.strictEqual(part.content, "hello"); + }); + + it("image_url block with http URL → UriPart", () => { + const part = mapOpenAIContentBlock({ + type: "image_url", + image_url: { url: "https://example.com/img.png" }, + }) as any; + assert.strictEqual(part.type, "uri"); + assert.strictEqual(part.modality, "image"); + assert.strictEqual(part.uri, "https://example.com/img.png"); + }); + + it("image_url block with base64 data URI → BlobPart", () => { + const part = mapOpenAIContentBlock({ + type: "image_url", + image_url: { url: "data:image/png;base64,abc123" }, + }) as any; + assert.strictEqual(part.type, "blob"); + assert.strictEqual(part.modality, "image"); + assert.strictEqual(part.mime_type, "image/png"); + assert.strictEqual(part.content, "abc123"); + }); + + it("empty text string → TextPart with empty content", () => { + const part = mapOpenAIContentBlock({ type: "text", text: "" }) as any; + assert.strictEqual(part.type, "text"); + assert.strictEqual(part.content, ""); + }); + + it("unknown block type → preserved as GenericPart with all fields", () => { + const block = { type: "future_type", field1: "a", field2: 42 }; + const part = mapOpenAIContentBlock(block) as any; + assert.strictEqual(part.type, "future_type"); + assert.strictEqual(part.field1, "a"); + assert.strictEqual(part.field2, 42); + }); +}); + +// ───────────────────────────────────────────────────────────────────────────── +// P2-1: formatInputMessages produces valid OTel gen_ai.input.messages JSON +// ───────────────────────────────────────────────────────────────────────────── + +describe("formatInputMessages produces valid gen_ai.input.messages JSON", () => { + it("simple string content user message", () => { + const json = formatInputMessages( + [{ role: "user", content: "What is 2+2?" }], + mapOpenAIContentBlock, + ); + const messages = assertValidOtelJsonArray(json, "gen_ai.input.messages"); + assert.strictEqual(messages[0].role, "user"); + assert.strictEqual(messages[0].parts[0].type, "text"); + assert.strictEqual(messages[0].parts[0].content, "What is 2+2?"); + }); + + it("multi-turn conversation", () => { + const json = formatInputMessages( + [ + { role: "user", content: "Hello" }, + { role: "assistant", content: [{ type: "text", text: "Hi there!" }] }, + { role: "user", content: "How are you?" }, + ], + mapOpenAIContentBlock, + ); + const messages = assertValidOtelJsonArray(json, "gen_ai.input.messages"); + assert.strictEqual(messages.length, 3); + assert.strictEqual(messages[0].role, "user"); + assert.strictEqual(messages[1].role, "assistant"); + assert.strictEqual(messages[1].parts[0].type, "text"); + }); + + it("traceContent=false → role preserved, parts empty", () => { + const json = formatInputMessages( + [{ role: "user", content: "secret content" }], + mapOpenAIContentBlock, + false, + ); + const messages = JSON.parse(json); + assert.strictEqual(messages[0].role, "user"); + assert.ok( + !messages[0].parts || messages[0].parts.length === 0, + "parts must be empty when traceContent=false", + ); + }); +}); + +// ───────────────────────────────────────────────────────────────────────────── +// P2-2: formatOutputMessage produces valid OTel gen_ai.output.messages JSON +// ───────────────────────────────────────────────────────────────────────────── + +describe("formatOutputMessage produces valid gen_ai.output.messages JSON", () => { + it("string content wrapped in array → TextPart", () => { + const json = formatOutputMessage( + ["The answer is 4."], + "stop", + openAIFinishReasonMap, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + mapOpenAIContentBlock, + ); + const messages = assertValidOtelJsonArray(json, "gen_ai.output.messages"); + assert.strictEqual(messages[0].role, "assistant"); + assert.strictEqual(messages[0].finish_reason, FinishReasons.STOP); + assert.strictEqual(messages[0].parts[0].type, "text"); + assert.strictEqual(messages[0].parts[0].content, "The answer is 4."); + }); + + it("block array content → mapped parts", () => { + const json = formatOutputMessage( + [{ type: "text", text: "Hello!" }], + "stop", + openAIFinishReasonMap, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + mapOpenAIContentBlock, + ); + const messages = JSON.parse(json); + assert.strictEqual(messages[0].parts[0].type, "text"); + assert.strictEqual(messages[0].parts[0].content, "Hello!"); + }); + + it("maps OpenAI finish reasons to OTel values", () => { + const cases: Array<[string, string]> = [ + ["stop", FinishReasons.STOP], + ["length", FinishReasons.LENGTH], + ["tool_calls", FinishReasons.TOOL_CALL], + ["content_filter", FinishReasons.CONTENT_FILTER], + ["function_call", FinishReasons.TOOL_CALL], + ]; + for (const [raw, expected] of cases) { + const json = formatOutputMessage( + ["ok"], + raw, + openAIFinishReasonMap, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + mapOpenAIContentBlock, + ); + const messages = JSON.parse(json); + assert.strictEqual( + messages[0].finish_reason, + expected, + `"${raw}" should map to "${expected}"`, + ); + } + }); + + it("null finish_reason → finish_reason is empty string (not null)", () => { + const json = formatOutputMessage( + ["ok"], + null, + openAIFinishReasonMap, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + mapOpenAIContentBlock, + ); + const messages = JSON.parse(json); + assert.strictEqual(messages[0].finish_reason, ""); + }); + + it("unknown finish_reason passes through unchanged", () => { + const json = formatOutputMessage( + ["ok"], + "some_future_reason", + openAIFinishReasonMap, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + mapOpenAIContentBlock, + ); + const messages = JSON.parse(json); + assert.strictEqual(messages[0].finish_reason, "some_future_reason"); + }); + + it("traceContent=false → finish_reason preserved, parts empty", () => { + const json = formatOutputMessage( + ["secret content"], + "stop", + openAIFinishReasonMap, + GEN_AI_OPERATION_NAME_VALUE_CHAT, + mapOpenAIContentBlock, + false, + ); + const messages = JSON.parse(json); + assert.strictEqual(messages[0].finish_reason, FinishReasons.STOP); + assert.ok( + !messages[0].parts || messages[0].parts.length === 0, + "parts must be empty when traceContent=false", + ); + }); +}); diff --git a/packages/sample-app/src/sample_llama_index_openai_agent.ts b/packages/sample-app/src/sample_llama_index_openai_agent.ts index 751f1baaf..2f7324e75 100644 --- a/packages/sample-app/src/sample_llama_index_openai_agent.ts +++ b/packages/sample-app/src/sample_llama_index_openai_agent.ts @@ -1,5 +1,6 @@ import * as llamaindex from "llamaindex"; import * as traceloop from "@traceloop/node-server-sdk"; +import * as llamaIndexOpenAI from "@llamaindex/openai"; import { OpenAIAgent, OpenAI as LLamaOpenAI } from "@llamaindex/openai"; traceloop.initialize({ @@ -8,6 +9,7 @@ traceloop.initialize({ disableBatch: true, instrumentModules: { llamaIndex: llamaindex, + llamaIndexOpenAI, }, }); diff --git a/packages/sample-app/src/sample_llamaindex.ts b/packages/sample-app/src/sample_llamaindex.ts index b03ea5054..2efbc8d5b 100644 --- a/packages/sample-app/src/sample_llamaindex.ts +++ b/packages/sample-app/src/sample_llamaindex.ts @@ -1,5 +1,7 @@ import * as traceloop from "@traceloop/node-server-sdk"; +import * as llamaindex from "llamaindex"; import { VectorStoreIndex, Document, Settings } from "llamaindex"; +import * as llamaIndexOpenAI from "@llamaindex/openai"; import { OpenAIEmbedding, OpenAI } from "@llamaindex/openai"; import { readFile } from "fs/promises"; @@ -7,10 +9,17 @@ traceloop.initialize({ appName: "sample_llamaindex", apiKey: process.env.TRACELOOP_API_KEY, disableBatch: true, + instrumentModules: { + llamaIndex: llamaindex, + llamaIndexOpenAI, + }, }); Settings.embedModel = new OpenAIEmbedding(); -Settings.llm = new OpenAI(); +// OpenAI only sends usage in the final streaming chunk if stream_options: { include_usage: true } +Settings.llm = new OpenAI({ + additionalChatOptions: { stream_options: { include_usage: true } }, +}); class SampleLlamaIndex { async query() { diff --git a/packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts b/packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts index 2293c5ca0..22d991ce5 100644 --- a/packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts +++ b/packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts @@ -97,6 +97,8 @@ export interface InitializeOptions { together?: typeof together.Together; langchain?: boolean; llamaIndex?: typeof llamaindex; + /** Only meaningful when `llamaIndex` is also provided. */ + llamaIndexOpenAI?: any; chromadb?: typeof chromadb; qdrant?: typeof qdrant; mcp?: typeof mcp; diff --git a/packages/traceloop-sdk/src/lib/tracing/index.ts b/packages/traceloop-sdk/src/lib/tracing/index.ts index ac7bf40cf..2d2c92a57 100644 --- a/packages/traceloop-sdk/src/lib/tracing/index.ts +++ b/packages/traceloop-sdk/src/lib/tracing/index.ts @@ -218,7 +218,10 @@ export const manuallyInitInstrumentations = ( exceptionLogger, }); instrumentations.push(llamaIndexInstrumentation); - llamaIndexInstrumentation.manuallyInstrument(instrumentModules.llamaIndex); + llamaIndexInstrumentation.manuallyInstrument( + instrumentModules.llamaIndex, + instrumentModules.llamaIndexOpenAI, + ); } if (instrumentModules?.chromadb) { diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index d18f2cf0c..f537a3bcd 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -415,11 +415,14 @@ importers: specifier: ^0.203.0 version: 0.203.0(@opentelemetry/api@1.9.0) '@opentelemetry/semantic-conventions': - specifier: ^1.38.0 - version: 1.38.0 + specifier: ^1.40.0 + version: 1.40.0 '@traceloop/ai-semantic-conventions': specifier: workspace:* version: link:../ai-semantic-conventions + '@traceloop/instrumentation-utils': + specifier: workspace:* + version: link:../instrumentation-utils lodash: specifier: ^4.17.21 version: 4.17.21