luxferre/bantam
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

token accounting cleanup

Browse Source
This commit is contained in:
Luxferre
2026-08-18 10:17:52 +03:00
parent 4a936e850d
commit a34ee1cf3c
6 changed files with 251 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main.go
+71 -18
View File
@@ -19,6 +19,7 @@ import (
"path/filepath"
"regexp"
"sort"
"sync"
"strconv"
"strings"
"time"
@@ -49,6 +50,23 @@ type Cfg struct {
Raw map[string]string
}
// internalKey reports whether a model.cfg key is an agent-internal parameter
// that must never be forwarded to the chat completions API.
func internalKey(k string) bool {
switch k {
case "endpoint", "model", "temperature", "stream", "api_key", "timeout",
"shell_timeout", "max_al_iterations", "color", "context_window":
return true
}
return false
}
// llmTransport is a shared HTTP transport reused across all LLM calls so that
// connections are pooled instead of recreated per request.
var llmTransport = &http.Transport{
DialContext: (&net.Dialer{Timeout: 300 * time.Second}).DialContext,
}
var defCfg = Cfg{"https://opencode.ai/zen/v1", "big-pickle", "-", 0.7, 300, 120, 1000, true, "auto", 200000, nil}
func atoiD(s string, d int) int {
@@ -92,8 +110,24 @@ func queryModelsContextWindow(cfg *Cfg) int {
return 0
}
var cwCacheMu sync.Mutex
var cwCache = map[string]int{}
func fetchContextWindow(cfg *Cfg) int {
// Only values discovered from the /models endpoint are cached, keyed by
// endpoint+model. The context_window-override and 200000 default are derived
// per call from cfg so they never shadow each other across configs.
key := cfg.Endpoint + "\x00" + cfg.Model
cwCacheMu.Lock()
if cw, ok := cwCache[key]; ok {
cwCacheMu.Unlock()
return cw
}
cwCacheMu.Unlock()
if cw := queryModelsContextWindow(cfg); cw > 0 {
cwCacheMu.Lock()
cwCache[key] = cw
cwCacheMu.Unlock()
return cw
}
if v, ok := cfg.Raw["context_window"]; ok {
@@ -653,6 +687,8 @@ func (u Usage) Cached() int {
return u.CachedTokens
}
// estTokens is a coarse chars/4 fallback used only when the provider omits
// usage in its response; when real usage is present it is never used.
func estTokens(msgs []Message) int {
chars := 0
for _, m := range msgs {
@@ -668,9 +704,13 @@ func estTokens(msgs []Message) int {
return t
}
func formatUsage(u Usage, cw int) string {
func contextPct(u Usage, cw int) float64 {
if cw <= 0 { cw = 200000 }
pct := float64(u.PromptTokens) * 100.0 / float64(cw)
return float64(u.PromptTokens) * 100.0 / float64(cw)
}
func formatUsage(u Usage, cw int) string {
pct := contextPct(u, cw)
cached := u.Cached()
if cached > 0 {
uncached := u.PromptTokens - cached
@@ -714,7 +754,7 @@ func filterText(s string) string {
for _, r := range s {
if r == ' ' || r == '\t' || r == '\n' {
b.WriteRune(r)
} else if !unicode.Is(unicode.Z, r) && !unicode.IsControl(r) && !unicode.Is(unicode.C, r) && unicode.IsPrint(r) {
} else if unicode.IsPrint(r) {
b.WriteRune(r)
}
}
@@ -741,8 +781,12 @@ func sanitizeMessages(msgs []Message) {
}
func isInvalidAssistantErr(err error) bool {
s := err.Error()
return strings.Contains(s, "Invalid assistant message") || strings.Contains(s, "content or tool_calls must be set")
if err == nil { return false }
s := strings.ToLower(err.Error())
return strings.Contains(s, "invalid assistant message") ||
strings.Contains(s, "content or tool_calls must be set") ||
strings.Contains(s, "tool_calls must be set") ||
strings.Contains(s, "content must be set")
}
func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any) (Message, Usage, error) {
@@ -753,10 +797,10 @@ func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any)
if tools != nil { p["tools"] = tools }
if cfg.Stream { p["stream_options"] = map[string]any{"include_usage": true} }
for k, v := range cfg.Raw {
switch k {
case "endpoint", "model", "temperature", "stream", "api_key", "timeout", "shell_timeout", "max_al_iterations", "color", "context_window":
if internalKey(k) {
continue
default:
}
{
var jv any
if err := json.Unmarshal([]byte(v), &jv); err == nil {
p[k] = jv
@@ -766,11 +810,8 @@ func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any)
}
}
body, _ := json.Marshal(p)
client := &http.Client{Transport: &http.Transport{
DialContext: (&net.Dialer{Timeout: time.Duration(cfg.Timeout) * time.Second}).DialContext,
ResponseHeaderTimeout: time.Duration(cfg.Timeout) * time.Second,
}}
defer client.CloseIdleConnections()
llmTransport.ResponseHeaderTimeout = time.Duration(cfg.Timeout) * time.Second
client := &http.Client{Transport: llmTransport}
fib := []int{1, 1, 2, 3, 5, 8, 13, 21, 34}
pend := c("...requesting...", 1, 2)
var resp *http.Response
@@ -1143,8 +1184,14 @@ func saveSession(msgs []Message) (string, string) {
path = filepath.Join(d, sid+".json")
}
s := Session{sid, time.Now().Format("2006-01-02 15:04:05"), summary(msgs), msgs}
b, _ := json.MarshalIndent(s, "", " ")
os.WriteFile(path, b, 0644)
b, err := json.MarshalIndent(s, "", " ")
if err != nil {
fmt.Fprintf(os.Stderr, "saveSession: marshal error: %v\n", err)
return sid, s.Summary
}
if err := os.WriteFile(path, b, 0644); err != nil {
fmt.Fprintf(os.Stderr, "saveSession: write error: %v\n", err)
}
return sid, s.Summary
}
@@ -1188,8 +1235,14 @@ func loadSession(sid string) ([]Message, error) {
func autosave(msgs []Message) {
s := Session{"autosave", time.Now().Format("2006-01-02 15:04:05"), summary(msgs), msgs}
b, _ := json.MarshalIndent(s, "", " ")
os.WriteFile(filepath.Join(sdir(), "autosave.json"), b, 0644)
b, err := json.MarshalIndent(s, "", " ")
if err != nil {
fmt.Fprintf(os.Stderr, "autosave: marshal error: %v\n", err)
return
}
if err := os.WriteFile(filepath.Join(sdir(), "autosave.json"), b, 0644); err != nil {
fmt.Fprintf(os.Stderr, "autosave: write error: %v\n", err)
}
}
const compactionPrompt = "You are now acting as a compaction engine. Summarize the preceding conversation concisely but completely, preserving all important facts, decisions, code snippets, tool outputs, errors, and current task state so work can seamlessly continue. Output only the summary."
@@ -1602,7 +1655,7 @@ func main() {
msgs = resMsgs
autosave(msgs)
fmt.Println(c(formatUsage(usg, cfg.ContextWindow), 2))
pct := float64(usg.PromptTokens) * 100.0 / float64(cfg.ContextWindow)
pct := contextPct(usg, cfg.ContextWindow)
if pct >= 60.0 && len(msgs) > 1 {
fmt.Print(c(fmt.Sprintf("Context usage is at %.1f%% (%d / %d tokens). Compact conversation? [Y/n]: ", pct, usg.PromptTokens, cfg.ContextWindow), 33))
if ans, ok := readPlain(""); ok {