1- import { ChatMessage } from 'gpt-tokenizer/GptEncoding' ;
21import { GenkitPluginV2 } from 'genkit/plugin' ;
32import { openAI } from '@genkit-ai/compat-oai/openai' ;
43import { RateLimiter } from 'limiter' ;
76 PromptDataForCounting ,
87 RateLimitConfig ,
98} from '../model-provider.js' ;
10- import o3 from 'gpt-tokenizer/model/o3' ;
11- import o4Mini from 'gpt-tokenizer/model/o4-mini' ;
9+ import { encoding_for_model } from 'tiktoken' ;
1210
1311export class OpenAiModelProvider extends GenkitModelProvider {
1412 readonly apiKeyVariableName = 'OPENAI_API_KEY' ;
@@ -30,8 +28,17 @@ export class OpenAiModelProvider extends GenkitModelProvider {
3028 tokensPerInterval : 30_000 * 0.75 , // *0.75 to be more resilient to token count deviations
3129 interval : 1000 * 60 * 1.5 , // Refresh tokens after 1.5 minutes to be on the safe side.
3230 } ) ,
33- countTokens : async ( prompt ) =>
34- o3 . countTokens ( this . genkitPromptToOpenAi ( prompt ) ) ,
31+ countTokens : async ( prompt ) => {
32+ const encoding = encoding_for_model ( 'gpt-4o' ) ;
33+ try {
34+ const messages = this . genkitPromptToOpenAi ( prompt ) ;
35+ const text = messages . map ( m => `${ m . role } : ${ m . content } ` ) . join ( '\n' ) ;
36+ const tokens = encoding . encode ( text ) ;
37+ return tokens . length ;
38+ } finally {
39+ encoding . free ( ) ;
40+ }
41+ } ,
3542 } ,
3643 // See https://platform.openai.com/docs/models/o4-mini
3744 'openai/o4-mini' : {
@@ -43,8 +50,17 @@ export class OpenAiModelProvider extends GenkitModelProvider {
4350 tokensPerInterval : 100_000 * 0.75 , // *0.75 to be more resilient to token count deviations
4451 interval : 1000 * 60 * 1.5 , // Refresh tokens after 1.5 minutes to be on the safe side.
4552 } ) ,
46- countTokens : async ( prompt ) =>
47- o4Mini . countTokens ( this . genkitPromptToOpenAi ( prompt ) ) ,
53+ countTokens : async ( prompt ) => {
54+ const encoding = encoding_for_model ( 'gpt-4o-mini' ) ;
55+ try {
56+ const messages = this . genkitPromptToOpenAi ( prompt ) ;
57+ const text = messages . map ( m => `${ m . role } : ${ m . content } ` ) . join ( '\n' ) ;
58+ const tokens = encoding . encode ( text ) ;
59+ return tokens . length ;
60+ } finally {
61+ encoding . free ( ) ;
62+ }
63+ } ,
4864 } ,
4965 // See: https://platform.openai.com/docs/models/gpt-5
5066 'openai/gpt-5' : {
@@ -56,10 +72,17 @@ export class OpenAiModelProvider extends GenkitModelProvider {
5672 tokensPerInterval : 30_000 * 0.75 , // *0.75 to be more resilient to token count deviations
5773 interval : 1000 * 60 * 1.5 , // Refresh tokens after 1.5 minutes to be on the safe side.
5874 } ) ,
59- // TODO: at the time of writing, the `gpt-tokenizer` doesn't support gpt-5.
60- // See https://github.com/niieani/gpt-tokenizer/issues/73
61- countTokens : async ( prompt ) =>
62- o3 . countTokens ( this . genkitPromptToOpenAi ( prompt ) ) ,
75+ countTokens : async ( prompt ) => {
76+ const encoding = encoding_for_model ( 'gpt-5' ) ;
77+ try {
78+ const messages = this . genkitPromptToOpenAi ( prompt ) ;
79+ const text = messages . map ( m => `${ m . role } : ${ m . content } ` ) . join ( '\n' ) ;
80+ const tokens = encoding . encode ( text ) ;
81+ return tokens . length ;
82+ } finally {
83+ encoding . free ( ) ;
84+ }
85+ } ,
6386 } ,
6487 } ;
6588
@@ -72,8 +95,8 @@ export class OpenAiModelProvider extends GenkitModelProvider {
7295 return { } ;
7396 }
7497
75- private genkitPromptToOpenAi ( prompt : PromptDataForCounting ) : ChatMessage [ ] {
76- const openAiPrompt : string | ChatMessage [ ] = [ ] ;
98+ private genkitPromptToOpenAi ( prompt : PromptDataForCounting ) : Array < { role : string ; content : string } > {
99+ const openAiPrompt : Array < { role : string ; content : string } > = [ ] ;
77100 for ( const part of prompt . messages ) {
78101 for ( const c of part . content ) {
79102 openAiPrompt . push ( {
@@ -82,6 +105,6 @@ export class OpenAiModelProvider extends GenkitModelProvider {
82105 } ) ;
83106 }
84107 }
85- return [ ...openAiPrompt , { content : prompt . prompt } ] ;
108+ return [ ...openAiPrompt , { role : 'user' , content : prompt . prompt } ] ;
86109 }
87110}
0 commit comments