280 lines
6.3 KiB
Go
280 lines
6.3 KiB
Go
package openai
|
|
|
|
import (
|
|
"encoding/json"
|
|
"strconv"
|
|
"strings"
|
|
"unicode/utf8"
|
|
)
|
|
|
|
// estimateInputTokens returns a conservative approximation of the token count
|
|
// for a request payload. The formula is based on the plan:
|
|
//
|
|
// runes/4 + runes/16
|
|
//
|
|
// where `runes` is the total character count of all serialisable input
|
|
// elements (messages, tools, metadata) and structural_overhead is a
|
|
// proportional structural overhead.
|
|
func estimateInputTokens(input string, metadata map[string]string, tools []any, toolChoice any) int {
|
|
est := estimateFromStrings(input, metadata, tools, toolChoice)
|
|
// Minimum 1 token to avoid division-by-zero issues downstream.
|
|
if est < 1 {
|
|
return 1
|
|
}
|
|
return est
|
|
}
|
|
|
|
// estimateFromStrings produces a conservative token estimate by summing
|
|
// string-length proxies plus a small structural overhead.
|
|
func estimateFromStrings(input string, metadata map[string]string, tools []any, toolChoice any) int {
|
|
// Count runes in the main text payload.
|
|
totalRunes := 0
|
|
|
|
// input is the concatenated prompt/chat text.
|
|
totalRunes += utf8.RuneCountInString(input)
|
|
|
|
// metadata values contribute tokens proportional to their size.
|
|
for _, v := range metadata {
|
|
totalRunes += utf8.RuneCountInString(v)
|
|
}
|
|
|
|
// tool schemas are typically JSON-heavy; each tool entry adds a base
|
|
// overhead plus the size of its name/description.
|
|
for _, t := range tools {
|
|
totalRunes += 16 // base structural overhead per tool entry
|
|
if fn, ok := t.(map[string]any); ok {
|
|
if name, ok := fn["name"].(string); ok {
|
|
totalRunes += utf8.RuneCountInString(name)
|
|
}
|
|
if parameters, ok := fn["parameters"].(map[string]any); ok {
|
|
totalRunes += utf8.RuneCountInString(toolsSchemaToJSON(parameters))
|
|
}
|
|
}
|
|
if choice, ok := toolChoice.(map[string]any); ok {
|
|
if fn, ok := choice["function"].(map[string]any); ok {
|
|
if name, ok := fn["name"].(string); ok {
|
|
totalRunes += utf8.RuneCountInString(name)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// Conservative token estimate: each ~4 runes is roughly 1 token.
|
|
tokens := totalRunes / 4
|
|
// Structural overhead: add 1 token per ~16 runes to account for
|
|
// system tokens, delimiter tokens, message boundary tokens, etc.
|
|
tokens += totalRunes / 16
|
|
return tokens
|
|
}
|
|
|
|
// toolsSchemaToJSON is a lightweight JSON stringifier for the tools[].parameters
|
|
// object that avoids an import cycle with encoding/json.
|
|
func toolsSchemaToJSON(v any) string {
|
|
if m, ok := v.(map[string]any); ok {
|
|
var b stringsBuilder
|
|
b.WriteByte('{')
|
|
first := true
|
|
for k, val := range m {
|
|
if !first {
|
|
b.WriteByte(',')
|
|
}
|
|
first = false
|
|
encodeJSONString(&b, k)
|
|
b.WriteByte(':')
|
|
encodeJSONValue(&b, val)
|
|
}
|
|
b.WriteByte('}')
|
|
return b.String()
|
|
}
|
|
return "{}"
|
|
}
|
|
|
|
// stringsBuilder is a minimal []byte-based string builder to avoid import cycles.
|
|
type stringsBuilder struct {
|
|
b []byte
|
|
}
|
|
|
|
func (s *stringsBuilder) WriteByte(b byte) {
|
|
s.b = append(s.b, b)
|
|
}
|
|
|
|
func (s *stringsBuilder) WriteString(str string) {
|
|
s.b = append(s.b, str...)
|
|
}
|
|
|
|
func (s *stringsBuilder) WriteRune(r rune) {
|
|
buf := make([]byte, utf8.UTFMax)
|
|
n := utf8.EncodeRune(buf, r)
|
|
s.b = append(s.b, buf[:n]...)
|
|
}
|
|
|
|
func (s *stringsBuilder) Len() int {
|
|
return len(s.b)
|
|
}
|
|
|
|
func (s *stringsBuilder) String() string {
|
|
return string(s.b)
|
|
}
|
|
|
|
func encodeJSONString(b *stringsBuilder, v string) {
|
|
b.WriteByte('"')
|
|
for i := 0; i < len(v); i++ {
|
|
c := v[i]
|
|
switch c {
|
|
case '"', '\\':
|
|
b.WriteByte('\\')
|
|
b.WriteByte(c)
|
|
default:
|
|
b.WriteByte(c)
|
|
}
|
|
}
|
|
b.WriteByte('"')
|
|
}
|
|
|
|
func encodeJSONValue(b *stringsBuilder, v any) {
|
|
switch val := v.(type) {
|
|
case string:
|
|
encodeJSONString(b, val)
|
|
case bool:
|
|
if val {
|
|
b.WriteString("true")
|
|
} else {
|
|
b.WriteString("false")
|
|
}
|
|
case float64:
|
|
b.WriteString(ftoa(val))
|
|
case int:
|
|
b.WriteString(itoa(val))
|
|
case int64:
|
|
b.WriteString(itoa64(val))
|
|
case map[string]any:
|
|
encodeJSONObject(b, val)
|
|
case []any:
|
|
encodeJSONArray(b, val)
|
|
case nil:
|
|
b.WriteString("null")
|
|
default:
|
|
b.WriteString("null")
|
|
}
|
|
}
|
|
|
|
func encodeJSONObject(b *stringsBuilder, m map[string]any) {
|
|
b.WriteByte('{')
|
|
first := true
|
|
for k, v := range m {
|
|
if !first {
|
|
b.WriteByte(',')
|
|
}
|
|
first = false
|
|
encodeJSONString(b, k)
|
|
b.WriteByte(':')
|
|
encodeJSONValue(b, v)
|
|
}
|
|
b.WriteByte('}')
|
|
}
|
|
|
|
func encodeJSONArray(b *stringsBuilder, items []any) {
|
|
b.WriteByte('[')
|
|
for i, item := range items {
|
|
if i > 0 {
|
|
b.WriteByte(',')
|
|
}
|
|
encodeJSONValue(b, item)
|
|
}
|
|
b.WriteByte(']')
|
|
}
|
|
|
|
func ftoa(v float64) string {
|
|
return strconv.FormatFloat(v, 'f', -1, 64)
|
|
}
|
|
|
|
// lightweight integer/float formatters to avoid fmt import.
|
|
func itoa(v int) string {
|
|
if v == 0 {
|
|
return "0"
|
|
}
|
|
negative := v < 0
|
|
if negative {
|
|
v = -v
|
|
}
|
|
var buf [20]byte
|
|
i := len(buf)
|
|
for v > 0 {
|
|
i--
|
|
buf[i] = byte('0' + v%10)
|
|
v /= 10
|
|
}
|
|
if negative {
|
|
i--
|
|
buf[i] = '-'
|
|
}
|
|
return string(buf[i:])
|
|
}
|
|
|
|
func itoa64(v int64) string {
|
|
if v == 0 {
|
|
return "0"
|
|
}
|
|
negative := v < 0
|
|
if negative {
|
|
v = -v
|
|
}
|
|
var buf [20]byte
|
|
i := len(buf)
|
|
for v > 0 {
|
|
i--
|
|
buf[i] = byte('0' + v%10)
|
|
v /= 10
|
|
}
|
|
if negative {
|
|
i--
|
|
buf[i] = '-'
|
|
}
|
|
return string(buf[i:])
|
|
}
|
|
|
|
func zeroPadLeft(pad int, s string) string {
|
|
for len(s) < pad {
|
|
s = "0" + s
|
|
}
|
|
return s
|
|
}
|
|
|
|
// classifyContext returns "normal" or "long" based on the token estimate
|
|
// and the server's threshold.
|
|
func classifyContext(tokens int, threshold int) string {
|
|
if tokens >= threshold {
|
|
return "long"
|
|
}
|
|
return "normal"
|
|
}
|
|
|
|
// estimateChatInputTokens computes a conservative token-count approximation for
|
|
// a chat completion request. The prompt argument is the already-computed
|
|
// payload (including any strict-output contract instruction) so the estimate
|
|
// does not double-count the contract text that is also sent to the node.
|
|
func (s *Server) estimateChatInputTokens(prompt string, runMeta map[string]string, tools []any, toolChoice any) int {
|
|
textParts := []string{prompt}
|
|
// Metadata keys+values.
|
|
for k, v := range runMeta {
|
|
textParts = append(textParts, k, v)
|
|
}
|
|
// Tool schemas.
|
|
for _, t := range tools {
|
|
textParts = append(textParts, toJSON(t))
|
|
}
|
|
if toolChoice != nil {
|
|
textParts = append(textParts, toJSON(toolChoice))
|
|
}
|
|
input := strings.Join(textParts, "\n")
|
|
return estimateInputTokens(input, runMeta, tools, toolChoice)
|
|
}
|
|
|
|
// toJSON is a lightweight serialiser for any json.RawMessage-compatible value.
|
|
func toJSON(v any) string {
|
|
b, err := json.Marshal(v)
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
return string(b)
|
|
}
|