agentu/pkg/llm/token_estimate.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
}