feat: implement 'estimateCost'

Author: niieani

src/GptEncoding.test.ts

@@ -10,6 +10,7 @@ import {
   chatModelParams,
   encodingNames,
 } from './mapping.js'
+import { models } from './models.js'
 import { resolveEncoding } from './resolveEncoding.js'
 import { EndOfText } from './specialTokens.js'
 
@@ -287,6 +288,84 @@ describe.each(chatModelNames)('%s', (modelName) => {
   })
 })
 
+describe('estimateCost functionality', () => {
+  const gpt4oEncoding = GptEncoding.getEncodingApiForModel(
+    'gpt-4o',
+    resolveEncoding,
+  )
+  const gpt35Encoding = GptEncoding.getEncodingApiForModel(
+    'gpt-3.5-turbo',
+    resolveEncoding,
+  )
+
+  test('estimates cost correctly for gpt-4o model', () => {
+    const tokenCount = 1_000
+    const cost = gpt4oEncoding.estimateCost(tokenCount)
+
+    // gpt-4o has $2.5 per million tokens for input and $10 per million tokens for output
+    expect(cost.input).toBeCloseTo(0.002_5, 6) // 1000/1M * $2.5
+    expect(cost.output).toBeCloseTo(0.01, 6) // 1000/1M * $10
+    expect(cost.batchInput).toBeCloseTo(0.001_25, 6) // 1000/1M * $1.25
+    expect(cost.batchOutput).toBeCloseTo(0.005, 6) // 1000/1M * $5
+  })
+
+  test('estimates cost correctly for gpt-3.5-turbo model', () => {
+    const tokenCount = 1_000
+    const cost = gpt35Encoding.estimateCost(tokenCount)
+
+    // gpt-3.5-turbo has $0.5 per million tokens for input and $1.5 per million tokens for output
+    expect(cost.input).toBeCloseTo(0.000_5, 6) // 1000/1M * $0.5
+    expect(cost.output).toBeCloseTo(0.001_5, 6) // 1000/1M * $1.5
+    expect(cost.batchInput).toBeCloseTo(0.000_25, 6) // 1000/1M * $0.25
+    expect(cost.batchOutput).toBeCloseTo(0.000_75, 6) // 1000/1M * $0.75
+  })
+
+  test('allows overriding model name', () => {
+    const tokenCount = 1_000
+    // Use gpt-4o encoding but override with gpt-3.5-turbo model name
+    const cost = gpt4oEncoding.estimateCost(tokenCount, 'gpt-3.5-turbo')
+
+    expect(cost.input).toBeCloseTo(0.000_5, 6) // 1000/1M * $0.5
+    expect(cost.output).toBeCloseTo(0.001_5, 6) // 1000/1M * $1.5
+  })
+
+  test('throws error when model name is not provided', () => {
+    const encoding = GptEncoding.getEncodingApi('cl100k_base', resolveEncoding)
+    const tokenCount = 1_000
+
+    // No model name was provided during initialization or function call
+    expect(() => encoding.estimateCost(tokenCount)).toThrow(
+      'Model name must be provided either during initialization or passed in to the method.',
+    )
+  })
+
+  test('throws error for unknown model', () => {
+    const tokenCount = 1_000
+    expect(() =>
+      gpt4oEncoding.estimateCost(tokenCount, 'non-existent-model' as any),
+    ).toThrow('Unknown model: non-existent-model')
+  })
+
+  test('only includes properties that exist for the model', () => {
+    // Find a model that only has input cost but no output cost
+    const modelWithInputOnly = Object.entries(models).find(
+      ([_, model]) =>
+        model.cost?.input !== undefined && model.cost?.output === undefined,
+    )
+
+    if (modelWithInputOnly) {
+      const [modelName] = modelWithInputOnly
+      const cost = gpt4oEncoding.estimateCost(1_000, modelName as any)
+
+      expect(cost.input).toBeDefined()
+      expect(cost.output).toBeUndefined()
+    } else {
+      // Skip test if we can't find an appropriate model
+      console.log('Skipping test: no model with input-only cost found')
+    }
+  })
+})
+
 function loadTestPlans() {
   const testPlanPath = path.join(__dirname, '../data/TestPlans.txt')
   const testPlanData = fs.readFileSync(testPlanPath, 'utf8')

src/GptEncoding.ts

@@ -15,6 +15,7 @@ import {
   type GetMergeableRanksFn,
   getEncodingParams,
 } from './modelParams.js'
+import { type CostEstimate, models } from './models.js'
 import {
   EndOfPrompt,
   EndOfText,
@@ -129,6 +130,7 @@ export class GptEncoding {
     this.countTokens = this.countTokens.bind(this)
     this.setMergeCacheSize = this.setMergeCacheSize.bind(this)
     this.clearMergeCache = this.clearMergeCache.bind(this)
+    this.estimateCost = this.estimateCost.bind(this)
     this.modelName = modelName
   }
 
@@ -460,4 +462,53 @@ export class GptEncoding {
 
     return buffer
   }
+
+  /**
+   * Estimates the cost of processing a given token count using the model's pricing.
+   *
+   * @param tokenCount - The number of tokens to estimate cost for
+   * @param modelName - Optional model name to use for cost calculation (defaults to this.modelName)
+   * @returns Cost estimate object with applicable price components (input, output, batchInput, batchOutput)
+   */
+  estimateCost(tokenCount: number, modelName = this.modelName): CostEstimate {
+    if (!modelName) {
+      throw new Error(
+        'Model name must be provided either during initialization or passed in to the method.',
+      )
+    }
+
+    const model = models[modelName]
+    if (!model) {
+      throw new Error(`Unknown model: ${modelName}`)
+    }
+
+    if (!model.cost) {
+      throw new Error(`No cost information available for model: ${modelName}`)
+    }
+
+    const costPerMillion = model.cost
+    const result: CostEstimate = {}
+
+    // Calculate cost per token and multiply by token count
+    // eslint-disable-next-line no-magic-numbers
+    const millionTokens = tokenCount / 1_000_000
+
+    if (costPerMillion.input !== undefined) {
+      result.input = costPerMillion.input * millionTokens
+    }
+
+    if (costPerMillion.output !== undefined) {
+      result.output = costPerMillion.output * millionTokens
+    }
+
+    if (costPerMillion.batchInput !== undefined) {
+      result.batchInput = costPerMillion.batchInput * millionTokens
+    }
+
+    if (costPerMillion.batchOutput !== undefined) {
+      result.batchOutput = costPerMillion.batchOutput * millionTokens
+    }
+
+    return result
+  }
 }

src/encoding/cl100k_base.ts

@@ -19,6 +19,7 @@ const {
   vocabularySize,
   setMergeCacheSize,
   clearMergeCache,
+  estimateCost,
 } = api
 export {
   clearMergeCache,
@@ -30,6 +31,7 @@ export {
   encodeChat,
   encodeChatGenerator,
   encodeGenerator,
+  estimateCost,
   isWithinTokenLimit,
   setMergeCacheSize,
   vocabularySize,

src/encoding/o200k_base.ts

@@ -19,6 +19,7 @@ const {
   vocabularySize,
   setMergeCacheSize,
   clearMergeCache,
+  estimateCost,
 } = api
 export {
   clearMergeCache,
@@ -30,6 +31,7 @@ export {
   encodeChat,
   encodeChatGenerator,
   encodeGenerator,
+  estimateCost,
   isWithinTokenLimit,
   setMergeCacheSize,
   vocabularySize,

src/encoding/p50k_base.ts

@@ -17,6 +17,7 @@ const {
   vocabularySize,
   setMergeCacheSize,
   clearMergeCache,
+  estimateCost,
 } = api
 export {
   clearMergeCache,
@@ -26,6 +27,7 @@ export {
   decodeGenerator,
   encode,
   encodeGenerator,
+  estimateCost,
   isWithinTokenLimit,
   setMergeCacheSize,
   vocabularySize,

src/encoding/p50k_edit.ts

@@ -14,9 +14,12 @@ const {
   encodeGenerator,
   isWithinTokenLimit,
   countTokens,
+  encodeChat,
+  encodeChatGenerator,
   vocabularySize,
   setMergeCacheSize,
   clearMergeCache,
+  estimateCost,
 } = api
 export {
   clearMergeCache,
@@ -25,7 +28,10 @@ export {
   decodeAsyncGenerator,
   decodeGenerator,
   encode,
+  encodeChat,
+  encodeChatGenerator,
   encodeGenerator,
+  estimateCost,
   isWithinTokenLimit,
   setMergeCacheSize,
   vocabularySize,

src/encoding/r50k_base.ts

@@ -17,6 +17,7 @@ const {
   vocabularySize,
   setMergeCacheSize,
   clearMergeCache,
+  estimateCost,
 } = api
 export {
   clearMergeCache,
@@ -26,6 +27,7 @@ export {
   decodeGenerator,
   encode,
   encodeGenerator,
+  estimateCost,
   isWithinTokenLimit,
   setMergeCacheSize,
   vocabularySize,

src/model/gpt-3.5-turbo-0125.ts

@@ -18,8 +18,11 @@ const {
   encodeChatGenerator,
   vocabularySize,
   setMergeCacheSize,
+  clearMergeCache,
+  estimateCost,
 } = api
 export {
+  clearMergeCache,
   countTokens,
   decode,
   decodeAsyncGenerator,
@@ -28,6 +31,7 @@ export {
   encodeChat,
   encodeChatGenerator,
   encodeGenerator,
+  estimateCost,
   isWithinTokenLimit,
   setMergeCacheSize,
   vocabularySize,

src/model/gpt-3.5-turbo-0301.ts

@@ -18,8 +18,11 @@ const {
   encodeChatGenerator,
   vocabularySize,
   setMergeCacheSize,
+  clearMergeCache,
+  estimateCost,
 } = api
 export {
+  clearMergeCache,
   countTokens,
   decode,
   decodeAsyncGenerator,
@@ -28,6 +31,7 @@ export {
   encodeChat,
   encodeChatGenerator,
   encodeGenerator,
+  estimateCost,
   isWithinTokenLimit,
   setMergeCacheSize,
   vocabularySize,

src/model/gpt-3.5-turbo-0613.ts

@@ -18,8 +18,11 @@ const {
   encodeChatGenerator,
   vocabularySize,
   setMergeCacheSize,
+  clearMergeCache,
+  estimateCost,
 } = api
 export {
+  clearMergeCache,
   countTokens,
   decode,
   decodeAsyncGenerator,
@@ -28,6 +31,7 @@ export {
   encodeChat,
   encodeChatGenerator,
   encodeGenerator,
+  estimateCost,
   isWithinTokenLimit,
   setMergeCacheSize,
   vocabularySize,