diff --git a/scripts/modules/ai-services-unified.js b/scripts/modules/ai-services-unified.js index e1ba0b6b..ac64d101 100644 --- a/scripts/modules/ai-services-unified.js +++ b/scripts/modules/ai-services-unified.js @@ -93,31 +93,55 @@ function _getProvider(providerName) { // Helper function to get cost for a specific model function _getCostForModel(providerName, modelId) { - const DEFAULT_COST = { inputCost: 0, outputCost: 0, currency: 'USD' }; + const DEFAULT_COST = { + inputCost: 0, + outputCost: 0, + currency: 'USD', + isUnknown: false + }; if (!MODEL_MAP || !MODEL_MAP[providerName]) { log( 'warn', `Provider "${providerName}" not found in MODEL_MAP. Cannot determine cost for model ${modelId}.` ); - return DEFAULT_COST; + return { ...DEFAULT_COST, isUnknown: true }; } const modelData = MODEL_MAP[providerName].find((m) => m.id === modelId); - if (!modelData?.cost_per_1m_tokens) { + if (!modelData) { log( 'debug', - `Cost data not found for model "${modelId}" under provider "${providerName}". Assuming zero cost.` + `Model "${modelId}" not found under provider "${providerName}". Assuming unknown cost.` ); - return DEFAULT_COST; + return { ...DEFAULT_COST, isUnknown: true }; + } + + // Check if cost_per_1m_tokens is explicitly null (unknown pricing) + if (modelData.cost_per_1m_tokens === null) { + log( + 'debug', + `Cost data is null for model "${modelId}" under provider "${providerName}". Pricing unknown.` + ); + return { ...DEFAULT_COST, isUnknown: true }; + } + + // Check if cost_per_1m_tokens is missing/undefined (also unknown) + if (modelData.cost_per_1m_tokens === undefined) { + log( + 'debug', + `Cost data not found for model "${modelId}" under provider "${providerName}". Pricing unknown.` + ); + return { ...DEFAULT_COST, isUnknown: true }; } const costs = modelData.cost_per_1m_tokens; return { inputCost: costs.input || 0, outputCost: costs.output || 0, - currency: costs.currency || 'USD' + currency: costs.currency || 'USD', + isUnknown: false }; } @@ -867,8 +891,8 @@ async function logAiUsage({ const timestamp = new Date().toISOString(); const totalTokens = (inputTokens || 0) + (outputTokens || 0); - // Destructure currency along with costs - const { inputCost, outputCost, currency } = _getCostForModel( + // Destructure currency along with costs and unknown flag + const { inputCost, outputCost, currency, isUnknown } = _getCostForModel( providerName, modelId ); @@ -890,7 +914,8 @@ async function logAiUsage({ outputTokens: outputTokens || 0, totalTokens, totalCost, - currency // Add currency to the telemetry data + currency, // Add currency to the telemetry data + isUnknownCost: isUnknown // Flag to indicate if pricing is unknown }; if (getDebugFlag()) { diff --git a/scripts/modules/ui.js b/scripts/modules/ui.js index 65450d0b..a82dc530 100644 --- a/scripts/modules/ui.js +++ b/scripts/modules/ui.js @@ -2310,7 +2310,8 @@ function displayAiUsageSummary(telemetryData, outputType = 'cli') { outputTokens, totalTokens, totalCost, - commandName + commandName, + isUnknownCost } = telemetryData; let summary = chalk.bold.blue('AI Usage Summary:') + '\n'; @@ -2320,7 +2321,10 @@ function displayAiUsageSummary(telemetryData, outputType = 'cli') { summary += chalk.gray( ` Tokens: ${totalTokens} (Input: ${inputTokens}, Output: ${outputTokens})\n` ); - summary += chalk.gray(` Est. Cost: $${totalCost.toFixed(6)}`); + + // Show "Unknown" if pricing data is not available, otherwise show the cost + const costDisplay = isUnknownCost ? 'Unknown' : `$${totalCost.toFixed(6)}`; + summary += chalk.gray(` Est. Cost: ${costDisplay}`); console.log( boxen(summary, { diff --git a/src/ai-providers/base-provider.js b/src/ai-providers/base-provider.js index 7409ad20..d9543400 100644 --- a/src/ai-providers/base-provider.js +++ b/src/ai-providers/base-provider.js @@ -176,12 +176,19 @@ export class BaseAIProvider { `${this.name} generateText completed successfully for model: ${params.modelId}` ); + const inputTokens = + result.usage?.inputTokens ?? result.usage?.promptTokens ?? 0; + const outputTokens = + result.usage?.outputTokens ?? result.usage?.completionTokens ?? 0; + const totalTokens = + result.usage?.totalTokens ?? inputTokens + outputTokens; + return { text: result.text, usage: { - inputTokens: result.usage?.promptTokens, - outputTokens: result.usage?.completionTokens, - totalTokens: result.usage?.totalTokens + inputTokens, + outputTokens, + totalTokens } }; } catch (error) { @@ -296,12 +303,19 @@ export class BaseAIProvider { `${this.name} generateObject completed successfully for model: ${params.modelId}` ); + const inputTokens = + result.usage?.inputTokens ?? result.usage?.promptTokens ?? 0; + const outputTokens = + result.usage?.outputTokens ?? result.usage?.completionTokens ?? 0; + const totalTokens = + result.usage?.totalTokens ?? inputTokens + outputTokens; + return { object: result.object, usage: { - inputTokens: result.usage?.promptTokens, - outputTokens: result.usage?.completionTokens, - totalTokens: result.usage?.totalTokens + inputTokens, + outputTokens, + totalTokens } }; } catch (error) {