@traceloop/instrumentation-bedrock 0.9.2 → 0.11.0

package/dist/index.js

@@ -5,7 +5,7 @@ var api = require('@opentelemetry/api');
 var instrumentation = require('@opentelemetry/instrumentation');
 var aiSemanticConventions = require('@traceloop/ai-semantic-conventions');
 
-var version = "0.9.2";
+var version = "0.11.0";
 
 class BedrockInstrumentation extends instrumentation.InstrumentationBase {
     constructor(config = {}) {

package/dist/index.mjs

@@ -0,0 +1,285 @@
+import { __awaiter, __asyncValues } from 'tslib';
+import { trace, context, SpanStatusCode, SpanKind } from '@opentelemetry/api';
+import { InstrumentationBase, InstrumentationNodeModuleDefinition, safeExecuteInTheMiddle } from '@opentelemetry/instrumentation';
+import { SpanAttributes, LLMRequestTypeValues, CONTEXT_KEY_ALLOW_TRACE_CONTENT } from '@traceloop/ai-semantic-conventions';
+
+var version = "0.11.0";
+
+class BedrockInstrumentation extends InstrumentationBase {
+    constructor(config = {}) {
+        super("@traceloop/instrumentation-bedrock", version, config);
+    }
+    setConfig(config = {}) {
+        super.setConfig(config);
+    }
+    init() {
+        const module = new InstrumentationNodeModuleDefinition("@aws-sdk/client-bedrock-runtime", [">=3.499.0"], this.wrap.bind(this), this.unwrap.bind(this));
+        return module;
+    }
+    manuallyInstrument(module) {
+        this._diag.debug(`Patching @aws-sdk/client-bedrock-runtime manually`);
+        this._wrap(module.BedrockRuntimeClient.prototype, "send", this.wrapperMethod());
+    }
+    wrap(module, moduleVersion) {
+        this._diag.debug(`Patching @aws-sdk/client-bedrock-runtime@${moduleVersion}`);
+        this._wrap(module.BedrockRuntimeClient.prototype, "send", this.wrapperMethod());
+        return module;
+    }
+    unwrap(module, moduleVersion) {
+        this._diag.debug(`Unpatching @aws-sdk/client-bedrock-runtime@${moduleVersion}`);
+        this._unwrap(module.BedrockRuntimeClient.prototype, "send");
+    }
+    wrapperMethod() {
+        // eslint-disable-next-line @typescript-eslint/no-this-alias
+        const plugin = this;
+        // eslint-disable-next-line @typescript-eslint/ban-types
+        return (original) => {
+            return function method(...args) {
+                const span = plugin._startSpan({
+                    params: args[0],
+                });
+                const execContext = trace.setSpan(context.active(), span);
+                const execPromise = safeExecuteInTheMiddle(() => {
+                    return context.with(execContext, () => {
+                        return original.apply(this, args);
+                    });
+                }, (e) => {
+                    if (e) {
+                        plugin._diag.error(`Error in bedrock instrumentation`, e);
+                    }
+                });
+                const wrappedPromise = plugin._wrapPromise(span, execPromise);
+                return context.bind(execContext, wrappedPromise);
+            };
+        };
+    }
+    _wrapPromise(span, promise) {
+        return promise
+            .then((result) => __awaiter(this, void 0, void 0, function* () {
+            yield this._endSpan({
+                span,
+                result: result,
+            });
+            return new Promise((resolve) => resolve(result));
+        }))
+            .catch((error) => {
+            return new Promise((_, reject) => {
+                span.setStatus({
+                    code: SpanStatusCode.ERROR,
+                    message: error.message,
+                });
+                span.recordException(error);
+                span.end();
+                reject(error);
+            });
+        });
+    }
+    _startSpan({ params, }) {
+        var _a, _b;
+        let attributes = {};
+        try {
+            const [vendor, model] = params.input.modelId
+                ? params.input.modelId.split(".")
+                : ["", ""];
+            attributes = {
+                [SpanAttributes.LLM_SYSTEM]: vendor,
+                [SpanAttributes.LLM_REQUEST_MODEL]: model,
+                [SpanAttributes.LLM_RESPONSE_MODEL]: model,
+                [SpanAttributes.LLM_REQUEST_TYPE]: LLMRequestTypeValues.COMPLETION,
+            };
+            if (typeof params.input.body === "string") {
+                const requestBody = JSON.parse(params.input.body);
+                attributes = Object.assign(Object.assign({}, attributes), this._setRequestAttributes(vendor, requestBody));
+            }
+        }
+        catch (e) {
+            this._diag.debug(e);
+            (_b = (_a = this._config).exceptionLogger) === null || _b === void 0 ? void 0 : _b.call(_a, e);
+        }
+        return this.tracer.startSpan(`bedrock.completion`, {
+            kind: SpanKind.CLIENT,
+            attributes,
+        });
+    }
+    _endSpan({ span, result, }) {
+        var _a, e_1, _b, _c;
+        var _d, _e, _f;
+        return __awaiter(this, void 0, void 0, function* () {
+            try {
+                if ("body" in result) {
+                    const attributes = "attributes" in span
+                        ? span["attributes"]
+                        : {};
+                    if (SpanAttributes.LLM_SYSTEM in attributes) {
+                        if (!(result.body instanceof Object.getPrototypeOf(Uint8Array))) {
+                            const rawRes = result.body;
+                            let streamedContent = "";
+                            try {
+                                for (var _g = true, rawRes_1 = __asyncValues(rawRes), rawRes_1_1; rawRes_1_1 = yield rawRes_1.next(), _a = rawRes_1_1.done, !_a; _g = true) {
+                                    _c = rawRes_1_1.value;
+                                    _g = false;
+                                    const value = _c;
+                                    // Convert it to a JSON String
+                                    const jsonString = new TextDecoder().decode((_d = value.chunk) === null || _d === void 0 ? void 0 : _d.bytes);
+                                    // Parse the JSON string
+                                    const parsedResponse = JSON.parse(jsonString);
+                                    if ("amazon-bedrock-invocationMetrics" in parsedResponse) {
+                                        span.setAttribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, parsedResponse["amazon-bedrock-invocationMetrics"]["inputTokenCount"]);
+                                        span.setAttribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, parsedResponse["amazon-bedrock-invocationMetrics"]["outputTokenCount"]);
+                                        span.setAttribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, parsedResponse["amazon-bedrock-invocationMetrics"]["inputTokenCount"] +
+                                            parsedResponse["amazon-bedrock-invocationMetrics"]["outputTokenCount"]);
+                                    }
+                                    let responseAttributes = this._setResponseAttributes(attributes[SpanAttributes.LLM_SYSTEM], parsedResponse, true);
+                                    // ! NOTE: This make sure the content always have all streamed chunks
+                                    if (this._shouldSendPrompts()) {
+                                        // Update local value with attribute value that was set by _setResponseAttributes
+                                        streamedContent +=
+                                            responseAttributes[`${SpanAttributes.LLM_COMPLETIONS}.0.content`];
+                                        // re-assign the new value to responseAttributes
+                                        responseAttributes = Object.assign(Object.assign({}, responseAttributes), { [`${SpanAttributes.LLM_COMPLETIONS}.0.content`]: streamedContent });
+                                    }
+                                    span.setAttributes(responseAttributes);
+                                }
+                            }
+                            catch (e_1_1) { e_1 = { error: e_1_1 }; }
+                            finally {
+                                try {
+                                    if (!_g && !_a && (_b = rawRes_1.return)) yield _b.call(rawRes_1);
+                                }
+                                finally { if (e_1) throw e_1.error; }
+                            }
+                        }
+                        else if (result.body instanceof Object.getPrototypeOf(Uint8Array)) {
+                            // Convert it to a JSON String
+                            const jsonString = new TextDecoder().decode(result.body);
+                            // Parse the JSON string
+                            const parsedResponse = JSON.parse(jsonString);
+                            const responseAttributes = this._setResponseAttributes(attributes[SpanAttributes.LLM_SYSTEM], parsedResponse);
+                            span.setAttributes(responseAttributes);
+                        }
+                    }
+                }
+            }
+            catch (e) {
+                this._diag.debug(e);
+                (_f = (_e = this._config).exceptionLogger) === null || _f === void 0 ? void 0 : _f.call(_e, e);
+            }
+            span.setStatus({ code: SpanStatusCode.OK });
+            span.end();
+        });
+    }
+    _setRequestAttributes(vendor, requestBody) {
+        switch (vendor) {
+            case "ai21": {
+                return Object.assign({ [SpanAttributes.LLM_REQUEST_TOP_P]: requestBody["topP"], [SpanAttributes.LLM_REQUEST_TEMPERATURE]: requestBody["temperature"], [SpanAttributes.LLM_REQUEST_MAX_TOKENS]: requestBody["maxTokens"], [SpanAttributes.LLM_PRESENCE_PENALTY]: requestBody["presencePenalty"]["scale"], [SpanAttributes.LLM_FREQUENCY_PENALTY]: requestBody["frequencyPenalty"]["scale"] }, (this._shouldSendPrompts()
+                    ? {
+                        [`${SpanAttributes.LLM_PROMPTS}.0.role`]: "user",
+                        [`${SpanAttributes.LLM_PROMPTS}.0.content`]: requestBody["prompt"],
+                    }
+                    : {}));
+            }
+            case "amazon": {
+                return Object.assign({ [SpanAttributes.LLM_REQUEST_TOP_P]: requestBody["textGenerationConfig"]["topP"], [SpanAttributes.LLM_REQUEST_TEMPERATURE]: requestBody["textGenerationConfig"]["temperature"], [SpanAttributes.LLM_REQUEST_MAX_TOKENS]: requestBody["textGenerationConfig"]["maxTokenCount"] }, (this._shouldSendPrompts()
+                    ? {
+                        [`${SpanAttributes.LLM_PROMPTS}.0.role`]: "user",
+                        [`${SpanAttributes.LLM_PROMPTS}.0.content`]: requestBody["inputText"],
+                    }
+                    : {}));
+            }
+            case "anthropic": {
+                return Object.assign({ [SpanAttributes.LLM_REQUEST_TOP_P]: requestBody["top_p"], [SpanAttributes.LLM_TOP_K]: requestBody["top_k"], [SpanAttributes.LLM_REQUEST_TEMPERATURE]: requestBody["temperature"], [SpanAttributes.LLM_REQUEST_MAX_TOKENS]: requestBody["max_tokens_to_sample"] }, (this._shouldSendPrompts()
+                    ? {
+                        [`${SpanAttributes.LLM_PROMPTS}.0.role`]: "user",
+                        [`${SpanAttributes.LLM_PROMPTS}.0.content`]: requestBody["prompt"]
+                            // The format is removing when we are setting span attribute
+                            .replace("\n\nHuman:", "")
+                            .replace("\n\nAssistant:", ""),
+                    }
+                    : {}));
+            }
+            case "cohere": {
+                return Object.assign({ [SpanAttributes.LLM_REQUEST_TOP_P]: requestBody["p"], [SpanAttributes.LLM_TOP_K]: requestBody["k"], [SpanAttributes.LLM_REQUEST_TEMPERATURE]: requestBody["temperature"], [SpanAttributes.LLM_REQUEST_MAX_TOKENS]: requestBody["max_tokens"] }, (this._shouldSendPrompts()
+                    ? {
+                        [`${SpanAttributes.LLM_PROMPTS}.0.role`]: "user",
+                        [`${SpanAttributes.LLM_PROMPTS}.0.content`]: requestBody["prompt"],
+                    }
+                    : {}));
+            }
+            case "meta": {
+                return Object.assign({ [SpanAttributes.LLM_REQUEST_TOP_P]: requestBody["top_p"], [SpanAttributes.LLM_REQUEST_TEMPERATURE]: requestBody["temperature"], [SpanAttributes.LLM_REQUEST_MAX_TOKENS]: requestBody["max_gen_len"] }, (this._shouldSendPrompts()
+                    ? {
+                        [`${SpanAttributes.LLM_PROMPTS}.0.role`]: "user",
+                        [`${SpanAttributes.LLM_PROMPTS}.0.content`]: requestBody["prompt"],
+                    }
+                    : {}));
+            }
+            default:
+                return {};
+        }
+    }
+    _setResponseAttributes(vendor, response, isStream = false) {
+        switch (vendor) {
+            case "ai21": {
+                return Object.assign({ [`${SpanAttributes.LLM_COMPLETIONS}.0.finish_reason`]: response["completions"][0]["finishReason"]["reason"], [`${SpanAttributes.LLM_COMPLETIONS}.0.role`]: "assistant" }, (this._shouldSendPrompts()
+                    ? {
+                        [`${SpanAttributes.LLM_COMPLETIONS}.0.content`]: response["completions"][0]["data"]["text"],
+                    }
+                    : {}));
+            }
+            case "amazon": {
+                return Object.assign({ [`${SpanAttributes.LLM_COMPLETIONS}.0.finish_reason`]: isStream
+                        ? response["completionReason"]
+                        : response["results"][0]["completionReason"], [`${SpanAttributes.LLM_COMPLETIONS}.0.role`]: "assistant", [SpanAttributes.LLM_USAGE_PROMPT_TOKENS]: response["inputTextTokenCount"], [SpanAttributes.LLM_USAGE_COMPLETION_TOKENS]: isStream
+                        ? response["totalOutputTextTokenCount"]
+                        : response["results"][0]["tokenCount"], [SpanAttributes.LLM_USAGE_TOTAL_TOKENS]: isStream
+                        ? response["inputTextTokenCount"] +
+                            response["totalOutputTextTokenCount"]
+                        : response["inputTextTokenCount"] +
+                            response["results"][0]["tokenCount"] }, (this._shouldSendPrompts()
+                    ? {
+                        [`${SpanAttributes.LLM_COMPLETIONS}.0.content`]: isStream
+                            ? response["outputText"]
+                            : response["results"][0]["outputText"],
+                    }
+                    : {}));
+            }
+            case "anthropic": {
+                return Object.assign({ [`${SpanAttributes.LLM_COMPLETIONS}.0.finish_reason`]: response["stop_reason"], [`${SpanAttributes.LLM_COMPLETIONS}.0.role`]: "assistant" }, (this._shouldSendPrompts()
+                    ? {
+                        [`${SpanAttributes.LLM_COMPLETIONS}.0.content`]: response["completion"],
+                    }
+                    : {}));
+            }
+            case "cohere": {
+                return Object.assign({ [`${SpanAttributes.LLM_COMPLETIONS}.0.finish_reason`]: response["generations"][0]["finish_reason"], [`${SpanAttributes.LLM_COMPLETIONS}.0.role`]: "assistant" }, (this._shouldSendPrompts()
+                    ? {
+                        [`${SpanAttributes.LLM_COMPLETIONS}.0.content`]: response["generations"][0]["text"],
+                    }
+                    : {}));
+            }
+            case "meta": {
+                return Object.assign({ [`${SpanAttributes.LLM_COMPLETIONS}.0.finish_reason`]: response["stop_reason"], [`${SpanAttributes.LLM_COMPLETIONS}.0.role`]: "assistant", [SpanAttributes.LLM_USAGE_PROMPT_TOKENS]: response["prompt_token_count"], [SpanAttributes.LLM_USAGE_COMPLETION_TOKENS]: response["generation_token_count"], [SpanAttributes.LLM_USAGE_TOTAL_TOKENS]: response["prompt_token_count"] + response["generation_token_count"] }, (this._shouldSendPrompts()
+                    ? {
+                        [`${SpanAttributes.LLM_COMPLETIONS}.0.content`]: response["generation"],
+                    }
+                    : {}));
+            }
+            default:
+                return {};
+        }
+    }
+    _shouldSendPrompts() {
+        const contextShouldSendPrompts = context
+            .active()
+            .getValue(CONTEXT_KEY_ALLOW_TRACE_CONTENT);
+        if (contextShouldSendPrompts !== undefined) {
+            return contextShouldSendPrompts;
+        }
+        return this._config.traceContent !== undefined
+            ? this._config.traceContent
+            : true;
+    }
+}
+
+export { BedrockInstrumentation };
+//# sourceMappingURL=index.mjs.map

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@traceloop/instrumentation-bedrock",
-  "version": "0.9.2",
+  "version": "0.11.0",
   "description": "Amazon Bedrock Instrumentation",
   "main": "dist/index.js",
   "module": "dist/index.mjs",
@@ -26,6 +26,7 @@
   },
   "files": [
     "dist/**/*.js",
+    "dist/**/*.mjs",
     "dist/**/*.js.map",
     "dist/**/*.d.ts",
     "doc",
@@ -40,7 +41,7 @@
     "@opentelemetry/core": "^1.25.0",
     "@opentelemetry/instrumentation": "^0.52.0",
     "@opentelemetry/semantic-conventions": "^1.25.0",
-    "@traceloop/ai-semantic-conventions": "^0.9.0",
+    "@traceloop/ai-semantic-conventions": "^0.11.0",
     "tslib": "^2.3.0"
   },
   "devDependencies": {
@@ -50,5 +51,5 @@
     "@pollyjs/persister-fs": "^6.0.6"
   },
   "homepage": "https://github.com/traceloop/openllmetry-js/tree/main/packages/instrumentation-openai",
-  "gitHead": "5019c20e284836be10641e7d2c5692ef8621475d"
+  "gitHead": "77fe2d5853f90f5a9c1f1fedece8e3e8fc702adf"
 }