feat: add function calling token counting (#83)

- add a function-calling aware token counting API to `GptEncoding`
- introduce shared fixtures and a regression test suite that covers OpenAI function request shapes
- expose `countChatCompletionTokens` only on function-calling model entrypoints via the code generator
- extract the function calling token counting helpers and fixtures into dedicated modules with named token constants

Author: niieani

README.md

@@ -272,6 +272,51 @@ const text = 'Hello, world!'
 const tokenCount = countTokens(text)
 ```
 
+### `countChatCompletionTokens(request: ChatCompletionRequest): number`
+
+Counts the tokens that a function-calling chat completion request will consume, including message overhead, optional function definitions, and pinned function calls. This helper is only available on models that support the `function_calling` feature.
+
+Example:
+
+```typescript
+import {
+  countChatCompletionTokens,
+  type ChatCompletionRequest,
+} from 'gpt-tokenizer/model/gpt-4o'
+
+const request: ChatCompletionRequest = {
+  messages: [
+    { role: 'system', content: 'You are a helpful assistant.' },
+    { role: 'user', content: 'Find the weather for San Francisco.' },
+  ],
+  functions: [
+    {
+      name: 'get_weather',
+      description: 'Look up the weather for a city.',
+      parameters: {
+        type: 'object',
+        required: ['city'],
+        properties: {
+          city: { type: 'string' },
+          unit: { type: 'string', enum: ['celsius', 'fahrenheit'] },
+        },
+      },
+    },
+  ],
+}
+
+const promptTokenEstimate = countChatCompletionTokens(request)
+```
+
+You can also access the helper from the module's default export:
+
+```typescript
+import gpt4o from 'gpt-tokenizer/model/gpt-4o'
+
+// Reuse the `request` defined above
+const tokenCount = gpt4o.countChatCompletionTokens?.(request)
+```
+
 ### `encodeChat(chat: ChatMessage[], model?: ModelName, encodeOptions?: EncodeOptions): number[]`
 
 Encodes the given chat into a sequence of tokens. The optional `encodeOptions` parameter lets you configure special token handling.

src/GptEncoding.test.ts

@@ -18,6 +18,24 @@ import * as modelsMap from './modelsMap.js'
 import { resolveEncoding } from './resolveEncoding.js'
 import { EndOfText } from './specialTokens.js'
 
+describe('generated model exports', () => {
+  test('gpt-5 re-exports the chat token counter helper', async () => {
+    const mod = await import('./model/gpt-5.js')
+    const encoding = mod.default
+
+    expect('countChatCompletionTokens' in mod).toBe(true)
+    expect(mod.countChatCompletionTokens).toBe(
+      encoding.countChatCompletionTokens,
+    )
+  })
+
+  test('gpt-3.5-turbo-0613 omits the chat token counter helper', async () => {
+    const mod = await import('./model/gpt-3.5-turbo-0613.js')
+
+    expect('countChatCompletionTokens' in mod).toBe(false)
+  })
+})
+
 const sharedResults = {
   space: [220],
   tab: [197],

src/GptEncoding.ts

@@ -2,6 +2,13 @@
 /* eslint-disable no-param-reassign */
 import { BytePairEncodingCore, decoder } from './BytePairEncodingCore.js'
 import { ALL_SPECIAL_TOKENS } from './constants.js'
+import {
+  type ChatCompletionRequest,
+  type ChatMessage,
+  type EncodeChatOptions,
+  type HarmonyTerminator,
+  computeChatCompletionTokenCount,
+} from './functionCalling.js'
 import {
   type ChatModelName,
   type ChatParameters,
@@ -38,6 +45,25 @@ import {
 import { endsWithIncompleteUtfPairSurrogate } from './utfUtil.js'
 import { getMaxValueFromMap, getSpecialTokenRegex } from './util.js'
 
+export type {
+  ChatCompletionArrayProperty,
+  ChatCompletionBooleanProperty,
+  ChatCompletionFunctionCallOption,
+  ChatCompletionFunctionDefinition,
+  ChatCompletionFunctionParameters,
+  ChatCompletionFunctionProperty,
+  ChatCompletionFunctionType,
+  ChatCompletionNullProperty,
+  ChatCompletionNumberProperty,
+  ChatCompletionObjectProperty,
+  ChatCompletionRequest,
+  ChatCompletionStringProperty,
+  ChatMessage,
+  ChatMessageFunctionCall,
+  EncodeChatOptions,
+  HarmonyTerminator,
+} from './functionCalling.js'
+
 export interface CostEstimate {
   input?: number
   output?: number
@@ -66,28 +92,6 @@ export interface EncodeOptions {
   disallowedSpecial?: Set<string> | typeof ALL_SPECIAL_TOKENS
 }
 
-export type HarmonyTerminator = '<|end|>' | '<|return|>' | '<|call|>'
-
-export interface ChatMessage {
-  role?: 'system' | 'user' | 'assistant' | 'developer' | (string & {})
-  name?: string
-  content: string
-  /** Harmony-only: channel label such as `analysis`, `commentary`, or `final`. */
-  channel?: string
-  /** Harmony-only: recipient metadata, e.g. `functions.get_weather` or `assistant`. */
-  recipient?: string
-  /** Controls where the recipient metadata is rendered in Harmony headers. Defaults to `channel`. */
-  recipientPlacement?: 'role' | 'channel'
-  /** Harmony-only: constraint label, e.g. `json`. */
-  constraint?: string
-  /** Harmony-only: overrides the closing token, defaults to `<|end|>`. */
-  terminator?: HarmonyTerminator
-}
-
-export interface EncodeChatOptions {
-  primeWithAssistantResponse?: string
-}
-
 interface SpecialTokenConfig {
   allowedSpecial: Set<string> | undefined
   regexPattern: RegExp | undefined
@@ -109,6 +113,8 @@ export class GptEncoding {
   private defaultSpecialTokenConfig: SpecialTokenConfig
   private chatFormatter: ChatFormatter
 
+  countChatCompletionTokens?: (request: ChatCompletionRequest) => number
+
   readonly vocabularySize: number
 
   private constructor({
@@ -169,6 +175,10 @@ export class GptEncoding {
     this.setMergeCacheSize = this.setMergeCacheSize.bind(this)
     this.clearMergeCache = this.clearMergeCache.bind(this)
     this.estimateCost = this.estimateCost.bind(this)
+    if (modelSpec?.supported_features?.includes('function_calling')) {
+      this.countChatCompletionTokens =
+        this.countChatCompletionTokensInternal.bind(this)
+    }
     this.modelName = modelName
     this.modelSpec = modelSpec
     this.chatFormatter = chatFormatter ?? 'chatml'
@@ -522,6 +532,22 @@ export class GptEncoding {
     return count
   }
 
+  private countStringTokens(text: string): number {
+    if (!text) {
+      return 0
+    }
+
+    return this.bytePairEncodingCoreProcessor.countNative(text)
+  }
+
+  private countChatCompletionTokensInternal(
+    request: ChatCompletionRequest,
+  ): number {
+    return computeChatCompletionTokenCount(request, (text) =>
+      this.countStringTokens(text),
+    )
+  }
+
   setMergeCacheSize(size: number): void {
     this.bytePairEncodingCoreProcessor.setMergeCacheSize(size)
   }

src/codegen/generateByModel.ts

@@ -10,6 +10,7 @@ import {
   DEFAULT_ENCODING,
   modelToEncodingMap,
 } from '../mapping.js'
+import type { Feature } from '../modelTypes.js'
 import * as models from '../models.js'
 
 // eslint-disable-next-line no-underscore-dangle
@@ -71,7 +72,15 @@ await Promise.all(
       '',
     ]
 
-    const baseContent = isChatModel
+    const supportedFeatures = (
+      modelData as { supported_features?: readonly Feature[] }
+    ).supported_features
+
+    const supportsFunctionCalling =
+      Array.isArray(supportedFeatures) &&
+      supportedFeatures.includes('function_calling')
+
+    let baseContent = isChatModel
       ? template
           .replace(
             `getEncodingApi('cl100k_base', () => bpeRanks)`,
@@ -90,6 +99,14 @@ export { default } from '../encoding/${encoding}.js'
 export * from '../encoding/${encoding}.js'
 `
 
+    if (isChatModel && supportsFunctionCalling) {
+      const snippet = '  encodeChat,\n  encodeChatGenerator,\n'
+      const replacement =
+        '  encodeChat,\n  countChatCompletionTokens,\n  encodeChatGenerator,\n'
+      baseContent = baseContent.replace(snippet, replacement)
+      baseContent = baseContent.replace(snippet, replacement)
+    }
+
     const content = insertHeaderAfterDirectives(baseContent, headerLines)
     await fs.writeFile(
       path.join(__dirname, `../model/${modelName}.ts`),