Skip to content

Commit a872789

Browse files
khromovkhromov
committed
add pricing mapping
1 parent c0bd003 commit a872789

File tree

4 files changed

+206
-53
lines changed

4 files changed

+206
-53
lines changed

‎index.ts‎

Lines changed: 109 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,14 @@ import {
2323
} from "./lib/output-test-runner.ts";
2424
import { resultWriteTool, testComponentTool } from "./lib/tools/index.ts";
2525
import {
26+
lookupModelPricing,
27+
lookupModelPricingByKey,
2628
getModelPricingDisplay,
2729
calculateCost,
2830
formatCost,
2931
isPricingAvailable,
32+
type ModelPricing,
33+
type ModelPricingLookup,
3034
} from "./lib/pricing.ts";
3135
import type { LanguageModel } from "ai";
3236

@@ -97,8 +101,8 @@ function extractResultWriteContent(steps: unknown[]): string | null {
97101
*/
98102
function calculateTotalCost(
99103
tests: SingleTestResult[],
100-
modelString: string,
101-
): TotalCostInfo | null {
104+
pricing: ModelPricing,
105+
): TotalCostInfo {
102106
let totalInputTokens = 0;
103107
let totalOutputTokens = 0;
104108
let totalCachedInputTokens = 0;
@@ -112,14 +116,12 @@ function calculateTotalCost(
112116
}
113117

114118
const costResult = calculateCost(
115-
modelString,
119+
pricing,
116120
totalInputTokens,
117121
totalOutputTokens,
118122
totalCachedInputTokens,
119123
);
120124

121-
if (!costResult) return null;
122-
123125
return {
124126
inputCost: costResult.inputCost,
125127
outputCost: costResult.outputCost,
@@ -131,6 +133,69 @@ function calculateTotalCost(
131133
};
132134
}
133135

136+
/**
137+
* Resolve pricing lookup based on environment variables
138+
* Returns the pricing lookup result or null if disabled
139+
* Throws an error if pricing cannot be found and is not disabled
140+
*/
141+
function resolvePricingLookup(modelString: string): ModelPricingLookup | null {
142+
const costDisabled = process.env.MODEL_COST_DISABLED === "true";
143+
const explicitCostName = process.env.MODEL_COST_NAME;
144+
145+
// If cost calculation is explicitly disabled, return null
146+
if (costDisabled) {
147+
return null;
148+
}
149+
150+
// Check if pricing data file exists
151+
if (!isPricingAvailable()) {
152+
console.error(
153+
`\n✗ Model pricing file not found. Run 'bun run update-model-pricing' to download it.`,
154+
);
155+
console.error(
156+
` Or set MODEL_COST_DISABLED=true to skip cost calculation.\n`,
157+
);
158+
process.exit(1);
159+
}
160+
161+
// If explicit cost name is provided, use that
162+
if (explicitCostName) {
163+
const lookup = lookupModelPricingByKey(explicitCostName);
164+
if (!lookup) {
165+
console.error(
166+
`\n✗ Could not find pricing for MODEL_COST_NAME="${explicitCostName}" in model-pricing.json`,
167+
);
168+
console.error(
169+
` Check that the key exists in data/model-pricing.json.\n`,
170+
);
171+
process.exit(1);
172+
}
173+
return lookup;
174+
}
175+
176+
// Try automatic lookup
177+
const lookup = lookupModelPricing(modelString);
178+
if (!lookup) {
179+
console.error(
180+
`\n✗ Could not find pricing for model "${modelString}" in model-pricing.json`,
181+
);
182+
console.error(`\n Options:`);
183+
console.error(
184+
` 1. Set MODEL_COST_NAME=<key> to explicitly specify the pricing key`,
185+
);
186+
console.error(` Example: MODEL_COST_NAME=claude-sonnet-4-20250514`);
187+
console.error(
188+
` 2. Set MODEL_COST_DISABLED=true to skip cost calculation`,
189+
);
190+
console.error(
191+
`\n Browse data/model-pricing.json to find the correct key for your model.\n`,
192+
);
193+
process.exit(1);
194+
}
195+
196+
return lookup;
197+
}
198+
134199
/**
135200
* Run a single test with the AI agent
136201
*/
@@ -281,10 +346,24 @@ async function main() {
281346
const isHttpTransport = mcpEnabled && isHttpUrl(mcpServerUrl);
282347
const mcpTransportType = isHttpTransport ? "HTTP" : "StdIO";
283348

349+
// Load environment configuration
350+
const envConfig = loadEnvConfig();
351+
284352
console.log("╔════════════════════════════════════════════════════╗");
285353
console.log("║ SvelteBench 2.0 - Multi-Test ║");
286354
console.log("╚════════════════════════════════════════════════════╝");
287-
console.log(`Model: ${process.env.MODEL}`);
355+
console.log(`Model: ${envConfig.modelString}`);
356+
357+
// Resolve pricing lookup
358+
const costDisabled = process.env.MODEL_COST_DISABLED === "true";
359+
const pricingLookup = resolvePricingLookup(envConfig.modelString);
360+
361+
if (pricingLookup) {
362+
console.log(`Pricing model mapped: ${pricingLookup.matchedKey}`);
363+
} else if (costDisabled) {
364+
console.log(`Pricing: Disabled (MODEL_COST_DISABLED=true)`);
365+
}
366+
288367
console.log(`MCP Integration: ${mcpEnabled ? "Enabled" : "Disabled"}`);
289368
if (mcpEnabled) {
290369
console.log(`MCP Transport: ${mcpTransportType}`);
@@ -298,10 +377,6 @@ async function main() {
298377
`TestComponent Tool: ${testComponentEnabled ? "Enabled" : "Disabled"}`,
299378
);
300379

301-
// Check pricing availability
302-
const hasPricing = isPricingAvailable();
303-
console.log(`Pricing Data: ${hasPricing ? "Available" : "Not available (run 'bun run update-model-pricing' to download)"}`);
304-
305380
// Discover all tests
306381
console.log("\n📁 Discovering tests...");
307382
const tests = discoverTests();
@@ -340,8 +415,7 @@ async function main() {
340415
}
341416
}
342417

343-
// Load environment configuration and get model provider
344-
const envConfig = loadEnvConfig();
418+
// Get model provider
345419
const model = getModelProvider(envConfig);
346420

347421
// Run all tests
@@ -397,17 +471,31 @@ async function main() {
397471
`Total: ${passed} passed, ${failed} failed, ${skipped} skipped (${(totalDuration / 1000).toFixed(1)}s)`,
398472
);
399473

400-
// Calculate total cost
401-
const totalCost = calculateTotalCost(testResults, envConfig.modelString);
402-
const pricingDisplay = getModelPricingDisplay(envConfig.modelString);
474+
// Calculate total cost if pricing is available
475+
let totalCost: TotalCostInfo | null = null;
476+
let pricingInfo: PricingInfo | null = null;
477+
478+
if (pricingLookup) {
479+
totalCost = calculateTotalCost(testResults, pricingLookup.pricing);
480+
const pricingDisplay = getModelPricingDisplay(pricingLookup.pricing);
481+
pricingInfo = {
482+
inputCostPerMTok: pricingDisplay.inputCostPerMTok,
483+
outputCostPerMTok: pricingDisplay.outputCostPerMTok,
484+
cacheReadCostPerMTok: pricingDisplay.cacheReadCostPerMTok,
485+
};
403486

404-
if (totalCost) {
405487
console.log("\n💰 Cost Summary");
406488
console.log("─".repeat(50));
407-
console.log(`Input tokens: ${totalCost.inputTokens.toLocaleString()} (${formatCost(totalCost.inputCost)})`);
408-
console.log(`Output tokens: ${totalCost.outputTokens.toLocaleString()} (${formatCost(totalCost.outputCost)})`);
489+
console.log(
490+
`Input tokens: ${totalCost.inputTokens.toLocaleString()} (${formatCost(totalCost.inputCost)})`,
491+
);
492+
console.log(
493+
`Output tokens: ${totalCost.outputTokens.toLocaleString()} (${formatCost(totalCost.outputCost)})`,
494+
);
409495
if (totalCost.cachedInputTokens > 0) {
410-
console.log(`Cached tokens: ${totalCost.cachedInputTokens.toLocaleString()} (${formatCost(totalCost.cacheReadCost)})`);
496+
console.log(
497+
`Cached tokens: ${totalCost.cachedInputTokens.toLocaleString()} (${formatCost(totalCost.cacheReadCost)})`,
498+
);
411499
}
412500
console.log(`Total cost: ${formatCost(totalCost.totalCost)}`);
413501
}
@@ -424,15 +512,6 @@ async function main() {
424512
const jsonPath = `${resultsDir}/${jsonFilename}`;
425513
const htmlPath = `${resultsDir}/${htmlFilename}`;
426514

427-
// Build pricing info for metadata
428-
const pricing: PricingInfo | null = pricingDisplay
429-
? {
430-
inputCostPerMTok: pricingDisplay.inputCostPerMTok,
431-
outputCostPerMTok: pricingDisplay.outputCostPerMTok,
432-
cacheReadCostPerMTok: pricingDisplay.cacheReadCostPerMTok,
433-
}
434-
: null;
435-
436515
// Build the result data
437516
const resultData: MultiTestResultData = {
438517
tests: testResults,
@@ -442,7 +521,8 @@ async function main() {
442521
mcpTransportType: mcpEnabled ? mcpTransportType : null,
443522
timestamp: new Date().toISOString(),
444523
model: envConfig.modelString,
445-
pricing,
524+
pricingKey: pricingLookup?.matchedKey ?? null,
525+
pricing: pricingInfo,
446526
totalCost,
447527
},
448528
};

‎lib/pricing.ts‎

Lines changed: 80 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,14 @@ export interface ModelPricingDisplay {
3535
cacheReadCostPerMTok?: number;
3636
}
3737

38+
/**
39+
* Result of looking up a model's pricing
40+
*/
41+
export interface ModelPricingLookup {
42+
pricing: ModelPricing;
43+
matchedKey: string;
44+
}
45+
3846
// Cache the loaded pricing data
3947
let pricingData: Record<string, unknown> | null = null;
4048

@@ -108,10 +116,10 @@ function normalizeModelForLookup(modelString: string): string[] {
108116
}
109117

110118
/**
111-
* Get pricing information for a model
119+
* Look up pricing information for a model, returning both the pricing and the matched key
112120
* Returns null if pricing is not found
113121
*/
114-
export function getModelPricing(modelString: string): ModelPricing | null {
122+
export function lookupModelPricing(modelString: string): ModelPricingLookup | null {
115123
const data = loadPricingData();
116124
const candidates = normalizeModelForLookup(modelString);
117125

@@ -123,16 +131,19 @@ export function getModelPricing(modelString: string): ModelPricing | null {
123131

124132
if (typeof inputCost === "number" || typeof outputCost === "number") {
125133
return {
126-
inputCostPerToken: typeof inputCost === "number" ? inputCost : 0,
127-
outputCostPerToken: typeof outputCost === "number" ? outputCost : 0,
128-
cacheReadInputTokenCost:
129-
typeof modelData.cache_read_input_token_cost === "number"
130-
? modelData.cache_read_input_token_cost
131-
: undefined,
132-
cacheCreationInputTokenCost:
133-
typeof modelData.cache_creation_input_token_cost === "number"
134-
? modelData.cache_creation_input_token_cost
135-
: undefined,
134+
pricing: {
135+
inputCostPerToken: typeof inputCost === "number" ? inputCost : 0,
136+
outputCostPerToken: typeof outputCost === "number" ? outputCost : 0,
137+
cacheReadInputTokenCost:
138+
typeof modelData.cache_read_input_token_cost === "number"
139+
? modelData.cache_read_input_token_cost
140+
: undefined,
141+
cacheCreationInputTokenCost:
142+
typeof modelData.cache_creation_input_token_cost === "number"
143+
? modelData.cache_creation_input_token_cost
144+
: undefined,
145+
},
146+
matchedKey: candidate,
136147
};
137148
}
138149
}
@@ -141,15 +152,57 @@ export function getModelPricing(modelString: string): ModelPricing | null {
141152
return null;
142153
}
143154

155+
/**
156+
* Get pricing information for a model using an explicit key
157+
* Returns null if pricing is not found
158+
*/
159+
export function lookupModelPricingByKey(pricingKey: string): ModelPricingLookup | null {
160+
const data = loadPricingData();
161+
const modelData = data[pricingKey] as Record<string, unknown> | undefined;
162+
163+
if (!modelData) {
164+
return null;
165+
}
166+
167+
const inputCost = modelData.input_cost_per_token;
168+
const outputCost = modelData.output_cost_per_token;
169+
170+
if (typeof inputCost === "number" || typeof outputCost === "number") {
171+
return {
172+
pricing: {
173+
inputCostPerToken: typeof inputCost === "number" ? inputCost : 0,
174+
outputCostPerToken: typeof outputCost === "number" ? outputCost : 0,
175+
cacheReadInputTokenCost:
176+
typeof modelData.cache_read_input_token_cost === "number"
177+
? modelData.cache_read_input_token_cost
178+
: undefined,
179+
cacheCreationInputTokenCost:
180+
typeof modelData.cache_creation_input_token_cost === "number"
181+
? modelData.cache_creation_input_token_cost
182+
: undefined,
183+
},
184+
matchedKey: pricingKey,
185+
};
186+
}
187+
188+
return null;
189+
}
190+
191+
/**
192+
* Get pricing information for a model (legacy function for compatibility)
193+
* Returns null if pricing is not found
194+
*/
195+
export function getModelPricing(modelString: string): ModelPricing | null {
196+
const lookup = lookupModelPricing(modelString);
197+
return lookup?.pricing ?? null;
198+
}
199+
144200
/**
145201
* Get pricing display information (cost per million tokens)
146202
*/
147203
export function getModelPricingDisplay(
148-
modelString: string,
149-
): ModelPricingDisplay | null {
150-
const pricing = getModelPricing(modelString);
151-
if (!pricing) return null;
152-
204+
pricing: ModelPricing,
205+
): ModelPricingDisplay {
153206
return {
154207
inputCostPerMTok: pricing.inputCostPerToken * 1_000_000,
155208
outputCostPerMTok: pricing.outputCostPerToken * 1_000_000,
@@ -163,14 +216,11 @@ export function getModelPricingDisplay(
163216
* Calculate the cost for given token usage
164217
*/
165218
export function calculateCost(
166-
modelString: string,
219+
pricing: ModelPricing,
167220
inputTokens: number,
168221
outputTokens: number,
169222
cachedInputTokens: number = 0,
170-
): CostCalculation | null {
171-
const pricing = getModelPricing(modelString);
172-
if (!pricing) return null;
173-
223+
): CostCalculation {
174224
// For cached tokens, we subtract them from input tokens for billing purposes
175225
// and bill them at the cache read rate (if available, otherwise free)
176226
const uncachedInputTokens = inputTokens - cachedInputTokens;
@@ -222,3 +272,11 @@ export function formatMTokCost(costPerMTok: number): string {
222272
export function isPricingAvailable(): boolean {
223273
return existsSync(PRICING_FILE_PATH);
224274
}
275+
276+
/**
277+
* Get all available model keys (for debugging/listing)
278+
*/
279+
export function getAllModelKeys(): string[] {
280+
const data = loadPricingData();
281+
return Object.keys(data).filter((key) => key !== "sample_spec");
282+
}

0 commit comments

Comments
 (0)