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

improved compacting algo and token visibility

Browse Source
This commit is contained in:
Luxferre
2026-08-16 08:12:49 +03:00
parent 3c55ce14d7
commit d1597444d0
3 changed files with 429 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main.go
+206 -66
View File
@@ -44,10 +44,11 @@ type Cfg struct {
MaxALIterations int
Stream bool
Color string
ContextWindow int
Raw map[string]string
}
var defCfg = Cfg{"https://opencode.ai/zen/v1", "big-pickle", "-", 0.7, 300, 120, 1000, true, "auto", nil}
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 {
if v, e := strconv.Atoi(strings.TrimSpace(s)); e == nil {
@@ -56,6 +57,50 @@ func atoiD(s string, d int) int {
return d
}
func queryModelsContextWindow(cfg *Cfg) int {
client := &http.Client{Timeout: 3 * time.Second}
req, err := http.NewRequest("GET", strings.TrimRight(cfg.Endpoint, "/")+"/models", nil)
if err != nil { return 0 }
req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; Bantam/1.0)")
if cfg.APIKey != "" && cfg.APIKey != "-" { req.Header.Set("Authorization", "Bearer "+cfg.APIKey) }
resp, err := client.Do(req)
if err != nil || resp.StatusCode >= 400 { return 0 }
defer resp.Body.Close()
var res struct {
Data []map[string]any `json:"data"`
Models []map[string]any `json:"models"`
}
if json.NewDecoder(resp.Body).Decode(&res) != nil { return 0 }
list := res.Data
if len(list) == 0 { list = res.Models }
for _, item := range list {
id, _ := item["id"].(string)
if id == cfg.Model || strings.EqualFold(id, cfg.Model) {
for _, key := range []string{"context_window", "context_length", "max_context_length", "max_model_len", "context_size", "max_tokens", "max_input_tokens"} {
if val, ok := item[key]; ok {
switch v := val.(type) {
case float64:
if v > 0 { return int(v) }
case string:
if n := atoiD(v, 0); n > 0 { return n }
}
}
}
}
}
return 0
}
func fetchContextWindow(cfg *Cfg) int {
if cw := queryModelsContextWindow(cfg); cw > 0 {
return cw
}
if v, ok := cfg.Raw["context_window"]; ok {
return atoiD(v, 200000)
}
return 200000
}
func getCfg(path string) Cfg {
cfg := defCfg
cfg.Raw = map[string]string{
@@ -63,6 +108,7 @@ func getCfg(path string) Cfg {
"api_key": cfg.APIKey, "stream": strconv.FormatBool(cfg.Stream), "color": cfg.Color,
"timeout": strconv.Itoa(cfg.Timeout), "shell_timeout": strconv.Itoa(cfg.ShellTimeout),
"max_al_iterations": strconv.Itoa(cfg.MaxALIterations),
"context_window": strconv.Itoa(cfg.ContextWindow),
}
if d, err := os.ReadFile(path); err == nil {
for _, ln := range strings.Split(string(d), "\n") {
@@ -81,6 +127,7 @@ func getCfg(path string) Cfg {
case "max_al_iterations": cfg.MaxALIterations = atoiD(v, cfg.MaxALIterations)
case "stream": cfg.Stream = v == "true" || v == "1" || v == "yes"
case "color": cfg.Color = v
case "context_window": cfg.ContextWindow = atoiD(v, cfg.ContextWindow)
}
}
}
@@ -590,6 +637,48 @@ var TOOLS = []map[string]any{
func strp(s string) *string { return &s }
type Usage struct {
PromptTokens int `json:"prompt_tokens"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
PromptTokensDetails struct {
CachedTokens int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
CachedTokens int `json:"cached_tokens"`
}
func (u Usage) Cached() int {
if u.PromptTokensDetails.CachedTokens > 0 { return u.PromptTokensDetails.CachedTokens }
return u.CachedTokens
}
func estTokens(msgs []Message) int {
chars := 0
for _, m := range msgs {
if m.Content != nil { chars += len(*m.Content) }
chars += len(m.ReasoningContent)
for _, tc := range m.ToolCalls {
chars += len(tc.Function.Name) + len(tc.Function.Arguments)
}
}
if chars == 0 { return 0 }
t := chars / 4
if t == 0 { t = 1 }
return t
}
func formatUsage(u Usage, cw int) string {
if cw <= 0 { cw = 200000 }
pct := float64(u.PromptTokens) * 100.0 / float64(cw)
cached := u.Cached()
if cached > 0 {
uncached := u.PromptTokens - cached
if uncached < 0 { uncached = 0 }
return fmt.Sprintf("[tokens: %d prompt (%d cached, %d uncached) + %d completion | context: %d/%d (%.1f%%)]", u.PromptTokens, cached, uncached, u.CompletionTokens, u.PromptTokens, cw, pct)
}
return fmt.Sprintf("[tokens: %d prompt + %d completion | context: %d/%d (%.1f%%)]", u.PromptTokens, u.CompletionTokens, u.PromptTokens, cw, pct)
}
type streamDelta struct {
Choices []struct {
Delta struct {
@@ -606,6 +695,7 @@ type streamDelta struct {
} `json:"tool_calls"`
} `json:"delta"`
} `json:"choices"`
Usage *Usage `json:"usage"`
}
func sanitizeMessages(msgs []Message) {
@@ -629,14 +719,15 @@ func isInvalidAssistantErr(err error) bool {
return strings.Contains(s, "Invalid assistant message") || strings.Contains(s, "content or tool_calls must be set")
}
func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any) (Message, error) {
if err := ctx.Err(); err != nil { return Message{}, err }
func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any) (Message, Usage, error) {
if err := ctx.Err(); err != nil { return Message{}, Usage{}, err }
sanitizeMessages(msgs)
p := map[string]any{"model": cfg.Model, "temperature": cfg.Temperature, "messages": msgs, "stream": cfg.Stream}
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", "api_key", "timeout", "shell_timeout", "max_al_iterations", "color":
case "endpoint", "api_key", "timeout", "shell_timeout", "max_al_iterations", "color", "context_window":
continue
default:
var jv any
@@ -660,7 +751,7 @@ func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any)
for i := 0; i <= len(fib); i++ {
if err := ctx.Err(); err != nil {
if COL { fmt.Print("\r\033[K") }
return Message{}, err
return Message{}, Usage{}, err
}
if COL { fmt.Print("\r" + pend) } else { fmt.Println(pend) }
req, _ := http.NewRequestWithContext(ctx, "POST", strings.TrimRight(cfg.Endpoint, "/")+"/chat/completions", bytes.NewReader(body))
@@ -681,23 +772,23 @@ func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any)
if err == nil { break }
if COL { fmt.Print("\r\033[K") }
if errors.Is(err, context.Canceled) || ctx.Err() != nil {
return Message{}, ctx.Err()
return Message{}, Usage{}, ctx.Err()
}
if is4xxClientErr {
return Message{}, err
return Message{}, Usage{}, err
}
if i < len(fib) {
fmt.Println(c(fmt.Sprintf("[network error: %v, retrying in %ds...]", err, fib[i]), 31))
select {
case <-ctx.Done():
return Message{}, ctx.Err()
return Message{}, Usage{}, ctx.Err()
case <-time.After(time.Duration(fib[i]) * time.Second):
}
}
}
if err != nil {
if COL { fmt.Print("\r\033[K") }
return Message{}, err
return Message{}, Usage{}, err
}
defer resp.Body.Close()
if COL { fmt.Print("\r\033[K") }
@@ -709,25 +800,39 @@ func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any)
Reasoning string `json:"reasoning"`
} `json:"message"`
} `json:"choices"`
Usage Usage `json:"usage"`
}
if err := json.NewDecoder(resp.Body).Decode(&cr); err != nil {
if errors.Is(err, context.Canceled) || ctx.Err() != nil { return Message{}, ctx.Err() }
return Message{}, err
if errors.Is(err, context.Canceled) || ctx.Err() != nil { return Message{}, Usage{}, ctx.Err() }
return Message{}, Usage{}, err
}
if len(cr.Choices) == 0 { return Message{}, errors.New("empty choices in LLM response") }
if len(cr.Choices) == 0 { return Message{}, Usage{}, errors.New("empty choices in LLM response") }
m := cr.Choices[0].Message.Message
if m.ReasoningContent == "" { m.ReasoningContent = cr.Choices[0].Message.Reasoning }
return m, nil
u := cr.Usage
if u.PromptTokens == 0 {
u.PromptTokens = estTokens(msgs)
u.CompletionTokens = estTokens([]Message{m})
u.TotalTokens = u.PromptTokens + u.CompletionTokens
}
return m, u, nil
}
return parseStream(ctx, resp.Body)
m, u, err := parseStream(ctx, resp.Body)
if err == nil && u.PromptTokens == 0 {
u.PromptTokens = estTokens(msgs)
u.CompletionTokens = estTokens([]Message{m})
u.TotalTokens = u.PromptTokens + u.CompletionTokens
}
return m, u, err
}
func parseStream(ctx context.Context, r io.Reader) (Message, error) {
func parseStream(ctx context.Context, r io.Reader) (Message, Usage, error) {
var content, reas string
var rh, ch bool
var lineBuf string
var mdSt mdState
var tblBuf []string
var lastUsage Usage
tcs := map[int]*ToolCall{}
var order []int
@@ -741,13 +846,17 @@ func parseStream(ctx context.Context, r io.Reader) (Message, error) {
sc := bufio.NewScanner(r)
sc.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)
for sc.Scan() {
if err := ctx.Err(); err != nil { return Message{}, err }
if err := ctx.Err(); err != nil { return Message{}, lastUsage, err }
ln := strings.TrimSpace(sc.Text())
if !strings.HasPrefix(ln, "data:") { continue }
data := strings.TrimSpace(ln[5:])
if data == "[DONE]" { break }
var d streamDelta
if json.Unmarshal([]byte(data), &d) != nil || len(d.Choices) == 0 { continue }
if json.Unmarshal([]byte(data), &d) != nil { continue }
if d.Usage != nil && (d.Usage.PromptTokens > 0 || d.Usage.TotalTokens > 0) {
lastUsage = *d.Usage
}
if len(d.Choices) == 0 { continue }
dl := d.Choices[0].Delta
rc := dl.ReasoningContent
if rc == "" { rc = dl.Reasoning }
@@ -792,7 +901,7 @@ func parseStream(ctx context.Context, r io.Reader) (Message, error) {
if tc.Function.Arguments != "" { t.Function.Arguments += tc.Function.Arguments }
}
}
if err := ctx.Err(); err != nil { return Message{}, err }
if err := ctx.Err(); err != nil { return Message{}, lastUsage, err }
flushTable()
if lineBuf != "" {
if !mdSt.inCode && isTableLine(lineBuf) {
@@ -812,7 +921,10 @@ func parseStream(ctx context.Context, r io.Reader) (Message, error) {
m.ToolCalls = make([]ToolCall, 0, len(order))
for _, idx := range order { m.ToolCalls = append(m.ToolCalls, *tcs[idx]) }
}
return m, sc.Err()
if lastUsage.TotalTokens == 0 && lastUsage.PromptTokens > 0 {
lastUsage.TotalTokens = lastUsage.PromptTokens + lastUsage.CompletionTokens
}
return m, lastUsage, sc.Err()
}
func shell(ctx context.Context, cmd string, timeout int) string {
@@ -848,14 +960,15 @@ func last(msgs []Message) string {
return ""
}
func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]Message, error) {
func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]Message, Usage, error) {
done := false
var turnUsage Usage
for i := 0; i < cfg.MaxALIterations && !done; i++ {
if err := ctx.Err(); err != nil { return msgs, err }
m, err := llm(ctx, cfg, msgs, TOOLS)
if err := ctx.Err(); err != nil { return msgs, turnUsage, err }
m, u, err := llm(ctx, cfg, msgs, TOOLS)
if err != nil {
if errors.Is(err, context.Canceled) || ctx.Err() != nil {
return msgs, err
return msgs, turnUsage, err
}
if isInvalidAssistantErr(err) {
stripped := false
@@ -869,8 +982,12 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]
}
if stripped { continue }
}
return msgs, err
return msgs, turnUsage, err
}
turnUsage.PromptTokens = u.PromptTokens
turnUsage.CompletionTokens += u.CompletionTokens
turnUsage.TotalTokens += u.TotalTokens
if u.Cached() > 0 { turnUsage.CachedTokens = u.Cached() }
for j := range m.ToolCalls {
tc := &m.ToolCalls[j]
astr := tc.Function.Arguments
@@ -889,7 +1006,7 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]
}
if len(m.ToolCalls) == 0 { done = true; break }
for _, tc := range m.ToolCalls {
if err := ctx.Err(); err != nil { return msgs, err }
if err := ctx.Err(); err != nil { return msgs, turnUsage, err }
fn, astr := tc.Function.Name, tc.Function.Arguments
fmt.Println(c(fmt.Sprintf("[tool call: %s(%s)]", fn, astr), 33))
res, sty := "", 2
@@ -901,7 +1018,7 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]
case "shell_exec":
cmd, _ := a["command"].(string)
res = shell(ctx, cmd, cfg.ShellTimeout)
if err := ctx.Err(); err != nil { return msgs, err }
if err := ctx.Err(); err != nil { return msgs, turnUsage, err }
case "run_subagent":
pr, _ := a["prompt"].(string)
if depth >= MAX_DEPTH {
@@ -911,12 +1028,14 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]
{Role: "system", Content: strp(sp + "\n\nImportant: this is a child agent")},
{Role: "user", Content: strp(pr)},
}
if subr, err := AL(ctx, cfg, subMsgs, sp, depth+1); err != nil {
if subr, subu, err := AL(ctx, cfg, subMsgs, sp, depth+1); err != nil {
if errors.Is(err, context.Canceled) || ctx.Err() != nil {
return msgs, err
return msgs, turnUsage, err
}
res, sty = "[subagent error: "+err.Error()+"]", 31
} else {
turnUsage.CompletionTokens += subu.CompletionTokens
turnUsage.TotalTokens += subu.TotalTokens
res, sty = last(subr), 2
}
}
@@ -931,7 +1050,7 @@ func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]
if !done {
msgs = append(msgs, Message{Role: "assistant", Content: strp(fmt.Sprintf("[max AL iterations (%d) reached]", cfg.MaxALIterations))})
}
return msgs, nil
return msgs, turnUsage, nil
}
func homeDir() string {
@@ -1028,46 +1147,38 @@ func autosave(msgs []Message) {
os.WriteFile(filepath.Join(sdir(), "autosave.json"), b, 0644)
}
func summarize(ctx context.Context, cfg *Cfg, msgs []Message) (string, error) {
var sb strings.Builder
for _, m := range msgs {
if m.Role == "system" { continue }
ct := ""
if m.Content != nil { ct = *m.Content }
if ct == "" && len(m.ToolCalls) > 0 {
jc := make([]map[string]any, 0, len(m.ToolCalls))
for _, tc := range m.ToolCalls {
jc = append(jc, map[string]any{"function": map[string]any{"name": tc.Function.Name, "arguments": tc.Function.Arguments}})
}
b, _ := json.Marshal(jc)
ct = string(b)
}
if ct == "" { continue }
if len(ct) > 4000 { ct = ct[:4000] + "...[truncated]" }
sb.WriteString(m.Role + ": " + ct + "\n\n")
}
if sb.Len() == 0 { return "", errors.New("no conversation to summarize") }
joined := sb.String()
if len(joined) > 100000 { joined = joined[len(joined)-100000:] + "\n...[earlier parts truncated]" }
sys := "You are a conversation summarizer for an AI agent's context window. Summarize concisely but completely, preserving all important facts, decisions, code, errors, and the current task state, so the agent can continue the work without the original messages. Output only the summary."
cc := *cfg
cc.Stream = false
m, err := llm(ctx, &cc, []Message{{Role: "system", Content: strp(sys)}, {Role: "user", Content: strp("Summarize this conversation:\n\n" + joined)}}, nil)
if err != nil { return "", err }
s := ""
if m.Content != nil { s = *m.Content }
if s == "" { s = m.ReasoningContent }
if strings.TrimSpace(s) == "" { return "", errors.New("LLM returned an empty summary") }
return strings.TrimSpace(s), nil
}
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."
func compact(ctx context.Context, cfg *Cfg, msgs []Message) ([]Message, string, error) {
if len(msgs) == 0 || msgs[0].Role != "system" {
return msgs, "", errors.New("session has no system message")
}
s, err := summarize(ctx, cfg, msgs)
if len(msgs) <= 1 {
return msgs, "", errors.New("nothing to compact")
}
cMsgs := append(append([]Message{}, msgs...), Message{
Role: "user",
Content: strp(compactionPrompt),
})
cc := *cfg
cc.Stream = false
m, _, err := llm(ctx, &cc, cMsgs, nil)
if err != nil { return msgs, "", err }
return []Message{{Role: "system", Content: msgs[0].Content}, {Role: "user", Content: strp("Summary of the previous conversation:\n" + s + "\n\nPlease continue from here.")}}, s, nil
s := ""
if m.Content != nil { s = *m.Content }
if s == "" { s = m.ReasoningContent }
s = strings.TrimSpace(s)
if s == "" { return msgs, "", errors.New("LLM returned an empty summary") }
newMsgs := []Message{
{Role: "system", Content: msgs[0].Content},
{Role: "user", Content: strp("Summary of the previous conversation:\n" + s + "\n\nPlease continue from here.")},
}
return newMsgs, s, nil
}
func summarize(ctx context.Context, cfg *Cfg, msgs []Message) (string, error) {
_, s, err := compact(ctx, cfg, msgs)
return s, err
}
func loadHistory() {
@@ -1270,6 +1381,7 @@ func readLine(prompt string) (string, bool) {
func main() {
sp := prompt("system.txt")
cfg := getCfg("model.cfg")
cfg.ContextWindow = fetchContextWindow(&cfg)
COL = col(cfg)
stdin = bufio.NewReader(os.Stdin)
msgs := []Message{{Role: "system", Content: strp(sp)}}
@@ -1295,7 +1407,8 @@ func main() {
}
msgs = append(msgs, Message{Role: "user", Content: strp(u)})
sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
msgs, err = AL(sigCtx, &cfg, msgs, sp, 0)
var usg Usage
msgs, usg, err = AL(sigCtx, &cfg, msgs, sp, 0)
cancel()
if err != nil {
if errors.Is(err, context.Canceled) || sigCtx.Err() != nil {
@@ -1305,12 +1418,13 @@ func main() {
}
os.Exit(1)
}
fmt.Println(c(formatUsage(usg, cfg.ContextWindow), 2))
autosave(msgs)
return
}
loadHistory()
fmt.Println(c("Bantam Agent ready", 1, 32) + c(" (Ctrl+J = new line)", 2))
fmt.Println(c(fmt.Sprintf("endpoint: %s model: %s temp: %v", cfg.Endpoint, cfg.Model, cfg.Temperature), 2))
fmt.Println(c(fmt.Sprintf("endpoint: %s model: %s temp: %v context: %d", cfg.Endpoint, cfg.Model, cfg.Temperature, cfg.ContextWindow), 2))
for {
u, ok := readLine(c("> ", 1, 36))
if !ok {
@@ -1397,6 +1511,9 @@ func main() {
continue
}
cfg = getCfg("model.cfg")
if k == "model" || k == "endpoint" || k == "api_key" {
cfg.ContextWindow = fetchContextWindow(&cfg)
}
COL = col(cfg)
fmt.Println(c(fmt.Sprintf("[config updated: %s=%s]", k, v), 32))
} else {
@@ -1424,7 +1541,7 @@ func main() {
turnMsgs := append([]Message{}, msgs...)
turnMsgs = append(turnMsgs, Message{Role: "user", Content: strp(u)})
sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
resMsgs, err := AL(sigCtx, &cfg, turnMsgs, sp, 0)
resMsgs, usg, err := AL(sigCtx, &cfg, turnMsgs, sp, 0)
cancel()
if err != nil {
if errors.Is(err, context.Canceled) || sigCtx.Err() != nil {
@@ -1436,6 +1553,29 @@ func main() {
}
msgs = resMsgs
autosave(msgs)
fmt.Println(c(formatUsage(usg, cfg.ContextWindow), 2))
pct := float64(usg.PromptTokens) * 100.0 / float64(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))
ans, ok := readPlain("")
if ok {
ans = strings.TrimSpace(strings.ToLower(ans))
if ans == "" || ans == "y" || ans == "yes" {
fmt.Println(c("[compacting conversation...]", 33))
sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
nm, sm, err := compact(sigCtx, &cfg, msgs)
cancel()
if err != nil {
fmt.Println(c("[compact failed: "+err.Error()+"]", 31))
} else {
msgs = nm
autosave(msgs)
fmt.Println(c(fmt.Sprintf("[compacted to %d messages]", len(msgs)), 32))
fmt.Println(c("--- summary ---", 33) + "\n" + c(sm, 2))
}
}
}
}
}
done:
autosave(msgs)