From ac3cedaa529fb66ef6fee3b821cdb33eda92404f Mon Sep 17 00:00:00 2001 From: Nir Gazit Date: Fri, 1 Dec 2023 13:20:27 +0100 Subject: [PATCH] fix: allow disabling tracing of prompts through config --- package-lock.json | 7 - .../src/instrumentation.ts | 227 ++++++++++-------- packages/instrumentation-openai/src/types.ts | 8 +- packages/traceloop-sdk/package.json | 1 - .../initialize-options.interface.ts | 6 + .../traceloop-sdk/src/lib/tracing/index.ts | 7 + 6 files changed, 141 insertions(+), 115 deletions(-) diff --git a/package-lock.json b/package-lock.json index b1889766..55d8410e 100644 --- a/package-lock.json +++ b/package-lock.json @@ -4105,12 +4105,6 @@ "version": "1.1.0", "license": "BSD-3-Clause" }, - "node_modules/@scarf/scarf": { - "version": "1.3.0", - "resolved": "https://registry.npmjs.org/@scarf/scarf/-/scarf-1.3.0.tgz", - "integrity": "sha512-lHKK8M5CTcpFj2hZDB3wIjb0KAbEOgDmiJGDv1WBRfQgRm/a8/XMEkG/N1iM01xgbUDsPQwi42D+dFo1XPAKew==", - "hasInstallScript": true - }, "node_modules/@sigstore/bundle": { "version": "1.1.0", "dev": true, @@ -12179,7 +12173,6 @@ "dependencies": { "@opentelemetry/exporter-trace-otlp-proto": "^0.44.0", "@opentelemetry/sdk-node": "^0.44.0", - "@scarf/scarf": "^1.3.0", "@traceloop/instrumentation-openai": "^0.0.22", "@types/nunjucks": "^3.2.5", "fetch-retry": "^5.0.6", diff --git a/packages/instrumentation-openai/src/instrumentation.ts b/packages/instrumentation-openai/src/instrumentation.ts index c4cd565c..2b5cfb14 100644 --- a/packages/instrumentation-openai/src/instrumentation.ts +++ b/packages/instrumentation-openai/src/instrumentation.ts @@ -38,10 +38,16 @@ import { } from "openai/resources"; export class OpenAIInstrumentation extends InstrumentationBase { + protected override _config!: OpenAIInstrumentationConfig; + constructor(config: OpenAIInstrumentationConfig = {}) { super("@traceloop/instrumentation-openai", "0.0.17", config); } + public override setConfig(config: OpenAIInstrumentationConfig = {}) { + super.setConfig(config); + } + public manuallyInstrument( module: typeof openai.OpenAI & { openLLMetryPatched?: boolean }, ) { @@ -177,7 +183,12 @@ export class OpenAIInstrumentation extends InstrumentationBase { }, ); - const wrappedPromise = wrapPromise(type, version, span, execPromise); + const wrappedPromise = plugin._wrapPromise( + type, + version, + span, + execPromise, + ); return context.bind(execContext, wrappedPromise as any); }; @@ -232,7 +243,7 @@ export class OpenAIInstrumentation extends InstrumentationBase { }); } - if (shouldSendPrompts()) { + if (this._shouldSendPrompts()) { if (type === "chat") { params.messages.forEach((message, index) => { attributes[`${SpanAttributes.LLM_PROMPTS}.${index}.role`] = @@ -253,124 +264,128 @@ export class OpenAIInstrumentation extends InstrumentationBase { attributes, }); } -} -function wrapPromise( - type: "chat" | "completion", - version: "v3" | "v4", - span: Span, - promise: Promise, -): Promise { - return promise - .then((result) => { - return new Promise((resolve) => { - if (version === "v3") { - if (type === "chat") { - endSpan({ - type, - span, - result: (result as any).data as ChatCompletion, - }); - } else { - endSpan({ type, span, result: (result as any).data as Completion }); - } - } else { - if (type === "chat") { - endSpan({ type, span, result: result as ChatCompletion }); + private _wrapPromise( + type: "chat" | "completion", + version: "v3" | "v4", + span: Span, + promise: Promise, + ): Promise { + return promise + .then((result) => { + return new Promise((resolve) => { + if (version === "v3") { + if (type === "chat") { + this._endSpan({ + type, + span, + result: (result as any).data as ChatCompletion, + }); + } else { + this._endSpan({ + type, + span, + result: (result as any).data as Completion, + }); + } } else { - endSpan({ type, span, result: result as Completion }); + if (type === "chat") { + this._endSpan({ type, span, result: result as ChatCompletion }); + } else { + this._endSpan({ type, span, result: result as Completion }); + } } - } - resolve(result); - }); - }) - .catch((error: Error) => { - return new Promise((_, reject) => { - span.setStatus({ - code: SpanStatusCode.ERROR, - message: error.message, + resolve(result); }); - span.recordException(error); - span.end(); + }) + .catch((error: Error) => { + return new Promise((_, reject) => { + span.setStatus({ + code: SpanStatusCode.ERROR, + message: error.message, + }); + span.recordException(error); + span.end(); - reject(error); + reject(error); + }); }); - }); -} - -function endSpan({ - span, - type, - result, -}: - | { span: Span; type: "chat"; result: ChatCompletion } - | { span: Span; type: "completion"; result: Completion }) { - span.setAttribute(SpanAttributes.LLM_RESPONSE_MODEL, result.model); - if (result.usage) { - span.setAttribute( - SpanAttributes.LLM_USAGE_TOTAL_TOKENS, - result.usage?.total_tokens, - ); - span.setAttribute( - SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, - result.usage?.completion_tokens, - ); - span.setAttribute( - SpanAttributes.LLM_USAGE_PROMPT_TOKENS, - result.usage?.prompt_tokens, - ); } - if (shouldSendPrompts()) { - if (type === "chat") { - result.choices.forEach((choice, index) => { - span.setAttribute( - `${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`, - choice.finish_reason, - ); - span.setAttribute( - `${SpanAttributes.LLM_COMPLETIONS}.${index}.role`, - choice.message.role, - ); - span.setAttribute( - `${SpanAttributes.LLM_COMPLETIONS}.${index}.content`, - choice.message.content ?? "", - ); + private _endSpan({ + span, + type, + result, + }: + | { span: Span; type: "chat"; result: ChatCompletion } + | { span: Span; type: "completion"; result: Completion }) { + span.setAttribute(SpanAttributes.LLM_RESPONSE_MODEL, result.model); + if (result.usage) { + span.setAttribute( + SpanAttributes.LLM_USAGE_TOTAL_TOKENS, + result.usage?.total_tokens, + ); + span.setAttribute( + SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, + result.usage?.completion_tokens, + ); + span.setAttribute( + SpanAttributes.LLM_USAGE_PROMPT_TOKENS, + result.usage?.prompt_tokens, + ); + } - if (choice.message.function_call) { + if (this._shouldSendPrompts()) { + if (type === "chat") { + result.choices.forEach((choice, index) => { span.setAttribute( - `${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.name`, - choice.message.function_call.name, + `${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`, + choice.finish_reason, ); span.setAttribute( - `${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.arguments`, - choice.message.function_call.arguments, + `${SpanAttributes.LLM_COMPLETIONS}.${index}.role`, + choice.message.role, ); - } - }); - } else { - result.choices.forEach((choice, index) => { - span.setAttribute( - `${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`, - choice.finish_reason, - ); - span.setAttribute( - `${SpanAttributes.LLM_COMPLETIONS}.${index}.role`, - "assistant", - ); - span.setAttribute( - `${SpanAttributes.LLM_COMPLETIONS}.${index}.content`, - choice.text, - ); - }); + span.setAttribute( + `${SpanAttributes.LLM_COMPLETIONS}.${index}.content`, + choice.message.content ?? "", + ); + + if (choice.message.function_call) { + span.setAttribute( + `${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.name`, + choice.message.function_call.name, + ); + span.setAttribute( + `${SpanAttributes.LLM_COMPLETIONS}.${index}.function_call.arguments`, + choice.message.function_call.arguments, + ); + } + }); + } else { + result.choices.forEach((choice, index) => { + span.setAttribute( + `${SpanAttributes.LLM_COMPLETIONS}.${index}.finish_reason`, + choice.finish_reason, + ); + span.setAttribute( + `${SpanAttributes.LLM_COMPLETIONS}.${index}.role`, + "assistant", + ); + span.setAttribute( + `${SpanAttributes.LLM_COMPLETIONS}.${index}.content`, + choice.text, + ); + }); + } } - } - span.end(); -} + span.end(); + } -function shouldSendPrompts() { - return ( - (process.env.TRACELOOP_TRACE_CONTENT || "true").toLowerCase() === "true" - ); + private _shouldSendPrompts() { + return this._config.traceContent !== undefined + ? this._config.traceContent + : true; + } } diff --git a/packages/instrumentation-openai/src/types.ts b/packages/instrumentation-openai/src/types.ts index d67bb7e0..c441d1c4 100644 --- a/packages/instrumentation-openai/src/types.ts +++ b/packages/instrumentation-openai/src/types.ts @@ -1,3 +1,9 @@ import { InstrumentationConfig } from "@opentelemetry/instrumentation"; -export interface OpenAIInstrumentationConfig extends InstrumentationConfig {} +export interface OpenAIInstrumentationConfig extends InstrumentationConfig { + /** + * Whether to log prompts, completions and embeddings on traces. + * @default true + */ + traceContent?: boolean; +} diff --git a/packages/traceloop-sdk/package.json b/packages/traceloop-sdk/package.json index 40ab744b..8123fd43 100644 --- a/packages/traceloop-sdk/package.json +++ b/packages/traceloop-sdk/package.json @@ -34,7 +34,6 @@ "dependencies": { "@opentelemetry/exporter-trace-otlp-proto": "^0.44.0", "@opentelemetry/sdk-node": "^0.44.0", - "@scarf/scarf": "^1.3.0", "@traceloop/instrumentation-openai": "^0.0.22", "@types/nunjucks": "^3.2.5", "fetch-retry": "^5.0.6", 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 751b9756..763b906a 100644 --- a/packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts +++ b/packages/traceloop-sdk/src/lib/interfaces/initialize-options.interface.ts @@ -35,6 +35,12 @@ export interface InitializeOptions { */ suppressLogs?: boolean; + /** + * Whether to log prompts, completions and embeddings on traces. Optional. + * Defaults to true. + */ + traceContent?: boolean; + /** * The OpenTelemetry SpanExporter to be used for sending traces data. Optional. * Defaults to the OTLP exporter. diff --git a/packages/traceloop-sdk/src/lib/tracing/index.ts b/packages/traceloop-sdk/src/lib/tracing/index.ts index 69c0dfd4..d1d04ea8 100644 --- a/packages/traceloop-sdk/src/lib/tracing/index.ts +++ b/packages/traceloop-sdk/src/lib/tracing/index.ts @@ -32,6 +32,13 @@ export const initInstrumentations = () => { * @throws {InitializationError} if the configuration is invalid or if failed to fetch feature data. */ export const startTracing = (options: InitializeOptions) => { + if ( + options.traceContent === false || + (process.env.TRACELOOP_TRACE_CONTENT || "true").toLowerCase() === "false" + ) { + openAIInstrumentation.setConfig({ traceContent: false }); + } + const traceExporter = options.exporter ?? new OTLPTraceExporter({