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

package/dist/esm/tracing/signal.mjs

@@ -1,21 +1,38 @@
 import { SpanStatusCode } from "@opentelemetry/api";
-import { getCustomerIdStorage, getAgentIdStorage, getTokenStorage } from "./tracing.mjs";
-import { paidTracer } from "./tracing.mjs";
-export function _signal(eventName, enableCostTracing, data) {
-    if (!eventName) {
-        throw new Error("Event name is required for signal.");
+import { getPaidTracer, getToken } from "./tracing.mjs";
+import { getTracingContext } from "./tracingContext.mjs";
+/**
+ * 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.
+ */
+export function signal(eventName, enableCostTracing = false, data) {
+    const paidTracer = getPaidTracer();
+    const token = getToken();
+    const { externalCustomerId, externalProductId } = getTracingContext();
+    if (!token || !paidTracer) {
+        throw new Error(`Tracing is not initialized. Make sure you called 'initializeTracing()'`);
     }
-    const externalCustomerId = getCustomerIdStorage();
-    const externalAgentId = getAgentIdStorage();
-    const token = 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()`);
     }
     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,
             };
@@ -34,7 +51,6 @@ export function _signal(eventName, enableCostTracing, data) {
                 attributes["data"] = JSON.stringify(data);
             }
             span.setAttributes(attributes);
-            // Mark span as successful
             span.setStatus({ code: SpanStatusCode.OK });
         }
         catch (error) {

package/dist/esm/tracing/spanProcessor.d.mts

@@ -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/esm/tracing/spanProcessor.mjs

@@ -0,0 +1,78 @@
+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());
+    });
+};
+import { getTracingContext } from "./tracingContext.mjs";
+import { getToken } from "./tracing.mjs";
+export class PaidSpanProcessor {
+    onStart(span, _parentContext) {
+        const { name } = span;
+        const { storePrompt, externalCustomerId, externalProductId: externalAgentId } = 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 = 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;
+        });
+    }
+}
+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/esm/tracing/tracing.d.mts

@@ -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/esm/tracing/tracing.mjs

@@ -7,12 +7,12 @@ var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, ge
         step((generator = generator.apply(thisArg, _arguments || [])).next());
     });
 };
-import { NodeTracerProvider } from "@opentelemetry/sdk-trace-node";
-import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-http";
-import { BatchSpanProcessor } from "@opentelemetry/sdk-trace-base";
 import { SpanStatusCode } from "@opentelemetry/api";
-import { AsyncLocalStorage } from "async_hooks";
+import { OTLPTraceExporter } from "@opentelemetry/exporter-trace-otlp-http";
+import { NodeTracerProvider, SimpleSpanProcessor } from "@opentelemetry/sdk-trace-node";
 import winston from "winston";
+import { runWithTracingContext } from "./tracingContext.mjs";
+import { PaidSpanProcessor } from "./spanProcessor.mjs";
 export const logger = winston.createLogger({
     level: "silent", // Default to 'silent' to avoid logging unless set via environment variable
     format: winston.format.simple(),
@@ -22,29 +22,21 @@ const logLevel = process.env.PAID_LOG_LEVEL;
 if (logLevel) {
     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 OTLPTraceExporter({ url: COLLECTOR_ENDPOINT });
-let spanProcessor = new BatchSpanProcessor(paidExporter);
-let paidTracerProvider = new NodeTracerProvider({ spanProcessors: [spanProcessor] });
-paidTracerProvider.register();
-export let paidTracer = paidTracerProvider.getTracer("paid.node");
-// storage for passing info to child spans
-const customerIdStorage = new AsyncLocalStorage();
-const agentIdStorage = new AsyncLocalStorage();
-const tokenStorage = new AsyncLocalStorage();
-export const getCustomerIdStorage = () => customerIdStorage.getStore();
-export const getAgentIdStorage = () => agentIdStorage.getStore();
-export const getTokenStorage = () => tokenStorage.getStore();
-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;
+export function getToken() {
+    return paidApiToken;
+}
+let paidTracerProvider = undefined;
+export function getPaidTracerProvider() {
+    return paidTracerProvider;
+}
+let paidTracer = undefined;
+export 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) {
@@ -57,48 +49,63 @@ function setupGracefulShutdown(shuttable) {
                 .catch((error) => logger.error(`Error shutting down Paid tracing SDK ${error}`));
         });
     });
-}
-export function _initializeTracing(apiKey, collectorEndpoint) {
-    if (getToken()) {
-        throw new Error("Tracing SDK is already initialized.");
+};
+export function initializeTracing(apiKey, collectorEndpoint) {
+    const paidEnabled = (process.env.PAID_ENABLED || "true") !== "false";
+    if (!paidEnabled) {
+        logger.info("Paid tracing is disabled via PAID_ENABLED environment variable");
+        return;
+    }
+    const token = getToken();
+    if (!!token) {
+        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) {
+            logger.error("API key must be provided via PAID_API_KEY environment variable");
+            return;
+        }
+        paidApiToken = envKey;
     }
-    if (collectorEndpoint) {
-        paidExporter = new OTLPTraceExporter({ url: collectorEndpoint });
-        spanProcessor = new BatchSpanProcessor(paidExporter);
-        paidTracerProvider = new NodeTracerProvider({ spanProcessors: [spanProcessor] });
-        paidTracerProvider.register();
-        paidTracer = paidTracerProvider.getTracer("paid.node");
+    else {
+        paidApiToken = apiKey;
     }
+    const url = collectorEndpoint || DEFAULT_COLLECTOR_ENDPOINT;
+    const exporter = new OTLPTraceExporter({ url });
+    const spanProcessor = new SimpleSpanProcessor(exporter);
+    paidTracerProvider = new NodeTracerProvider({
+        spanProcessors: [spanProcessor, new PaidSpanProcessor()],
+    });
+    paidTracerProvider.register();
+    paidTracer = paidTracerProvider.getTracer("paid.node");
     setupGracefulShutdown(spanProcessor);
-    setToken(apiKey);
-    const collectorAddress = collectorEndpoint || COLLECTOR_ENDPOINT;
-    logger.info(`Paid tracing SDK initialized with collector endpoint: ${collectorAddress}`);
+    logger.info(`Paid tracing SDK initialized with collector endpoint: ${url}`);
 }
-export function _trace(externalCustomerId, fn, externalAgentId, ...args) {
+export 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 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 runWithTracingContext({
+                    externalCustomerId,
+                    externalProductId: externalAgentId,
+                    storePrompt,
+                    metadata,
+                }, () => __awaiter(this, void 0, void 0, function* () { return yield fn(...args); }));
                 span.setStatus({ code: SpanStatusCode.OK });
-                return result;
+                return res;
             }
             catch (error) {
                 span.setStatus({

package/dist/esm/tracing/tracingContext.d.mts

@@ -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/esm/tracing/tracingContext.mjs

@@ -0,0 +1,23 @@
+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());
+    });
+};
+import { AsyncLocalStorage } from "async_hooks";
+const getDefaultTracingContext = () => ({
+    externalCustomerId: undefined,
+    externalProductId: undefined,
+    storePrompt: false,
+    metadata: undefined,
+});
+const tracingContextStorage = new AsyncLocalStorage();
+export const getTracingContext = () => tracingContextStorage.getStore() || getDefaultTracingContext();
+export function runWithTracingContext(ctx, fn) {
+    return __awaiter(this, void 0, void 0, function* () {
+        return yield tracingContextStorage.run(ctx, fn);
+    });
+}

package/dist/esm/tracing/wrappers/anthropicWrapper.mjs

@@ -8,11 +8,16 @@ var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, ge
     });
 };
 import { SpanStatusCode } from "@opentelemetry/api";
-import { getCustomerIdStorage, getAgentIdStorage, getTokenStorage, paidTracer } from "../tracing.mjs";
+import { getPaidTracer, getToken } from "tracing/tracing.js";
+import { getTracingContext } from "tracing/tracingContext.js";
 export class PaidAnthropic {
     constructor(anthropicClient) {
         this.anthropic = anthropicClient;
-        this.tracer = paidTracer;
+        const tracer = 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);
@@ -25,9 +30,8 @@ class MessagesWrapper {
     }
     create(params) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = getCustomerIdStorage();
-            const externalAgentId = getAgentIdStorage();
-            const token = getTokenStorage();
+            const token = getToken();
+            const { externalProductId, externalCustomerId } = getTracingContext();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -38,8 +42,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/esm/tracing/wrappers/langchainCallback.mjs

@@ -9,13 +9,18 @@ var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, ge
 };
 import { BaseCallbackHandler } from "@langchain/core/callbacks/base";
 import { SpanStatusCode } from "@opentelemetry/api";
-import { getCustomerIdStorage, getAgentIdStorage, getTokenStorage, paidTracer } from "../tracing.mjs";
+import { getPaidTracer, getToken } from "tracing/tracing.js";
+import { getTracingContext } from "tracing/tracingContext.js";
 export class PaidLangChainCallback extends BaseCallbackHandler {
     constructor() {
         super();
         this.name = "PaidLangChainCallback";
         this.spans = new Map();
-        this.tracer = paidTracer;
+        const tracer = 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))
@@ -37,9 +42,8 @@ export class PaidLangChainCallback extends BaseCallbackHandler {
     }
     handleLLMStart(llm, prompts, runId, parentRunId, extraParams, tags, metadata) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = getCustomerIdStorage();
-            const externalAgentId = getAgentIdStorage();
-            const token = getTokenStorage();
+            const { externalProductId, externalCustomerId } = getTracingContext();
+            const token = getToken();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -54,8 +58,8 @@ export class PaidLangChainCallback extends 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/esm/tracing/wrappers/mistralWrapper.mjs

@@ -8,11 +8,16 @@ var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, ge
     });
 };
 import { SpanStatusCode } from "@opentelemetry/api";
-import { getCustomerIdStorage, getAgentIdStorage, getTokenStorage, paidTracer } from "../tracing.mjs";
+import { getPaidTracer, getToken } from "tracing/tracing.js";
+import { getTracingContext } from "tracing/tracingContext.js";
 export class PaidMistral {
     constructor(mistralClient) {
         this.mistral = mistralClient;
-        this.tracer = paidTracer;
+        const tracer = 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);
@@ -25,9 +30,8 @@ class OCRWrapper {
     }
     process(request, options) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = getCustomerIdStorage();
-            const externalAgentId = getAgentIdStorage();
-            const token = getTokenStorage();
+            const token = getToken();
+            const { externalProductId, externalCustomerId } = getTracingContext();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -39,8 +43,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/esm/tracing/wrappers/openAiWrapper.mjs

@@ -8,11 +8,16 @@ var __awaiter = (this && this.__awaiter) || function (thisArg, _arguments, P, ge
     });
 };
 import { SpanStatusCode } from "@opentelemetry/api";
-import { getCustomerIdStorage, getAgentIdStorage, getTokenStorage, paidTracer } from "../tracing.mjs";
+import { getPaidTracer, getToken } from "tracing/tracing.js";
+import { getTracingContext } from "tracing/tracingContext.js";
 export class PaidOpenAI {
     constructor(openaiClient) {
         this.openai = openaiClient;
-        this.tracer = paidTracer;
+        const tracer = 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);
@@ -43,9 +48,8 @@ class ChatCompletionsWrapper {
     }
     create(params) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = getCustomerIdStorage();
-            const externalAgentId = getAgentIdStorage();
-            const token = getTokenStorage();
+            const token = getToken();
+            const { externalProductId, externalCustomerId } = getTracingContext();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -56,8 +60,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 {
@@ -91,9 +95,8 @@ class ResponsesWrapper {
     }
     create(params) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = getCustomerIdStorage();
-            const externalAgentId = getAgentIdStorage();
-            const token = getTokenStorage();
+            const { externalCustomerId, externalProductId } = getTracingContext();
+            const token = getToken();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -105,8 +108,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 {
@@ -146,9 +149,8 @@ class EmbeddingsWrapper {
     }
     create(params) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = getCustomerIdStorage();
-            const externalAgentId = getAgentIdStorage();
-            const token = getTokenStorage();
+            const { externalCustomerId, externalProductId } = getTracingContext();
+            const token = getToken();
             if (!token || !externalCustomerId) {
                 throw new Error("No token or externalCustomerId: This wrapper should be used inside a callback to paid.trace().");
             }
@@ -159,8 +161,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 {
@@ -193,9 +195,8 @@ class ImagesWrapper {
     }
     generate(params) {
         return __awaiter(this, void 0, void 0, function* () {
-            const externalCustomerId = getCustomerIdStorage();
-            const externalAgentId = getAgentIdStorage();
-            const token = getTokenStorage();
+            const { externalCustomerId, externalProductId } = getTracingContext();
+            const token = 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().");
@@ -209,12 +210,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 || "",