@paid-ai/paid-node 0.5.0-alpha2 → 0.5.0

package/dist/cjs/tracing/signal.js

@@ -1,24 +1,41 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports._signal = _signal;
+exports.signal = signal;
 const api_1 = require("@opentelemetry/api");
 const tracing_js_1 = require("./tracing.js");
-const tracing_js_2 = require("./tracing.js");
-function _signal(eventName, enableCostTracing, data) {
-    if (!eventName) {
-        throw new Error("Event name is required for signal.");
+const tracingContext_js_1 = require("./tracingContext.js");
+/**
+ * Sends Paid signal. Needs to be called as part of callback to trace().
+ * When enableCostTracing flag is on, signal is associated
+ * with cost traces from the same Paid.trace() context.
+ *
+ * @param eventName - The name of the signal.
+ * @param enableCostTracing - Whether to associate this signal with cost traces
+ * from the current Paid.trace() context (default: false)
+ * @param data - Optional additional data to include with the signal
+ *
+ * @remarks
+ * When enableCostTracing is on, the signal will be associated with cost
+ * traces within the same Paid.trace() context.
+ * It is advised to only make one call to this function
+ * with enableCostTracing per Paid.trace() context.
+ * Otherwise, there will be multiple signals that refer to the same costs.
+ */
+function signal(eventName, enableCostTracing = false, data) {
+    const paidTracer = (0, tracing_js_1.getPaidTracer)();
+    const token = (0, tracing_js_1.getToken)();
+    const { externalCustomerId, externalProductId } = (0, tracingContext_js_1.getTracingContext)();
+    if (!token || !paidTracer) {
+        throw new Error(`Tracing is not initialized. Make sure you called 'initializeTracing()'`);
     }
-    const externalCustomerId = (0, tracing_js_1.getCustomerIdStorage)();
-    const externalAgentId = (0, tracing_js_1.getAgentIdStorage)();
-    const token = (0, tracing_js_1.getTokenStorage)();
-    if (!externalCustomerId || !externalAgentId || !token) {
-        throw new Error(`Missing some of: external_customer_id: ${externalCustomerId}, external_agent_id: ${externalAgentId}, or token. Make sure to call signal() within trace()`);
+    if (!externalCustomerId || !externalProductId) {
+        throw new Error(`Missing some of: external_customer_id: ${externalCustomerId}, external_product_id: ${externalProductId}, or token. Make sure to call signal() within trace()`);
     }
-    tracing_js_2.paidTracer.startActiveSpan("trace.signal", (span) => {
+    paidTracer.startActiveSpan("trace.signal", (span) => {
         try {
             const attributes = {
                 external_customer_id: externalCustomerId,
-                external_agent_id: externalAgentId,
+                external_agent_id: externalProductId,
                 event_name: eventName,
                 token: token,
             };
@@ -37,7 +54,6 @@ function _signal(eventName, enableCostTracing, data) {
                 attributes["data"] = JSON.stringify(data);
             }
             span.setAttributes(attributes);
-            // Mark span as successful
             span.setStatus({ code: api_1.SpanStatusCode.OK });
         }
         catch (error) {

package/dist/cjs/tracing/spanProcessor.d.ts

@@ -0,0 +1,10 @@
+import { Context } from "@opentelemetry/api";
+import { SpanProcessor, Span } from "@opentelemetry/sdk-trace-base";
+export declare class PaidSpanProcessor implements SpanProcessor {
+    private static readonly SPAN_NAME_PREFIX;
+    private static readonly PROMPT_ATTRIBUTES_SUBSTRINGS;
+    onStart(span: Span, _parentContext?: Context): void;
+    onEnd(): void;
+    shutdown(): Promise<void>;
+    forceFlush(): Promise<void>;
+}

package/dist/cjs/tracing/spanProcessor.js

@@ -0,0 +1,82 @@
+"use strict";
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.PaidSpanProcessor = void 0;
+const tracingContext_js_1 = require("./tracingContext.js");
+const tracing_js_1 = require("./tracing.js");
+class PaidSpanProcessor {
+    onStart(span, _parentContext) {
+        const { name } = span;
+        const { storePrompt, externalCustomerId, externalProductId: externalAgentId } = (0, tracingContext_js_1.getTracingContext)();
+        if (!storePrompt) {
+            const originalSetAttribute = span.setAttribute;
+            span.setAttribute = function (key, value) {
+                const isPromptRelated = PaidSpanProcessor.PROMPT_ATTRIBUTES_SUBSTRINGS.some((substr) => key.includes(substr));
+                if (isPromptRelated)
+                    return this;
+                return originalSetAttribute.call(this, key, value);
+            };
+            const originalSetAttributes = span.setAttributes;
+            span.setAttributes = function (attributes) {
+                const newAttributes = Object.entries(attributes).reduce((acc, [key, value]) => {
+                    const isPromptRelated = PaidSpanProcessor.PROMPT_ATTRIBUTES_SUBSTRINGS.some((substr) => key.includes(substr));
+                    if (isPromptRelated)
+                        return acc;
+                    return Object.assign(Object.assign({}, acc), { [key]: value });
+                }, {});
+                return originalSetAttributes.call(this, newAttributes);
+            };
+        }
+        if (name && !name.startsWith(PaidSpanProcessor.SPAN_NAME_PREFIX)) {
+            span.updateName(`${PaidSpanProcessor.SPAN_NAME_PREFIX}${name}`);
+        }
+        if (externalCustomerId) {
+            span.setAttribute("external_customer_id", externalCustomerId);
+        }
+        if (externalAgentId) {
+            span.setAttribute("external_agent_id", externalAgentId);
+        }
+        const token = (0, tracing_js_1.getToken)();
+        if (token) {
+            span.setAttribute("token", token);
+        }
+    }
+    onEnd() {
+        // mutating this span object doesn't do anything
+        // so we achieve filtering within the onStart method by monkey patching
+        // setAttribute method
+        return;
+    }
+    shutdown() {
+        return __awaiter(this, void 0, void 0, function* () {
+            return;
+        });
+    }
+    forceFlush() {
+        return __awaiter(this, void 0, void 0, function* () {
+            return;
+        });
+    }
+}
+exports.PaidSpanProcessor = PaidSpanProcessor;
+PaidSpanProcessor.SPAN_NAME_PREFIX = "paid.trace.";
+PaidSpanProcessor.PROMPT_ATTRIBUTES_SUBSTRINGS = [
+    "gen_ai.completion",
+    "gen_ai.request.messages",
+    "gen_ai.response.messages",
+    "llm.output_message",
+    "llm.input_message",
+    "llm.invocation_parameters",
+    "gen_ai.prompt",
+    "langchain.prompt",
+    "output.value",
+    "input.value",
+];

package/dist/cjs/tracing/tracing.d.ts

@@ -1,8 +1,14 @@
+import { Tracer } from "@opentelemetry/api";
+import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node";
 import winston from "winston";
 export declare const logger: winston.Logger;
-export declare let paidTracer: import("@opentelemetry/api").Tracer;
-export declare const getCustomerIdStorage: () => string | null | undefined;
-export declare const getAgentIdStorage: () => string | null | undefined;
-export declare const getTokenStorage: () => string | null | undefined;
-export declare function _initializeTracing(apiKey: string, collectorEndpoint?: string): void;
-export declare function _trace<T extends (...args: any[]) => any>(externalCustomerId: string, fn: T, externalAgentId: string | undefined, ...args: Parameters<T>): Promise<ReturnType<T>>;
+export declare function getToken(): string | undefined;
+export declare function getPaidTracerProvider(): NodeTracerProvider | undefined;
+export declare function getPaidTracer(): Tracer | undefined;
+export declare function initializeTracing(apiKey?: string, collectorEndpoint?: string): void;
+export declare function trace<F extends (...args: any[]) => any>(options: {
+    externalCustomerId: string;
+    externalProductId?: string;
+    storePrompt?: boolean;
+    metadata?: any;
+}, fn: F, ...args: Parameters<F>): Promise<ReturnType<F>>;

package/dist/cjs/tracing/tracing.js

@@ -12,15 +12,18 @@ var __importDefault = (this && this.__importDefault) || function (mod) {
     return (mod && mod.__esModule) ? mod : { "default": mod };
 };
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.getTokenStorage = exports.getAgentIdStorage = exports.getCustomerIdStorage = exports.paidTracer = exports.logger = void 0;
-exports._initializeTracing = _initializeTracing;
-exports._trace = _trace;
-const sdk_trace_node_1 = require("@opentelemetry/sdk-trace-node");
-const exporter_trace_otlp_http_1 = require("@opentelemetry/exporter-trace-otlp-http");
-const sdk_trace_base_1 = require("@opentelemetry/sdk-trace-base");
+exports.logger = void 0;
+exports.getToken = getToken;
+exports.getPaidTracerProvider = getPaidTracerProvider;
+exports.getPaidTracer = getPaidTracer;
+exports.initializeTracing = initializeTracing;
+exports.trace = trace;
 const api_1 = require("@opentelemetry/api");
-const async_hooks_1 = require("async_hooks");
+const exporter_trace_otlp_http_1 = require("@opentelemetry/exporter-trace-otlp-http");
+const sdk_trace_node_1 = require("@opentelemetry/sdk-trace-node");
 const winston_1 = __importDefault(require("winston"));
+const tracingContext_js_1 = require("./tracingContext.js");
+const spanProcessor_js_1 = require("./spanProcessor.js");
 exports.logger = winston_1.default.createLogger({
     level: "silent", // Default to 'silent' to avoid logging unless set via environment variable
     format: winston_1.default.format.simple(),
@@ -30,32 +33,21 @@ const logLevel = process.env.PAID_LOG_LEVEL;
 if (logLevel) {
     exports.logger.level = logLevel;
 }
-const COLLECTOR_ENDPOINT = process.env.PAID_COLLECTOR_ENDPOINT || "https://collector.agentpaid.io:4318/v1/traces";
-// set up default tracing provider
-let paidExporter = new exporter_trace_otlp_http_1.OTLPTraceExporter({ url: COLLECTOR_ENDPOINT });
-let spanProcessor = new sdk_trace_base_1.BatchSpanProcessor(paidExporter);
-let paidTracerProvider = new sdk_trace_node_1.NodeTracerProvider({ spanProcessors: [spanProcessor] });
-paidTracerProvider.register();
-exports.paidTracer = paidTracerProvider.getTracer("paid.node");
-// storage for passing info to child spans
-const customerIdStorage = new async_hooks_1.AsyncLocalStorage();
-const agentIdStorage = new async_hooks_1.AsyncLocalStorage();
-const tokenStorage = new async_hooks_1.AsyncLocalStorage();
-const getCustomerIdStorage = () => customerIdStorage.getStore();
-exports.getCustomerIdStorage = getCustomerIdStorage;
-const getAgentIdStorage = () => agentIdStorage.getStore();
-exports.getAgentIdStorage = getAgentIdStorage;
-const getTokenStorage = () => tokenStorage.getStore();
-exports.getTokenStorage = getTokenStorage;
-let _token;
-const setToken = (token) => {
-    _token = token;
-};
-const getToken = () => {
-    return _token;
-};
+const DEFAULT_COLLECTOR_ENDPOINT = process.env["PAID_OTEL_COLLECTOR_ENDPOINT"] || "https://collector.agentpaid.io:4318/v1/traces";
+let paidApiToken = undefined;
+function getToken() {
+    return paidApiToken;
+}
+let paidTracerProvider = undefined;
+function getPaidTracerProvider() {
+    return paidTracerProvider;
+}
+let paidTracer = undefined;
+function getPaidTracer() {
+    return paidTracer;
+}
 let _isShuttingDown = false;
-function setupGracefulShutdown(shuttable) {
+const setupGracefulShutdown = (shuttable) => {
     ["SIGINT", "SIGTERM", "beforeExit", "uncaughtException", "unhandledRejection"].forEach((signal) => {
         process.on(signal, () => {
             if (_isShuttingDown) {
@@ -68,48 +60,63 @@ function setupGracefulShutdown(shuttable) {
                 .catch((error) => exports.logger.error(`Error shutting down Paid tracing SDK ${error}`));
         });
     });
-}
-function _initializeTracing(apiKey, collectorEndpoint) {
-    if (getToken()) {
-        throw new Error("Tracing SDK is already initialized.");
+};
+function initializeTracing(apiKey, collectorEndpoint) {
+    const paidEnabled = (process.env.PAID_ENABLED || "true") !== "false";
+    if (!paidEnabled) {
+        exports.logger.info("Paid tracing is disabled via PAID_ENABLED environment variable");
+        return;
+    }
+    const token = getToken();
+    if (!!token) {
+        exports.logger.info("Tracing is already initialized - skipping re-intialization");
+        return;
     }
     if (!apiKey) {
-        throw new Error("Cannot initialize tracing SDK - first initialize the PaidClient with an API key");
+        const envKey = process.env.PAID_API_KEY;
+        if (!envKey) {
+            exports.logger.error("API key must be provided via PAID_API_KEY environment variable");
+            return;
+        }
+        paidApiToken = envKey;
     }
-    if (collectorEndpoint) {
-        paidExporter = new exporter_trace_otlp_http_1.OTLPTraceExporter({ url: collectorEndpoint });
-        spanProcessor = new sdk_trace_base_1.BatchSpanProcessor(paidExporter);
-        paidTracerProvider = new sdk_trace_node_1.NodeTracerProvider({ spanProcessors: [spanProcessor] });
-        paidTracerProvider.register();
-        exports.paidTracer = paidTracerProvider.getTracer("paid.node");
+    else {
+        paidApiToken = apiKey;
     }
+    const url = collectorEndpoint || DEFAULT_COLLECTOR_ENDPOINT;
+    const exporter = new exporter_trace_otlp_http_1.OTLPTraceExporter({ url });
+    const spanProcessor = new sdk_trace_node_1.SimpleSpanProcessor(exporter);
+    paidTracerProvider = new sdk_trace_node_1.NodeTracerProvider({
+        spanProcessors: [spanProcessor, new spanProcessor_js_1.PaidSpanProcessor()],
+    });
+    paidTracerProvider.register();
+    paidTracer = paidTracerProvider.getTracer("paid.node");
     setupGracefulShutdown(spanProcessor);
-    setToken(apiKey);
-    const collectorAddress = collectorEndpoint || COLLECTOR_ENDPOINT;
-    exports.logger.info(`Paid tracing SDK initialized with collector endpoint: ${collectorAddress}`);
+    exports.logger.info(`Paid tracing SDK initialized with collector endpoint: ${url}`);
 }
-function _trace(externalCustomerId, fn, externalAgentId, ...args) {
+function trace(options, fn, ...args) {
     return __awaiter(this, void 0, void 0, function* () {
         const token = getToken();
-        if (!token || !externalCustomerId) {
-            throw new Error(`Paid tracing is not initialized. Make sure to call initializeTracing() first.`);
+        const tracer = getPaidTracer();
+        if (!token || !tracer) {
+            throw new Error("Paid tracing is not initialized. Make sure to call initializeTracing() first.");
         }
-        return exports.paidTracer.startActiveSpan("paid.node", (span) => __awaiter(this, void 0, void 0, function* () {
+        const { externalCustomerId, externalProductId: externalAgentId, storePrompt, metadata } = options;
+        return yield tracer.startActiveSpan("parent_span", (span) => __awaiter(this, void 0, void 0, function* () {
             span.setAttribute("external_customer_id", externalCustomerId);
             span.setAttribute("token", token);
             if (externalAgentId) {
                 span.setAttribute("external_agent_id", externalAgentId);
             }
             try {
-                const result = yield customerIdStorage.run(externalCustomerId, () => __awaiter(this, void 0, void 0, function* () {
-                    return yield agentIdStorage.run(externalAgentId !== null && externalAgentId !== void 0 ? externalAgentId : null, () => __awaiter(this, void 0, void 0, function* () {
-                        return yield tokenStorage.run(token, () => __awaiter(this, void 0, void 0, function* () {
-                            return yield fn(...args);
-                        }));
-                    }));
-                }));
+                const res = yield (0, tracingContext_js_1.runWithTracingContext)({
+                    externalCustomerId,
+                    externalProductId: externalAgentId,
+                    storePrompt,
+                    metadata,
+                }, () => __awaiter(this, void 0, void 0, function* () { return yield fn(...args); }));
                 span.setStatus({ code: api_1.SpanStatusCode.OK });
-                return result;
+                return res;
             }
             catch (error) {
                 span.setStatus({

package/dist/cjs/tracing/tracingContext.d.ts

@@ -0,0 +1,8 @@
+export interface TracingContext {
+    externalCustomerId?: string;
+    externalProductId?: string;
+    storePrompt?: boolean;
+    metadata?: Record<string, any>;
+}
+export declare const getTracingContext: () => TracingContext;
+export declare function runWithTracingContext<F extends (...args: any) => any>(ctx: TracingContext, fn: F): Promise<ReturnType<F>>;

package/dist/cjs/tracing/tracingContext.js

@@ -0,0 +1,28 @@
+"use strict";
+var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, generator) {
+    function adopt(value) { return value instanceof P ? value : new P(function (resolve) { resolve(value); }); }
+    return new (P || (P = Promise))(function (resolve, reject) {
+        function fulfilled(value) { try { step(generator.next(value)); } catch (e) { reject(e); } }
+        function rejected(value) { try { step(generator["throw"](value)); } catch (e) { reject(e); } }
+        function step(result) { result.done ? resolve(result.value) : adopt(result.value).then(fulfilled, rejected); }
+        step((generator = generator.apply(thisArg, _arguments || [])).next());
+    });
+};
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.getTracingContext = void 0;
+exports.runWithTracingContext = runWithTracingContext;
+const async_hooks_1 = require("async_hooks");
+const getDefaultTracingContext = () => ({
+    externalCustomerId: undefined,
+    externalProductId: undefined,
+    storePrompt: false,
+    metadata: undefined,
+});
+const tracingContextStorage = new async_hooks_1.AsyncLocalStorage();
+const getTracingContext = () => tracingContextStorage.getStore() || getDefaultTracingContext();
+exports.getTracingContext = getTracingContext;
+function runWithTracingContext(ctx, fn) {
+    return __awaiter(this, void 0, void 0, function* () {
+        return yield tracingContextStorage.run(ctx, fn);
+    });
+}

package/dist/cjs/tracing/wrappers/anthropicWrapper.js

@@ -11,11 +11,16 @@ var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, ge
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.PaidAnthropic = void 0;
 const api_1 = require("@opentelemetry/api");
-const tracing_js_1 = require("../tracing.js");
+const tracing_js_1 = require("tracing/tracing.js");
+const tracingContext_js_1 = require("tracing/tracingContext.js");
 class PaidAnthropic {
     constructor(anthropicClient) {
         this.anthropic = anthropicClient;
-        this.tracer = tracing_js_1.paidTracer;
+        const tracer = (0, tracing_js_1.getPaidTracer)();
+        if (!tracer) {
+            throw new Error("Paid tracer is not initialized, Make sure to call 'initializeTracing()' first");
+        }
+        this.tracer = tracer;
     }
     get messages() {
         return new MessagesWrapper(this.anthropic, this.tracer);
@@ -29,9 +34,8 @@ class MessagesWrapper {
     }
     create(params) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = (0, tracing_js_1.getCustomerIdStorage)();
-            const externalAgentId = (0, tracing_js_1.getAgentIdStorage)();
-            const token = (0, tracing_js_1.getTokenStorage)();
+            const token = (0, tracing_js_1.getToken)();
+            const { externalProductId, externalCustomerId } = (0, tracingContext_js_1.getTracingContext)();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -42,8 +46,8 @@ class MessagesWrapper {
                     external_customer_id: externalCustomerId,
                     token: token,
                 };
-                if (externalAgentId) {
-                    attributes["external_agent_id"] = externalAgentId;
+                if (externalProductId) {
+                    attributes["external_agent_id"] = externalProductId;
                 }
                 // Set request model from params
                 if (params.model) {

package/dist/cjs/tracing/wrappers/langchainCallback.js

@@ -12,13 +12,18 @@ Object.defineProperty(exports, "__esModule", { value: true });
 exports.PaidLangChainCallback = void 0;
 const base_1 = require("@langchain/core/callbacks/base");
 const api_1 = require("@opentelemetry/api");
-const tracing_js_1 = require("../tracing.js");
+const tracing_js_1 = require("tracing/tracing.js");
+const tracingContext_js_1 = require("tracing/tracingContext.js");
 class PaidLangChainCallback extends base_1.BaseCallbackHandler {
     constructor() {
         super();
         this.name = "PaidLangChainCallback";
         this.spans = new Map();
-        this.tracer = tracing_js_1.paidTracer;
+        const tracer = (0, tracing_js_1.getPaidTracer)();
+        if (!tracer) {
+            throw new Error("Paid tracer is not initialized, Make sure to call 'initializeTracing()' first");
+        }
+        this.tracer = tracer;
     }
     extractProvider(serialized) {
         if (!(serialized === null || serialized === void 0 ? void 0 : serialized.id))
@@ -40,9 +45,8 @@ class PaidLangChainCallback extends base_1.BaseCallbackHandler {
     }
     handleLLMStart(llm, prompts, runId, parentRunId, extraParams, tags, metadata) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = (0, tracing_js_1.getCustomerIdStorage)();
-            const externalAgentId = (0, tracing_js_1.getAgentIdStorage)();
-            const token = (0, tracing_js_1.getTokenStorage)();
+            const { externalProductId, externalCustomerId } = (0, tracingContext_js_1.getTracingContext)();
+            const token = (0, tracing_js_1.getToken)();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -57,8 +61,8 @@ class PaidLangChainCallback extends base_1.BaseCallbackHandler {
                 external_customer_id: externalCustomerId,
                 token: token,
             };
-            if (externalAgentId) {
-                attributes["external_agent_id"] = externalAgentId;
+            if (externalProductId) {
+                attributes["external_agent_id"] = externalProductId;
             }
             span.setAttributes(attributes);
             this.spans.set(runId, span);

package/dist/cjs/tracing/wrappers/mistralWrapper.js

@@ -11,11 +11,16 @@ var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, ge
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.PaidMistral = void 0;
 const api_1 = require("@opentelemetry/api");
-const tracing_js_1 = require("../tracing.js");
+const tracing_js_1 = require("tracing/tracing.js");
+const tracingContext_js_1 = require("tracing/tracingContext.js");
 class PaidMistral {
     constructor(mistralClient) {
         this.mistral = mistralClient;
-        this.tracer = tracing_js_1.paidTracer;
+        const tracer = (0, tracing_js_1.getPaidTracer)();
+        if (!tracer) {
+            throw new Error("Paid tracer is not initialized, Make sure to call 'initializeTracing()' first");
+        }
+        this.tracer = tracer;
     }
     get ocr() {
         return new OCRWrapper(this.mistral, this.tracer);
@@ -29,9 +34,8 @@ class OCRWrapper {
     }
     process(request, options) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = (0, tracing_js_1.getCustomerIdStorage)();
-            const externalAgentId = (0, tracing_js_1.getAgentIdStorage)();
-            const token = (0, tracing_js_1.getTokenStorage)();
+            const token = (0, tracing_js_1.getToken)();
+            const { externalProductId, externalCustomerId } = (0, tracingContext_js_1.getTracingContext)();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -43,8 +47,8 @@ class OCRWrapper {
                     external_customer_id: externalCustomerId,
                     token: token,
                 };
-                if (externalAgentId) {
-                    attributes["external_agent_id"] = externalAgentId;
+                if (externalProductId) {
+                    attributes["external_agent_id"] = externalProductId;
                 }
                 // Check if annotations are requested
                 if (request.bboxAnnotationFormat || request.documentAnnotationFormat) {

package/dist/cjs/tracing/wrappers/openAiWrapper.js

@@ -11,11 +11,16 @@ var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, ge
 Object.defineProperty(exports, "__esModule", { value: true });
 exports.PaidOpenAI = void 0;
 const api_1 = require("@opentelemetry/api");
-const tracing_js_1 = require("../tracing.js");
+const tracing_js_1 = require("tracing/tracing.js");
+const tracingContext_js_1 = require("tracing/tracingContext.js");
 class PaidOpenAI {
     constructor(openaiClient) {
         this.openai = openaiClient;
-        this.tracer = tracing_js_1.paidTracer;
+        const tracer = (0, tracing_js_1.getPaidTracer)();
+        if (!tracer) {
+            throw new Error("Paid tracer is not initialized, Make sure to call 'initializeTracing()' first");
+        }
+        this.tracer = tracer;
     }
     get chat() {
         return new ChatWrapper(this.openai, this.tracer);
@@ -47,9 +52,8 @@ class ChatCompletionsWrapper {
     }
     create(params) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = (0, tracing_js_1.getCustomerIdStorage)();
-            const externalAgentId = (0, tracing_js_1.getAgentIdStorage)();
-            const token = (0, tracing_js_1.getTokenStorage)();
+            const token = (0, tracing_js_1.getToken)();
+            const { externalProductId, externalCustomerId } = (0, tracingContext_js_1.getTracingContext)();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -60,8 +64,8 @@ class ChatCompletionsWrapper {
                 };
                 attributes["external_customer_id"] = externalCustomerId;
                 attributes["token"] = token;
-                if (externalAgentId) {
-                    attributes["external_agent_id"] = externalAgentId;
+                if (externalProductId) {
+                    attributes["external_agent_id"] = externalProductId;
                 }
                 span.setAttributes(attributes);
                 try {
@@ -95,9 +99,8 @@ class ResponsesWrapper {
     }
     create(params) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = (0, tracing_js_1.getCustomerIdStorage)();
-            const externalAgentId = (0, tracing_js_1.getAgentIdStorage)();
-            const token = (0, tracing_js_1.getTokenStorage)();
+            const { externalCustomerId, externalProductId } = (0, tracingContext_js_1.getTracingContext)();
+            const token = (0, tracing_js_1.getToken)();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -109,8 +112,8 @@ class ResponsesWrapper {
                 };
                 attributes["external_customer_id"] = externalCustomerId;
                 attributes["token"] = token;
-                if (externalAgentId) {
-                    attributes["external_agent_id"] = externalAgentId;
+                if (externalProductId) {
+                    attributes["external_agent_id"] = externalProductId;
                 }
                 span.setAttributes(attributes);
                 try {
@@ -150,9 +153,8 @@ class EmbeddingsWrapper {
     }
     create(params) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = (0, tracing_js_1.getCustomerIdStorage)();
-            const externalAgentId = (0, tracing_js_1.getAgentIdStorage)();
-            const token = (0, tracing_js_1.getTokenStorage)();
+            const { externalCustomerId, externalProductId } = (0, tracingContext_js_1.getTracingContext)();
+            const token = (0, tracing_js_1.getToken)();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -163,8 +165,8 @@ class EmbeddingsWrapper {
                 };
                 attributes["external_customer_id"] = externalCustomerId;
                 attributes["token"] = token;
-                if (externalAgentId) {
-                    attributes["external_agent_id"] = externalAgentId;
+                if (externalProductId) {
+                    attributes["external_agent_id"] = externalProductId;
                 }
                 span.setAttributes(attributes);
                 try {
@@ -197,9 +199,8 @@ class ImagesWrapper {
     }
     generate(params) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = (0, tracing_js_1.getCustomerIdStorage)();
-            const externalAgentId = (0, tracing_js_1.getAgentIdStorage)();
-            const token = (0, tracing_js_1.getTokenStorage)();
+            const { externalCustomerId, externalProductId } = (0, tracingContext_js_1.getTracingContext)();
+            const token = (0, tracing_js_1.getToken)();
             const model = params.model || "";
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
@@ -213,12 +214,12 @@ class ImagesWrapper {
                 };
                 attributes["external_customer_id"] = externalCustomerId;
                 attributes["token"] = token;
-                if (externalAgentId) {
-                    attributes["external_agent_id"] = externalAgentId;
+                if (externalProductId) {
+                    attributes["external_agent_id"] = externalProductId;
                 }
                 span.setAttributes(attributes);
                 try {
-                    const response = yield this.openai.images.generate(params);
+                    const response = (yield this.openai.images.generate(params));
                     span.setAttributes({
                         "gen_ai.image.count": (_a = params.n) !== null && _a !== void 0 ? _a : 1,
                         "gen_ai.image.size": params.size || "",