@traceai/cohere 0.1.0 → 0.1.1

package/dist/src/instrumentation.d.ts

@@ -1,34 +1,25 @@
 import { InstrumentationBase, InstrumentationConfig, InstrumentationNodeModuleDefinition } from "@opentelemetry/instrumentation";
 import { TraceConfigOptions } from "@traceai/fi-core";
-/**
- * Function to check if instrumentation is enabled / disabled
- */
+import { Data } from "./attributes";
+/** True when at least one SDK method is instrumented, across module copies. */
 export declare function isPatched(): boolean;
-/**
- * @param instrumentationConfig The config for the instrumentation
- * @param traceConfig The FI trace configuration for masking sensitive information
- */
+/** Cohere v1/v2 instrumentation. Install before constructing SDK clients. */
 export declare class CohereInstrumentation extends InstrumentationBase {
-    private fiTracer;
     private _traceConfig?;
+    private patches?;
+    private manualModules?;
     readonly instrumentationName = "@traceai/cohere";
-    constructor({ instrumentationConfig, traceConfig, }?: {
+    constructor({ instrumentationConfig, traceConfig }?: {
         instrumentationConfig?: InstrumentationConfig;
         traceConfig?: TraceConfigOptions;
     });
-    enable(): void;
     protected init(): InstrumentationNodeModuleDefinition;
-    /**
-     * Manually instruments the Cohere module
-     */
-    manuallyInstrument(module: Record<string, unknown>): void;
-    /**
-     * Patches the Cohere module
-     */
+    enable(): void;
+    disable(): void;
+    manuallyInstrument(module: Data): void;
     private patch;
-    /**
-     * Un-patches the Cohere module
-     */
     private unpatch;
+    private restorePatches;
+    private wrapMethod;
 }
 //# sourceMappingURL=instrumentation.d.ts.map

package/dist/src/instrumentation.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"instrumentation.d.ts","sourceRoot":"","sources":["../../src/instrumentation.ts"],"names":[],"mappings":"AAAA,OAAO,EACL,mBAAmB,EACnB,qBAAqB,EACrB,mCAAmC,EAEpC,MAAM,gCAAgC,CAAC;AAmBxC,OAAO,EAGL,kBAAkB,EACnB,MAAM,kBAAkB,CAAC;AAU1B;;GAEG;AACH,wBAAgB,SAAS,YAExB;AAyGD;;;GAGG;AACH,qBAAa,qBAAsB,SAAQ,mBAAmB;IAC5D,OAAO,CAAC,QAAQ,CAAY;IAC5B,OAAO,CAAC,YAAY,CAAC,CAAqB;IAC1C,SAAgB,mBAAmB,qBAAqB;gBAE5C,EACV,qBAAqB,EACrB,WAAW,GACZ,GAAE;QACD,qBAAqB,CAAC,EAAE,qBAAqB,CAAC;QAC9C,WAAW,CAAC,EAAE,kBAAkB,CAAC;KAC7B;IASU,MAAM,IAAI,IAAI;IAK9B,SAAS,CAAC,IAAI;IAUd;;OAEG;IACH,kBAAkB,CAAC,MAAM,EAAE,MAAM,CAAC,MAAM,EAAE,OAAO,CAAC;IAIlD;;OAEG;IACH,OAAO,CAAC,KAAK;IAoPb;;OAEG;IACH,OAAO,CAAC,OAAO;CA2BhB"}
+{"version":3,"file":"instrumentation.d.ts","sourceRoot":"","sources":["../../src/instrumentation.ts"],"names":[],"mappings":"AAAA,OAAO,EAAE,mBAAmB,EAAE,qBAAqB,EAAE,mCAAmC,EAAE,MAAM,gCAAgC,CAAC;AAIjI,OAAO,EAAiC,kBAAkB,EAAE,MAAM,kBAAkB,CAAC;AAErF,OAAO,EAAE,IAAI,EAAoE,MAAM,cAAc,CAAC;AAUtG,+EAA+E;AAC/E,wBAAgB,SAAS,IAAI,OAAO,CAA2B;AAE/D,6EAA6E;AAC7E,qBAAa,qBAAsB,SAAQ,mBAAmB;IAC5D,OAAO,CAAC,YAAY,CAAC,CAAqB;IAE1C,OAAO,CAAC,OAAO,CAAC,CAAU;IAC1B,OAAO,CAAC,aAAa,CAAC,CAAY;IAClC,SAAgB,mBAAmB,qBAAqB;gBAE5C,EAAE,qBAAqB,EAAE,WAAW,EAAE,GAAE;QAClD,qBAAqB,CAAC,EAAE,qBAAqB,CAAC;QAC9C,WAAW,CAAC,EAAE,kBAAkB,CAAC;KAC7B;IAKN,SAAS,CAAC,IAAI;IAIE,MAAM,IAAI,IAAI;IAKd,OAAO,IAAI,IAAI;IAK/B,kBAAkB,CAAC,MAAM,EAAE,IAAI,GAAG,IAAI;IAKtC,OAAO,CAAC,KAAK;IA+Bb,OAAO,CAAC,OAAO;IAEf,OAAO,CAAC,cAAc;IAYtB,OAAO,CAAC,UAAU;CA8EnB"}

package/dist/src/instrumentation.js

@@ -1,364 +1,211 @@
 "use strict";
-var __rest = (this && this.__rest) || function (s, e) {
-    var t = {};
-    for (var p in s) if (Object.prototype.hasOwnProperty.call(s, p) && e.indexOf(p) < 0)
-        t[p] = s[p];
-    if (s != null && typeof Object.getOwnPropertySymbols === "function")
-        for (var i = 0, p = Object.getOwnPropertySymbols(s); i < p.length; i++) {
-            if (e.indexOf(p[i]) < 0 && Object.prototype.propertyIsEnumerable.call(s, p[i]))
-                t[p[i]] = s[p[i]];
-        }
-    return t;
-};
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.CohereInstrumentation = void 0;
 exports.isPatched = isPatched;
 const instrumentation_1 = require("@opentelemetry/instrumentation");
 const api_1 = require("@opentelemetry/api");
-const version_1 = require("./version");
-const fi_semantic_conventions_1 = require("@traceai/fi-semantic-conventions");
 const core_1 = require("@opentelemetry/core");
+const node_perf_hooks_1 = require("node:perf_hooks");
 const fi_core_1 = require("@traceai/fi-core");
-const MODULE_NAME = "cohere-ai";
-/**
- * Flag to check if the Cohere module has been patched
- */
-let _isFIPatched = false;
-/**
- * Function to check if instrumentation is enabled / disabled
- */
-function isPatched() {
-    return _isFIPatched;
-}
-/**
- * Resolves the execution context for the current span
- */
-function getExecContext(span) {
-    const activeContext = api_1.context.active();
-    const suppressTracing = (0, core_1.isTracingSuppressed)(activeContext);
-    const execContext = suppressTracing
-        ? api_1.trace.setSpan(api_1.context.active(), span)
-        : activeContext;
-    if (suppressTracing) {
-        api_1.trace.deleteSpan(activeContext);
-    }
-    return execContext;
-}
-/**
- * @param instrumentationConfig The config for the instrumentation
- * @param traceConfig The FI trace configuration for masking sensitive information
- */
+const version_1 = require("./version");
+const attributes_1 = require("./attributes");
+const stream_1 = require("./stream");
+const OPERATIONS = ["chat", "chatStream", "embed", "rerank"];
+const IN_COHERE_CALL = (0, api_1.createContextKey)("traceai.cohere.call");
+const owners = new WeakMap();
+let patchCount = 0;
+/** True when at least one SDK method is instrumented, across module copies. */
+function isPatched() { return patchCount > 0; }
+/** Cohere v1/v2 instrumentation. Install before constructing SDK clients. */
 class CohereInstrumentation extends instrumentation_1.InstrumentationBase {
-    constructor({ instrumentationConfig, traceConfig, } = {}) {
+    constructor({ instrumentationConfig, traceConfig } = {}) {
         super("@traceai/cohere", version_1.VERSION, Object.assign({}, instrumentationConfig));
         this.instrumentationName = "@traceai/cohere";
         this._traceConfig = traceConfig;
     }
+    init() {
+        return new instrumentation_1.InstrumentationNodeModuleDefinition("cohere-ai", [">=7.0.0 <9"], this.patch.bind(this), this.unpatch.bind(this));
+    }
     enable() {
+        var _a;
         super.enable();
-        this.fiTracer = new fi_core_1.FITracer({ tracer: this.tracer, traceConfig: this._traceConfig });
+        for (const module of (_a = this.manualModules) !== null && _a !== void 0 ? _a : [])
+            this.patch(module);
     }
-    init() {
-        const module = new instrumentation_1.InstrumentationNodeModuleDefinition("cohere-ai", [">=7.0.0"], this.patch.bind(this), this.unpatch.bind(this));
-        return module;
+    disable() {
+        super.disable();
+        this.restorePatches();
     }
-    /**
-     * Manually instruments the Cohere module
-     */
     manuallyInstrument(module) {
-        this.patch(module);
+        var _a;
+        ((_a = this.manualModules) !== null && _a !== void 0 ? _a : (this.manualModules = new Set())).add(module);
+        if (this.isEnabled())
+            this.patch(module);
     }
-    /**
-     * Patches the Cohere module
-     */
-    patch(module, moduleVersion) {
+    patch(module) {
         var _a, _b;
-        if ((module === null || module === void 0 ? void 0 : module.fiPatched) || _isFIPatched) {
-            return module;
-        }
-        const instrumentation = this;
-        // Ensure fiTracer uses the most up-to-date tracer
-        if (!((_a = instrumentation.fiTracer) === null || _a === void 0 ? void 0 : _a.tracer) ||
-            ((_b = instrumentation.fiTracer) === null || _b === void 0 ? void 0 : _b.tracer) !== instrumentation.tracer) {
-            instrumentation.fiTracer = new fi_core_1.FITracer({
-                tracer: instrumentation.tracer,
-                traceConfig: instrumentation._traceConfig
-            });
-        }
-        // Patch CohereClient class
-        const CohereClient = module.CohereClient;
-        if (CohereClient === null || CohereClient === void 0 ? void 0 : CohereClient.prototype) {
-            // Patch chat method
-            if (typeof CohereClient.prototype.chat === 'function') {
-                this._wrap(CohereClient.prototype, "chat", 
-                // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                (original) => {
-                    return function patchedChat(...args) {
-                        var _a;
-                        const request = args[0];
-                        const modelName = request.model || 'command';
-                        const { message, chat_history, preamble } = request, invocationParameters = __rest(request, ["message", "chat_history", "preamble"]);
-                        const span = instrumentation.fiTracer.startSpan(`Cohere Chat`, {
-                            kind: api_1.SpanKind.INTERNAL,
-                            attributes: Object.assign(Object.assign(Object.assign({ [fi_semantic_conventions_1.SemanticConventions.FI_SPAN_KIND]: fi_semantic_conventions_1.FISpanKind.LLM, [fi_semantic_conventions_1.SemanticConventions.LLM_MODEL_NAME]: modelName, [fi_semantic_conventions_1.SemanticConventions.INPUT_VALUE]: request.message, [fi_semantic_conventions_1.SemanticConventions.INPUT_MIME_TYPE]: fi_semantic_conventions_1.MimeType.TEXT, [fi_semantic_conventions_1.SemanticConventions.LLM_INVOCATION_PARAMETERS]: JSON.stringify(invocationParameters), [fi_semantic_conventions_1.SemanticConventions.LLM_PROVIDER]: fi_semantic_conventions_1.LLMProvider.COHERE, [fi_semantic_conventions_1.SemanticConventions.GEN_AI_OPERATION_NAME]: "chat" }, getChatInputMessagesAttributes(request)), getChatToolsAttributes(request)), { [fi_semantic_conventions_1.SemanticConventions.RAW_INPUT]: (_a = (0, fi_core_1.safelyJSONStringify)(request)) !== null && _a !== void 0 ? _a : "" }),
-                        });
-                        const execContext = getExecContext(span);
-                        const execPromise = (0, instrumentation_1.safeExecuteInTheMiddle)(() => {
-                            return api_1.context.with(api_1.trace.setSpan(execContext, span), () => {
-                                return original.apply(this, args);
-                            });
-                        }, (error) => {
-                            if (error) {
-                                span.recordException(error);
-                                span.setStatus({ code: api_1.SpanStatusCode.ERROR, message: error.message });
-                                span.end();
-                            }
-                        });
-                        const wrappedPromise = execPromise.then((result) => {
-                            var _a;
-                            span.setAttributes(Object.assign(Object.assign(Object.assign({ [fi_semantic_conventions_1.SemanticConventions.OUTPUT_VALUE]: result.text, [fi_semantic_conventions_1.SemanticConventions.OUTPUT_MIME_TYPE]: fi_semantic_conventions_1.MimeType.TEXT }, getChatOutputMessagesAttributes(result)), getChatUsageAttributes(result)), { [fi_semantic_conventions_1.SemanticConventions.RAW_OUTPUT]: (_a = (0, fi_core_1.safelyJSONStringify)(result)) !== null && _a !== void 0 ? _a : "" }));
-                            span.setStatus({ code: api_1.SpanStatusCode.OK });
-                            span.end();
-                            return result;
-                        }).catch((error) => {
-                            span.recordException(error);
-                            span.setStatus({ code: api_1.SpanStatusCode.ERROR, message: error.message });
-                            span.end();
-                            throw error;
-                        });
-                        return api_1.context.bind(execContext, wrappedPromise);
-                    };
-                });
-            }
-            // Patch embed method
-            if (typeof CohereClient.prototype.embed === 'function') {
-                this._wrap(CohereClient.prototype, "embed", 
-                // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                (original) => {
-                    return function patchedEmbed(...args) {
-                        var _a;
-                        const request = args[0];
-                        const modelName = request.model || 'embed-english-v3.0';
-                        const span = instrumentation.fiTracer.startSpan(`Cohere Embed`, {
-                            kind: api_1.SpanKind.INTERNAL,
-                            attributes: Object.assign(Object.assign({ [fi_semantic_conventions_1.SemanticConventions.FI_SPAN_KIND]: fi_semantic_conventions_1.FISpanKind.EMBEDDING, [fi_semantic_conventions_1.SemanticConventions.EMBEDDING_MODEL_NAME]: modelName, [fi_semantic_conventions_1.SemanticConventions.INPUT_VALUE]: JSON.stringify(request.texts), [fi_semantic_conventions_1.SemanticConventions.INPUT_MIME_TYPE]: fi_semantic_conventions_1.MimeType.JSON, [fi_semantic_conventions_1.SemanticConventions.LLM_PROVIDER]: fi_semantic_conventions_1.LLMProvider.COHERE, [fi_semantic_conventions_1.SemanticConventions.GEN_AI_OPERATION_NAME]: "embeddings" }, getEmbedTextAttributes(request)), { [fi_semantic_conventions_1.SemanticConventions.RAW_INPUT]: (_a = (0, fi_core_1.safelyJSONStringify)(request)) !== null && _a !== void 0 ? _a : "" }),
-                        });
-                        const execContext = getExecContext(span);
-                        const execPromise = (0, instrumentation_1.safeExecuteInTheMiddle)(() => {
-                            return api_1.context.with(api_1.trace.setSpan(execContext, span), () => {
-                                return original.apply(this, args);
-                            });
-                        }, (error) => {
-                            if (error) {
-                                span.recordException(error);
-                                span.setStatus({ code: api_1.SpanStatusCode.ERROR, message: error.message });
-                                span.end();
-                            }
-                        });
-                        const wrappedPromise = execPromise.then((result) => {
-                            var _a;
-                            span.setAttributes(Object.assign(Object.assign({}, getEmbedVectorAttributes(result)), { [fi_semantic_conventions_1.SemanticConventions.RAW_OUTPUT]: (_a = (0, fi_core_1.safelyJSONStringify)(result)) !== null && _a !== void 0 ? _a : "" }));
-                            span.setStatus({ code: api_1.SpanStatusCode.OK });
-                            span.end();
-                            return result;
-                        }).catch((error) => {
-                            span.recordException(error);
-                            span.setStatus({ code: api_1.SpanStatusCode.ERROR, message: error.message });
-                            span.end();
-                            throw error;
-                        });
-                        return api_1.context.bind(execContext, wrappedPromise);
-                    };
-                });
+        const Client = module.CohereClient;
+        const targets = [Client === null || Client === void 0 ? void 0 : Client.prototype];
+        // V2Client is not exported at the module root. The public client.v2 getter
+        // exposes the same generated class used by CohereClientV2's bound methods.
+        // Construction performs no I/O; an explicit inert token avoids env lookup.
+        if (typeof Client === "function" && Client.prototype && "v2" in Client.prototype) {
+            try {
+                const v2 = new Client({ token: "traceai-instrumentation-placeholder" }).v2;
+                if (v2)
+                    targets.push(Object.getPrototypeOf(v2));
             }
-            // Patch rerank method
-            if (typeof CohereClient.prototype.rerank === 'function') {
-                this._wrap(CohereClient.prototype, "rerank", 
-                // eslint-disable-next-line @typescript-eslint/no-explicit-any
-                (original) => {
-                    return function patchedRerank(...args) {
-                        var _a;
-                        const request = args[0];
-                        const modelName = request.model || 'rerank-english-v3.0';
-                        const span = instrumentation.fiTracer.startSpan(`Cohere Rerank`, {
-                            kind: api_1.SpanKind.INTERNAL,
-                            attributes: {
-                                [fi_semantic_conventions_1.SemanticConventions.FI_SPAN_KIND]: fi_semantic_conventions_1.FISpanKind.LLM,
-                                [fi_semantic_conventions_1.SemanticConventions.LLM_MODEL_NAME]: modelName,
-                                [fi_semantic_conventions_1.SemanticConventions.INPUT_VALUE]: request.query,
-                                [fi_semantic_conventions_1.SemanticConventions.INPUT_MIME_TYPE]: fi_semantic_conventions_1.MimeType.TEXT,
-                                [fi_semantic_conventions_1.SemanticConventions.LLM_PROVIDER]: fi_semantic_conventions_1.LLMProvider.COHERE,
-                                [fi_semantic_conventions_1.SemanticConventions.GEN_AI_OPERATION_NAME]: "rerank",
-                                "cohere.rerank.query": request.query,
-                                "cohere.rerank.documents_count": request.documents.length,
-                                "cohere.rerank.top_n": request.top_n,
-                                [fi_semantic_conventions_1.SemanticConventions.RAW_INPUT]: (_a = (0, fi_core_1.safelyJSONStringify)(request)) !== null && _a !== void 0 ? _a : "",
-                            },
-                        });
-                        const execContext = getExecContext(span);
-                        const execPromise = (0, instrumentation_1.safeExecuteInTheMiddle)(() => {
-                            return api_1.context.with(api_1.trace.setSpan(execContext, span), () => {
-                                return original.apply(this, args);
-                            });
-                        }, (error) => {
-                            if (error) {
-                                span.recordException(error);
-                                span.setStatus({ code: api_1.SpanStatusCode.ERROR, message: error.message });
-                                span.end();
-                            }
-                        });
-                        const wrappedPromise = execPromise.then((result) => {
-                            var _a;
-                            span.setAttributes(Object.assign(Object.assign({ "cohere.rerank.results_count": result.results.length }, getRerankResultsAttributes(result)), { [fi_semantic_conventions_1.SemanticConventions.RAW_OUTPUT]: (_a = (0, fi_core_1.safelyJSONStringify)(result)) !== null && _a !== void 0 ? _a : "" }));
-                            span.setStatus({ code: api_1.SpanStatusCode.OK });
-                            span.end();
-                            return result;
-                        }).catch((error) => {
-                            span.recordException(error);
-                            span.setStatus({ code: api_1.SpanStatusCode.ERROR, message: error.message });
-                            span.end();
-                            throw error;
-                        });
-                        return api_1.context.bind(execContext, wrappedPromise);
-                    };
-                });
+            catch (_c) {
+                // Telemetry must not prevent module loading; v1 remains usable.
+                api_1.diag.warn("Cohere instrumentation could not discover the v2 client");
             }
         }
-        _isFIPatched = true;
-        try {
-            module.fiPatched = true;
-        }
-        catch (e) {
-            api_1.diag.warn(`Failed to set fiPatched flag on module '${MODULE_NAME}'. Error: ${e}`);
+        for (const target of targets) {
+            if (!target)
+                continue;
+            for (const method of OPERATIONS) {
+                if (typeof target[method] !== "function" || ((_a = owners.get(target)) === null || _a === void 0 ? void 0 : _a.has(method)))
+                    continue;
+                const hadOwn = Object.prototype.hasOwnProperty.call(target, method);
+                this._wrap(target, method, (original) => this.wrapMethod(method, original));
+                const wrapped = target[method];
+                ((_b = this.patches) !== null && _b !== void 0 ? _b : (this.patches = [])).push({ target, method, wrapped, hadOwn });
+                if (!owners.has(target))
+                    owners.set(target, new Map());
+                owners.get(target).set(method, this);
+                patchCount++;
+            }
         }
         return module;
     }
-    /**
-     * Un-patches the Cohere module
-     */
-    unpatch(moduleExports, moduleVersion) {
-        const CohereClient = moduleExports.CohereClient;
-        if (CohereClient === null || CohereClient === void 0 ? void 0 : CohereClient.prototype) {
-            if (CohereClient.prototype.chat) {
-                this._unwrap(CohereClient.prototype, "chat");
-            }
-            if (CohereClient.prototype.embed) {
-                this._unwrap(CohereClient.prototype, "embed");
-            }
-            if (CohereClient.prototype.rerank) {
-                this._unwrap(CohereClient.prototype, "rerank");
+    unpatch() { this.restorePatches(); }
+    restorePatches() {
+        var _a, _b;
+        for (const { target, method, wrapped, hadOwn } of (_a = this.patches) !== null && _a !== void 0 ? _a : []) {
+            if (target[method] === wrapped) {
+                this._unwrap(target, method);
+                if (!hadOwn)
+                    delete target[method];
             }
+            (_b = owners.get(target)) === null || _b === void 0 ? void 0 : _b.delete(method);
+            patchCount--;
         }
-        _isFIPatched = false;
-        try {
-            moduleExports.fiPatched = false;
-        }
-        catch (e) {
-            api_1.diag.warn(`Failed to unset fiPatched flag on module '${MODULE_NAME}'. Error: ${e}`);
-        }
+        this.patches = [];
     }
-}
-exports.CohereInstrumentation = CohereInstrumentation;
-/**
- * Converts the chat request to LLM input messages
- */
-function getChatInputMessagesAttributes(request) {
-    var _a;
-    const messages = [];
-    if (request.preamble) {
-        messages.push({ role: 'system', content: request.preamble });
-    }
-    if (request.chat_history) {
-        request.chat_history.forEach((msg) => {
-            messages.push({ role: msg.role, content: msg.message });
-        });
-    }
-    messages.push({ role: 'user', content: request.message });
-    return { [fi_semantic_conventions_1.SemanticConventions.LLM_INPUT_MESSAGES]: (_a = (0, fi_core_1.safelyJSONStringify)(messages)) !== null && _a !== void 0 ? _a : "[]" };
-}
-/**
- * Gets chat tools attributes
- */
-function getChatToolsAttributes(request) {
-    var _a;
-    if (!request.tools) {
-        return {};
-    }
-    return { [fi_semantic_conventions_1.SemanticConventions.LLM_TOOLS]: (_a = (0, fi_core_1.safelyJSONStringify)(request.tools)) !== null && _a !== void 0 ? _a : "[]" };
-}
-/**
- * Gets chat output messages attributes
- */
-function getChatOutputMessagesAttributes(response) {
-    var _a;
-    const msg = { role: 'assistant', content: response.text };
-    if (response.tool_calls) {
-        msg.tool_calls = response.tool_calls.map((tc) => ({
-            function: { name: tc.name, arguments: JSON.stringify(tc.parameters) },
-        }));
-    }
-    return { [fi_semantic_conventions_1.SemanticConventions.LLM_OUTPUT_MESSAGES]: (_a = (0, fi_core_1.safelyJSONStringify)([msg])) !== null && _a !== void 0 ? _a : "[]" };
-}
-/**
- * Gets chat usage attributes
- */
-function getChatUsageAttributes(response) {
-    const meta = response.meta;
-    if (!meta) {
-        return {};
-    }
-    const billedUnits = meta.billed_units;
-    const tokens = meta.tokens;
-    if (billedUnits) {
-        return {
-            [fi_semantic_conventions_1.SemanticConventions.LLM_TOKEN_COUNT_PROMPT]: billedUnits.input_tokens,
-            [fi_semantic_conventions_1.SemanticConventions.LLM_TOKEN_COUNT_COMPLETION]: billedUnits.output_tokens,
-            [fi_semantic_conventions_1.SemanticConventions.LLM_TOKEN_COUNT_TOTAL]: (billedUnits.input_tokens || 0) + (billedUnits.output_tokens || 0),
-        };
-    }
-    if (tokens) {
-        return {
-            [fi_semantic_conventions_1.SemanticConventions.LLM_TOKEN_COUNT_PROMPT]: tokens.input_tokens,
-            [fi_semantic_conventions_1.SemanticConventions.LLM_TOKEN_COUNT_COMPLETION]: tokens.output_tokens,
-            [fi_semantic_conventions_1.SemanticConventions.LLM_TOKEN_COUNT_TOTAL]: (tokens.input_tokens || 0) + (tokens.output_tokens || 0),
+    wrapMethod(operation, original) {
+        const instrumentation = this;
+        return function (...args) {
+            const parent = api_1.context.active();
+            // Bound v2 methods can outlive unpatch. Respect disable even for those.
+            if (!instrumentation.isEnabled() || (0, core_1.isTracingSuppressed)(parent) || parent.getValue(IN_COHERE_CALL)) {
+                return original.apply(this, args);
+            }
+            const config = (0, fi_core_1.generateTraceConfig)(instrumentation._traceConfig);
+            let span;
+            try {
+                const tracer = new fi_core_1.FITracer({ tracer: instrumentation.tracer, traceConfig: config });
+                const name = operation === "chatStream" ? "Chat Stream" : operation[0].toUpperCase() + operation.slice(1);
+                span = tracer.startSpan(`Cohere ${name}`, { kind: api_1.SpanKind.INTERNAL, attributes: (0, attributes_1.inputAttributes)(operation, (0, attributes_1.object)(args[0]), config) }, parent);
+            }
+            catch (_a) {
+                api_1.diag.warn("Cohere instrumentation could not start span");
+                return original.apply(this, args);
+            }
+            const execContext = api_1.trace.setSpan(parent, span).setValue(IN_COHERE_CALL, true);
+            let ended = false;
+            const finish = (response, error, cancelled = false, timeToFirstTokenMs) => {
+                var _a;
+                if (ended)
+                    return;
+                ended = true;
+                try {
+                    // Resolved HTTP-200 chat responses can still signal failure or user
+                    // cancellation. Use the same classification for both chat paths,
+                    // independent of content capture and without rejecting SDK results.
+                    if ((operation === "chat" || operation === "chatStream") && error === undefined && !cancelled) {
+                        const reason = (_a = response.finishReason) !== null && _a !== void 0 ? _a : response.finish_reason;
+                        switch (reason) {
+                            case "ERROR":
+                            case "ERROR_TOXIC":
+                            case "ERROR_LIMIT":
+                            case "TIMEOUT":
+                                error = new Error(`Cohere generation failed (${reason})`);
+                                break;
+                            case "USER_CANCEL":
+                                cancelled = true;
+                                break;
+                        }
+                    }
+                    span.setAttributes((0, attributes_1.outputAttributes)(operation, response, config));
+                    if (timeToFirstTokenMs !== undefined)
+                        span.setAttribute("cohere.stream.time_to_first_token_ms", timeToFirstTokenMs);
+                    if (error !== undefined) {
+                        span.recordException(error instanceof Error ? error : new Error(String(error)));
+                        span.setStatus({ code: api_1.SpanStatusCode.ERROR, message: error instanceof Error ? error.message : String(error) });
+                    }
+                    else if (cancelled) {
+                        span.setAttribute(operation === "chatStream" ? "cohere.stream.cancelled" : "cohere.request.cancelled", true);
+                    }
+                    else {
+                        span.setStatus({ code: api_1.SpanStatusCode.OK });
+                    }
+                }
+                catch (_b) {
+                    api_1.diag.warn("Cohere instrumentation could not record response");
+                }
+                finally {
+                    span.end();
+                }
+            };
+            const streams = new WeakMap();
+            const resolved = (result) => {
+                if (operation === "chatStream" && result && typeof result[Symbol.asyncIterator] === "function") {
+                    if (!streams.has(result))
+                        streams.set(result, (0, stream_1.wrapStream)(result, execContext, !(0, attributes_1.hideOutput)(config), finish, requestStart));
+                    return streams.get(result);
+                }
+                finish((0, attributes_1.object)(result));
+                return result;
+            };
+            const rejected = (error) => {
+                finish({}, error !== null && error !== void 0 ? error : new Error("Cohere request failed"));
+                throw error;
+            };
+            let pending;
+            const requestStart = node_perf_hooks_1.performance.now();
+            try {
+                pending = api_1.context.with(execContext, () => original.apply(this, args));
+            }
+            catch (error) {
+                return rejected(error);
+            }
+            return observePromise(pending, resolved, rejected);
         };
     }
-    return {};
-}
-/**
- * Gets embed text attributes
- */
-function getEmbedTextAttributes(request) {
-    return request.texts.reduce((acc, text, index) => {
-        acc[`${fi_semantic_conventions_1.SemanticConventions.EMBEDDING_EMBEDDINGS}.${index}.${fi_semantic_conventions_1.SemanticConventions.EMBEDDING_TEXT}`] = text;
-        return acc;
-    }, {});
-}
-/**
- * Gets embed vector attributes
- */
-function getEmbedVectorAttributes(response) {
-    var _a;
-    const embeddings = Array.isArray(response.embeddings)
-        ? response.embeddings
-        : ((_a = response.embeddings) === null || _a === void 0 ? void 0 : _a.float) || [];
-    return embeddings.reduce((acc, embedding, index) => {
-        acc[`${fi_semantic_conventions_1.SemanticConventions.EMBEDDING_EMBEDDINGS}.${index}.${fi_semantic_conventions_1.SemanticConventions.EMBEDDING_VECTOR}`] = embedding;
-        return acc;
-    }, {});
 }
-/**
- * Gets rerank results attributes
- */
-function getRerankResultsAttributes(response) {
-    const attributes = {};
-    response.results.forEach((result, index) => {
-        attributes[`cohere.rerank.results.${index}.index`] = result.index;
-        attributes[`cohere.rerank.results.${index}.relevance_score`] = result.relevance_score;
-        if (result.document) {
-            attributes[`cohere.rerank.results.${index}.document`] = result.document.text;
-        }
+exports.CohereInstrumentation = CohereInstrumentation;
+/** Keep SDK promise helpers (notably HttpResponsePromise.withRawResponse). */
+function observePromise(pending, resolved, rejected) {
+    const promise = pending;
+    const observed = promise.then(resolved, rejected);
+    if (typeof promise.withRawResponse !== "function")
+        return observed;
+    // A caller may exclusively use withRawResponse. Do not leave an unobserved
+    // rejected branch merely because telemetry also observes the normal promise.
+    void observed.catch(() => { });
+    return new Proxy(promise, {
+        get(target, key) {
+            if (key === "then")
+                return observed.then.bind(observed);
+            if (key === "catch")
+                return observed.catch.bind(observed);
+            if (key === "finally")
+                return observed.finally.bind(observed);
+            if (key === "withRawResponse")
+                return () => target.withRawResponse().then((raw) => (Object.assign(Object.assign({}, raw), { data: resolved(raw.data) })), rejected);
+            const value = Reflect.get(target, key, target);
+            return typeof value === "function" ? value.bind(target) : value;
+        },
     });
-    return attributes;
 }
 //# sourceMappingURL=instrumentation.js.map