platform/plugins/openai/src/utils.ts
2024-08-28 18:36:11 +07:00

42 lines
1.5 KiB
TypeScript

//
// Copyright © 2024 Hardcore Engineering Inc.
//
// Licensed under the Eclipse Public License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License. You may
// obtain a copy of the License at https://www.eclipse.org/legal/epl-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//
// See the License for the specific language governing permissions and
// limitations under the License.
//
import OpenAI from 'openai'
import { Tiktoken } from 'tiktoken'
// Based on https://github.com/openai/openai-cookbook/blob/main/examples/How_to_count_tokens_with_tiktoken.ipynb
export function countTokens (messages: OpenAI.ChatCompletionMessageParam[], encoding: Tiktoken): number {
const tokensPerMessage = 3 // every message follows <|start|>{role/name}\n{content}<|end|>\n
const tokensPerName = 1 // every name follows <|name|>{name}<|end|>\n
let result = 0
for (const message of messages) {
result += tokensPerMessage
for (const key in message) {
const value = message[key as keyof OpenAI.ChatCompletionMessageParam] as string
if (value == null) continue
result += encoding.encode(value).length
if (key === 'name') {
result += tokensPerName
}
}
}
result += 3 // every reply is primed with <|start|>assistant<|message|>
return result
}