@juspay/neurolink
Version:
Universal AI Development Platform with working MCP integration, multi-provider support, voice (TTS/STT/realtime), and professional CLI. 58+ external MCP servers discoverable, multimodal file processing, RAG pipelines. Build, test, and deploy AI applicatio
416 lines • 14.2 kB
JavaScript
/**
* Token Usage Tracker
* Aggregates token usage and cost across spans
*/
/**
* Built-in model pricing database (approximate, subject to change)
*/
const MODEL_PRICING = {
// OpenAI
"gpt-4o": {
inputPricePerMillion: 2.5,
outputPricePerMillion: 10.0,
cachedInputPricePerMillion: 1.25,
},
"gpt-4o-mini": {
inputPricePerMillion: 0.15,
outputPricePerMillion: 0.6,
cachedInputPricePerMillion: 0.075,
},
o1: {
inputPricePerMillion: 15.0,
outputPricePerMillion: 60.0,
cachedInputPricePerMillion: 7.5,
},
"o1-mini": {
inputPricePerMillion: 3.0,
outputPricePerMillion: 12.0,
cachedInputPricePerMillion: 1.5,
},
// Anthropic
"claude-sonnet-4-20250514": {
inputPricePerMillion: 3.0,
outputPricePerMillion: 15.0,
cachedInputPricePerMillion: 0.3,
},
"claude-3-5-sonnet-20241022": {
inputPricePerMillion: 3.0,
outputPricePerMillion: 15.0,
cachedInputPricePerMillion: 0.3,
},
"claude-3-5-haiku-20241022": {
inputPricePerMillion: 0.8,
outputPricePerMillion: 4.0,
cachedInputPricePerMillion: 0.08,
},
// Google
"gemini-2.5-flash": {
inputPricePerMillion: 0.15,
outputPricePerMillion: 0.6,
},
"gemini-2.5-pro": {
inputPricePerMillion: 1.25,
outputPricePerMillion: 10.0,
},
"gemini-2.0-flash-exp": {
inputPricePerMillion: 0.15,
outputPricePerMillion: 0.6,
},
// Mistral
"mistral-large-latest": {
inputPricePerMillion: 2.0,
outputPricePerMillion: 6.0,
},
"mistral-small-latest": {
inputPricePerMillion: 0.2,
outputPricePerMillion: 0.6,
},
};
/**
* Token tracker for aggregating usage across spans
*/
export class TokenTracker {
stats = {
totalInputTokens: 0,
totalOutputTokens: 0,
totalTokens: 0,
cacheReadTokens: 0,
cacheCreationTokens: 0,
reasoningTokens: 0,
totalCost: 0,
byProvider: new Map(),
byModel: new Map(),
bySpanType: new Map(),
};
customPricing = new Map();
/**
* Set custom pricing for a single model
* @param modelName - The model name (e.g., "gpt-4o", "claude-3-5-sonnet")
* @param pricing - The pricing information
*/
setObservabilityModelPricing(modelName, pricing) {
this.customPricing.set(modelName, pricing);
}
/**
* Update pricing for an existing model (alias for setObservabilityModelPricing)
* @param model - The model name
* @param pricing - The new pricing information
*/
updatePricing(model, pricing) {
this.customPricing.set(model, pricing);
}
/**
* Load pricing configuration from a config object
* Useful for loading pricing from environment or config files
* @param config - Record of model names to pricing information
*/
loadPricingFromConfig(config) {
for (const [model, pricing] of Object.entries(config)) {
this.customPricing.set(model, pricing);
}
}
/**
* Get pricing for a specific model
* @param model - The model name
* @returns The pricing information or undefined if not found
*/
getModelPricing(model) {
return this.customPricing.get(model) ?? MODEL_PRICING[model];
}
/**
* Get all available model pricing (custom + built-in)
* @returns Record of all model pricing
*/
getAllPricing() {
const allPricing = {
...MODEL_PRICING,
};
// Custom pricing takes precedence
const customPricingEntries = Array.from(this.customPricing.entries());
for (const [model, pricing] of customPricingEntries) {
allPricing[model] = pricing;
}
return allPricing;
}
/**
* Remove custom pricing for a model (falls back to built-in)
* @param model - The model name to remove custom pricing for
*/
removeCustomPricing(model) {
return this.customPricing.delete(model);
}
/**
* Track token usage from a span
*/
trackSpan(span) {
const attrs = span.attributes;
const inputTokens = attrs["ai.tokens.input"] ?? 0;
const outputTokens = attrs["ai.tokens.output"] ?? 0;
const totalTokens = attrs["ai.tokens.total"] ?? inputTokens + outputTokens;
const cacheRead = attrs["ai.tokens.cache_read"] ?? 0;
const cacheCreation = attrs["ai.tokens.cache_creation"] ?? 0;
const reasoning = attrs["ai.tokens.reasoning"] ?? 0;
const cost = attrs["ai.cost.total"] ??
this.calculateCost(attrs, inputTokens, outputTokens, cacheRead);
// Update totals
this.stats.totalInputTokens += inputTokens;
this.stats.totalOutputTokens += outputTokens;
this.stats.totalTokens += totalTokens;
this.stats.cacheReadTokens += cacheRead;
this.stats.cacheCreationTokens += cacheCreation;
this.stats.reasoningTokens += reasoning;
this.stats.totalCost += cost;
// Update by provider
const provider = attrs["ai.provider"];
if (provider) {
const existing = this.stats.byProvider.get(provider) ?? {
provider,
inputTokens: 0,
outputTokens: 0,
totalTokens: 0,
cost: 0,
requestCount: 0,
};
existing.inputTokens += inputTokens;
existing.outputTokens += outputTokens;
existing.totalTokens += totalTokens;
existing.cost += cost;
existing.requestCount += 1;
this.stats.byProvider.set(provider, existing);
}
// Update by model
const model = attrs["ai.model"];
if (model) {
const existing = this.stats.byModel.get(model) ?? {
model,
provider: provider ?? "unknown",
inputTokens: 0,
outputTokens: 0,
totalTokens: 0,
cost: 0,
requestCount: 0,
avgTokensPerRequest: 0,
};
existing.inputTokens += inputTokens;
existing.outputTokens += outputTokens;
existing.totalTokens += totalTokens;
existing.cost += cost;
existing.requestCount += 1;
existing.avgTokensPerRequest =
existing.totalTokens / existing.requestCount;
this.stats.byModel.set(model, existing);
}
// Update by span type
const currentTypeTotal = this.stats.bySpanType.get(span.type) ?? 0;
this.stats.bySpanType.set(span.type, currentTypeTotal + totalTokens);
}
/**
* Calculate cost from token counts and provider/model
*/
calculateCost(attrs, inputTokens, outputTokens, cacheReadTokens = 0) {
const model = attrs["ai.model"];
if (!model) {
return 0;
}
// Check custom pricing first, then fall back to built-in
const pricing = this.customPricing.get(model) ?? MODEL_PRICING[model];
if (!pricing) {
return 0;
}
const regularInputTokens = inputTokens - cacheReadTokens;
const regularCost = (regularInputTokens / 1_000_000) * pricing.inputPricePerMillion;
const cachedCost = (cacheReadTokens / 1_000_000) *
(pricing.cachedInputPricePerMillion ?? pricing.inputPricePerMillion);
const outputCost = (outputTokens / 1_000_000) * pricing.outputPricePerMillion;
return regularCost + cachedCost + outputCost;
}
/**
* Track token usage from a simple usage object
* This is a convenience method for tracking usage without a full span
* @param usage - Token usage data
*/
trackUsage(usage) {
const inputTokens = usage.promptTokens ?? 0;
const outputTokens = usage.completionTokens ?? 0;
const totalTokens = usage.totalTokens ?? inputTokens + outputTokens;
// Update totals
this.stats.totalInputTokens += inputTokens;
this.stats.totalOutputTokens += outputTokens;
this.stats.totalTokens += totalTokens;
// Calculate cost if model is provided
let cost = 0;
if (usage.model) {
const pricing = this.customPricing.get(usage.model) ?? MODEL_PRICING[usage.model];
if (pricing) {
cost =
(inputTokens / 1_000_000) * pricing.inputPricePerMillion +
(outputTokens / 1_000_000) * pricing.outputPricePerMillion;
}
}
this.stats.totalCost += cost;
// Update by provider if provided
if (usage.provider) {
const existing = this.stats.byProvider.get(usage.provider) ?? {
provider: usage.provider,
inputTokens: 0,
outputTokens: 0,
totalTokens: 0,
cost: 0,
requestCount: 0,
};
existing.inputTokens += inputTokens;
existing.outputTokens += outputTokens;
existing.totalTokens += totalTokens;
existing.cost += cost;
existing.requestCount += 1;
this.stats.byProvider.set(usage.provider, existing);
}
// Update by model if provided
if (usage.model) {
const existing = this.stats.byModel.get(usage.model) ?? {
model: usage.model,
provider: usage.provider ?? "unknown",
inputTokens: 0,
outputTokens: 0,
totalTokens: 0,
cost: 0,
requestCount: 0,
avgTokensPerRequest: 0,
};
existing.inputTokens += inputTokens;
existing.outputTokens += outputTokens;
existing.totalTokens += totalTokens;
existing.cost += cost;
existing.requestCount += 1;
existing.avgTokensPerRequest =
existing.totalTokens / existing.requestCount;
this.stats.byModel.set(usage.model, existing);
}
}
/**
* Get current stats
*/
getStats() {
return { ...this.stats };
}
/**
* Get stats for a specific time window of spans
*/
getStatsForWindow(spans) {
const tracker = new TokenTracker();
// Copy custom pricing so windowed calculations use the same rates
for (const [model, pricing] of this.customPricing) {
tracker.setObservabilityModelPricing(model, pricing);
}
for (const span of spans) {
tracker.trackSpan(span);
}
return tracker.getStats();
}
/**
* Reset all stats
*/
reset() {
this.stats = {
totalInputTokens: 0,
totalOutputTokens: 0,
totalTokens: 0,
cacheReadTokens: 0,
cacheCreationTokens: 0,
reasoningTokens: 0,
totalCost: 0,
byProvider: new Map(),
byModel: new Map(),
bySpanType: new Map(),
};
}
/**
* Export stats as JSON
*/
toJSON() {
return {
totalInputTokens: this.stats.totalInputTokens,
totalOutputTokens: this.stats.totalOutputTokens,
totalTokens: this.stats.totalTokens,
cacheReadTokens: this.stats.cacheReadTokens,
cacheCreationTokens: this.stats.cacheCreationTokens,
reasoningTokens: this.stats.reasoningTokens,
totalCost: this.stats.totalCost,
byProvider: Object.fromEntries(this.stats.byProvider),
byModel: Object.fromEntries(this.stats.byModel),
bySpanType: Object.fromEntries(this.stats.bySpanType),
};
}
/**
* Format cost as currency string
*/
formatCost(cost, currency = "USD") {
return new Intl.NumberFormat("en-US", {
style: "currency",
currency,
minimumFractionDigits: 4,
}).format(cost);
}
/**
* Get a summary string of current stats
*/
getSummary() {
const stats = this.stats;
return [
`Token Usage Summary:`,
` Input tokens: ${stats.totalInputTokens.toLocaleString()}`,
` Output tokens: ${stats.totalOutputTokens.toLocaleString()}`,
` Total tokens: ${stats.totalTokens.toLocaleString()}`,
` Total cost: ${this.formatCost(stats.totalCost)}`,
stats.cacheReadTokens > 0
? ` Cache read tokens: ${stats.cacheReadTokens.toLocaleString()}`
: "",
stats.reasoningTokens > 0
? ` Reasoning tokens: ${stats.reasoningTokens.toLocaleString()}`
: "",
]
.filter(Boolean)
.join("\n");
}
}
/**
* Enrich span with token usage attributes
*/
export function enrichSpanWithTokenUsage(span, usage) {
return {
...span,
attributes: {
...span.attributes,
"ai.tokens.input": usage.promptTokens ?? 0,
"ai.tokens.output": usage.completionTokens ?? 0,
"ai.tokens.total": usage.totalTokens ??
(usage.promptTokens ?? 0) + (usage.completionTokens ?? 0),
"ai.tokens.cache_creation": usage.cacheCreationTokens,
"ai.tokens.cache_read": usage.cacheReadTokens,
"ai.tokens.reasoning": usage.reasoningTokens,
},
};
}
/**
* Global token tracker instance (singleton pattern from main)
*/
let globalTokenTracker = null;
/**
* Get the global token tracker instance
*/
export function getTokenTracker() {
if (!globalTokenTracker) {
globalTokenTracker = new TokenTracker();
}
return globalTokenTracker;
}
/**
* Reset the global token tracker (for testing)
*/
export function resetTokenTracker() {
if (globalTokenTracker) {
globalTokenTracker.reset();
}
globalTokenTracker = null;
}
//# sourceMappingURL=tokenTracker.js.map