@@ -14,17 +14,25 @@ let encoder: Tiktoken | undefined;
1414 * @param {Function[] } prompt.functions OpenAI function definitions
1515 * @returns An estimate for the number of tokens the prompt will use
1616 */
17- export function promptTokensEstimate ( { messages, functions } : { messages : Message [ ] , functions ?: Function [ ] } ) : number {
17+ export function promptTokensEstimate ( {
18+ messages,
19+ functions,
20+ } : {
21+ messages : Message [ ] ;
22+ functions ?: Function [ ] ;
23+ } ) : number {
1824 // It appears that if functions are present, the first system message is padded with a trailing newline. This
1925 // was inferred by trying lots of combinations of messages and functions and seeing what the token counts were.
2026 let paddedSystem = false ;
21- let tokens = messages . map ( m => {
22- if ( m . role === "system" && functions && ! paddedSystem ) {
23- m = { ...m , content : m . content + "\n" }
24- paddedSystem = true ;
25- }
26- return messageTokensEstimate ( m ) ;
27- } ) . reduce ( ( a , b ) => a + b , 0 ) ;
27+ let tokens = messages
28+ . map ( ( m ) => {
29+ if ( m . role === "system" && functions && ! paddedSystem ) {
30+ m = { ...m , content : m . content + "\n" } ;
31+ paddedSystem = true ;
32+ }
33+ return messageTokensEstimate ( m ) ;
34+ } )
35+ . reduce ( ( a , b ) => a + b , 0 ) ;
2836
2937 // Each completion (vs message) seems to carry a 3-token overhead
3038 tokens += 3 ;
@@ -37,7 +45,7 @@ export function promptTokensEstimate({ messages, functions }: { messages: Messag
3745 // If there's a system message _and_ functions are present, subtract four tokens. I assume this is because
3846 // functions typically add a system message, but reuse the first one if it's already there. This offsets
3947 // the extra 9 tokens added by the function definitions.
40- if ( functions && messages . find ( m => m . role === "system" ) ) {
48+ if ( functions && messages . find ( ( m ) => m . role === "system" ) ) {
4149 tokens -= 4 ;
4250 }
4351
@@ -68,7 +76,7 @@ export function messageTokensEstimate(message: Message): number {
6876 message . content ,
6977 message . name ,
7078 message . function_call ?. name ,
71- message . function_call ?. arguments
79+ message . function_call ?. arguments ,
7280 ] . filter ( ( v ) : v is string => ! ! v ) ;
7381 let tokens = components . map ( stringTokens ) . reduce ( ( a , b ) => a + b , 0 ) ;
7482 tokens += 3 ; // Add three per message
@@ -85,7 +93,7 @@ export function messageTokensEstimate(message: Message): number {
8593}
8694
8795/**
88- * Estimate the number of tokens a function definition will use. Note that using the function definition within
96+ * Estimate the number of tokens a function definition will use. Note that using the function definition within
8997 * a prompt will add extra tokens, so you might want to use `promptTokensEstimate` instead.
9098 * @param funcs An array of OpenAI function definitions
9199 * @returns An estimate for the number of tokens the function definitions will use
0 commit comments