gpt-tokenizer 2.8.0 → 2.9.0

package/src/GptEncoding.test.ts

@@ -1,5 +1,7 @@
 import fs from 'fs'
 import path from 'path'
+// eslint-disable-next-line import/no-extraneous-dependencies
+import { describe, expect, test } from 'vitest'
 import { ALL_SPECIAL_TOKENS } from './constants.js'
 import { type ChatMessage, GptEncoding } from './GptEncoding.js'
 import {
@@ -8,7 +10,9 @@ import {
   chatModelParams,
   encodingNames,
 } from './mapping.js'
+import { models } from './models.js'
 import { resolveEncoding } from './resolveEncoding.js'
+import { EndOfText } from './specialTokens.js'
 
 const sharedResults = {
   space: [220],
@@ -66,8 +70,8 @@ const offsetPrompts = [
   // Basic prompt with "hello world"
   'hello world',
 
-  // Basic prompt with special token "<|endoftext|>"
-  'hello world<|endoftext|> green cow',
+  // Basic prompt with special token end of text token
+  `hello world${EndOfText} green cow`,
 
   // Chinese text: "我非常渴望与人工智能一起工作"
   '我非常渴望与人工智能一起工作',
@@ -96,7 +100,7 @@ describe.each(encodingNames)('%s', (encodingName: EncodingName) => {
   } = encoding
 
   describe('encode and decode', () => {
-    it.each(offsetPrompts)('offset prompt: %s', (str) => {
+    test.each(offsetPrompts)('offset prompt: %s', (str) => {
       expect(
         decode(encode(str, { allowedSpecial: ALL_SPECIAL_TOKENS })),
       ).toEqual(str)
@@ -105,7 +109,7 @@ describe.each(encodingNames)('%s', (encodingName: EncodingName) => {
 
   describe('basic functionality', () => {
     const result = results[encodingName]
-    it('empty string', () => {
+    test('empty string', () => {
       const str = ''
       expect(encode(str)).toEqual([])
       expect(decode(encode(str))).toEqual(str)
@@ -113,7 +117,7 @@ describe.each(encodingNames)('%s', (encodingName: EncodingName) => {
       expect(isWithinTokenLimit(str, 3)).toBe(0)
     })
 
-    it('space', () => {
+    test('space', () => {
       const str = ' '
       expect(encode(str)).toEqual(result.space)
       expect(decode(encode(str))).toEqual(str)
@@ -121,13 +125,13 @@ describe.each(encodingNames)('%s', (encodingName: EncodingName) => {
       expect(isWithinTokenLimit(str, 0)).toBe(false)
     })
 
-    it('tab', () => {
+    test('tab', () => {
       const str = '\t'
       expect(encode(str)).toEqual(result.tab)
       expect(decode(encode(str))).toEqual(str)
     })
 
-    it('simple text', () => {
+    test('simple text', () => {
       const str = 'This is some text'
       expect(encode(str)).toEqual(result[str])
       expect(decode(encode(str))).toEqual(str)
@@ -135,14 +139,14 @@ describe.each(encodingNames)('%s', (encodingName: EncodingName) => {
       expect(isWithinTokenLimit(str, 5)).toBe(result[str].length)
     })
 
-    it('multi-token word', () => {
+    test('multi-token word', () => {
       const str = 'indivisible'
       expect(encode(str)).toEqual(result.indivisible)
       expect(decode(encode(str))).toEqual(str)
       expect(isWithinTokenLimit(str, 3)).toBe(result.indivisible.length)
     })
 
-    it('emojis', () => {
+    test('emojis', () => {
       const str = 'hello 👋 world 🌍'
       expect(encode(str)).toEqual(result[str])
       expect(decode(encode(str))).toEqual(str)
@@ -150,7 +154,7 @@ describe.each(encodingNames)('%s', (encodingName: EncodingName) => {
       expect(isWithinTokenLimit(str, 400)).toBe(result[str].length)
     })
 
-    it('decode token-by-token via generator', () => {
+    test('decode token-by-token via generator', () => {
       const str = 'hello 👋 world 🌍'
       const generator = decodeGenerator(result[str])
       result.decodedHelloWorldTokens.forEach((token: string) => {
@@ -158,8 +162,8 @@ describe.each(encodingNames)('%s', (encodingName: EncodingName) => {
       })
     })
 
-    it('encodes and decodes special tokens', () => {
-      const str = 'hello <|endoftext|> world'
+    test('encodes and decodes special tokens', () => {
+      const str = `hello ${EndOfText} world`
       const encoded = encode(str, {
         allowedSpecial: ALL_SPECIAL_TOKENS,
       })
@@ -174,7 +178,7 @@ describe.each(encodingNames)('%s', (encodingName: EncodingName) => {
       }
     }
 
-    it('decode token-by-token via async generator', async () => {
+    test('decode token-by-token via async generator', async () => {
       const generator = decodeAsyncGenerator(getHelloWorldTokensAsync())
       const decoded = [...result.decodedHelloWorldTokens]
       for await (const value of generator) {
@@ -182,14 +186,14 @@ describe.each(encodingNames)('%s', (encodingName: EncodingName) => {
       }
     })
 
-    it('properties of Object', () => {
+    test('properties of Object', () => {
       const str = 'toString constructor hasOwnProperty valueOf'
 
       expect(encode(str)).toEqual(result[str])
       expect(decode(encode(str))).toEqual(str)
     })
 
-    it('text with commas', () => {
+    test('text with commas', () => {
       const str = 'hello, I am a text, and I have commas. a,b,c'
       expect(decode(encode(str))).toEqual(str)
       expect(encode(str)).toStrictEqual(result[str])
@@ -200,10 +204,10 @@ describe.each(encodingNames)('%s', (encodingName: EncodingName) => {
 
   describe('test plan', () => {
     testPlans[encodingName].forEach(({ sample, encoded }) => {
-      it(`encodes ${sample}`, () => {
+      test(`encodes ${sample}`, () => {
         expect(encode(sample)).toEqual(encoded)
       })
-      it(`decodes ${sample}`, () => {
+      test(`decodes ${sample}`, () => {
         expect(decode(encoded)).toEqual(sample)
       })
     })
@@ -258,7 +262,7 @@ describe.each(chatModelNames)('%s', (modelName) => {
     : 121
 
   describe('chat functionality', () => {
-    it('encodes a chat correctly', () => {
+    test('encodes a chat correctly', () => {
       const encoded = encoding.encodeChat(exampleMessages)
       expect(encoded).toMatchSnapshot()
       expect(encoded).toHaveLength(expectedEncodedLength)
@@ -267,14 +271,14 @@ describe.each(chatModelNames)('%s', (modelName) => {
       expect(decoded).toMatchSnapshot()
     })
 
-    it('isWithinTokenLimit: false', () => {
+    test('isWithinTokenLimit: false', () => {
       const isWithinTokenLimit = encoding.isWithinTokenLimit(
         exampleMessages,
         50,
       )
       expect(isWithinTokenLimit).toBe(false)
     })
-    it('isWithinTokenLimit: true (number)', () => {
+    test('isWithinTokenLimit: true (number)', () => {
       const isWithinTokenLimit = encoding.isWithinTokenLimit(
         exampleMessages,
         150,
@@ -284,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')

package/src/GptEncoding.ts

@@ -15,6 +15,7 @@ import {
   type GetMergeableRanksFn,
   getEncodingParams,
 } from './modelParams.js'
+import { type CostEstimate, models } from './models.js'
 import {
   EndOfPrompt,
   EndOfText,
@@ -128,6 +129,8 @@ export class GptEncoding {
     this.encodeChatGenerator = this.encodeChatGenerator.bind(this)
     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
   }
 
@@ -366,6 +369,10 @@ export class GptEncoding {
     this.bytePairEncodingCoreProcessor.setMergeCacheSize(size)
   }
 
+  clearMergeCache(): void {
+    this.bytePairEncodingCoreProcessor.clearMergeCache()
+  }
+
   decode(inputTokensToDecode: Iterable<number>): string {
     return this.bytePairEncodingCoreProcessor.decodeNative(inputTokensToDecode)
   }
@@ -455,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
+  }
 }