loveuer
44 lines
1018 B
Go
44 lines
1018 B
Go
package llm
|
|
|
|
import "encoding/json"
|
|
|
|
const tokenEstimateCharRatio = 4
|
|
|
|
// EstimateTokens estimates the token cost of messages and tool definitions in a
|
|
// chat completion request. It intentionally uses a simple JSON-size heuristic:
|
|
// roughly four UTF-8 bytes per token plus a small per-message framing overhead.
|
|
func EstimateTokens(messages []Message, tools []Tool) int {
|
|
if len(messages) == 0 && len(tools) == 0 {
|
|
return 0
|
|
}
|
|
|
|
payload := struct {
|
|
Messages []Message `json:"messages,omitempty"`
|
|
Tools []Tool `json:"tools,omitempty"`
|
|
}{
|
|
Messages: messages,
|
|
Tools: tools,
|
|
}
|
|
data, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return 0
|
|
}
|
|
|
|
return estimateBytesAsTokens(len(data)) + len(messages)*4
|
|
}
|
|
|
|
func estimateMessageTokens(msg Message) int {
|
|
data, err := json.Marshal(msg)
|
|
if err != nil {
|
|
return 0
|
|
}
|
|
return estimateBytesAsTokens(len(data)) + 4
|
|
}
|
|
|
|
func estimateBytesAsTokens(n int) int {
|
|
if n <= 0 {
|
|
return 0
|
|
}
|
|
return (n + tokenEstimateCharRatio - 1) / tokenEstimateCharRatio
|
|
}
|