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

added C-c handling

Browse Source
This commit is contained in:
Luxferre
2026-08-15 17:09:26 +03:00
parent 4509851243
commit 3c55ce14d7
2 changed files with 216 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
main.go
+98 -33
View File
@@ -15,6 +15,7 @@ import (
"net/http"
"os"
"os/exec"
"os/signal"
"path/filepath"
"regexp"
"sort"
@@ -628,7 +629,8 @@ func isInvalidAssistantErr(err error) bool {
return strings.Contains(s, "Invalid assistant message") || strings.Contains(s, "content or tool_calls must be set")
}
func llm(cfg *Cfg, msgs []Message, tools []map[string]any) (Message, error) {
func llm(ctx context.Context, cfg *Cfg, msgs []Message, tools []map[string]any) (Message, error) {
if err := ctx.Err(); err != nil { return Message{}, err }
sanitizeMessages(msgs)
p := map[string]any{"model": cfg.Model, "temperature": cfg.Temperature, "messages": msgs, "stream": cfg.Stream}
if tools != nil { p["tools"] = tools }
@@ -650,13 +652,18 @@ func llm(cfg *Cfg, msgs []Message, tools []map[string]any) (Message, error) {
DialContext: (&net.Dialer{Timeout: time.Duration(cfg.Timeout) * time.Second}).DialContext,
ResponseHeaderTimeout: time.Duration(cfg.Timeout) * time.Second,
}}
defer client.CloseIdleConnections()
fib := []int{1, 1, 2, 3, 5, 8, 13, 21, 34}
pend := c("...requesting...", 1, 2)
var resp *http.Response
var err error
for i := 0; i <= len(fib); i++ {
if err := ctx.Err(); err != nil {
if COL { fmt.Print("\r\033[K") }
return Message{}, err
}
if COL { fmt.Print("\r" + pend) } else { fmt.Println(pend) }
req, _ := http.NewRequest("POST", strings.TrimRight(cfg.Endpoint, "/")+"/chat/completions", bytes.NewReader(body))
req, _ := http.NewRequestWithContext(ctx, "POST", strings.TrimRight(cfg.Endpoint, "/")+"/chat/completions", bytes.NewReader(body))
req.Header.Set("Content-Type", "application/json")
req.Header.Set("User-Agent", "Mozilla/5.0 (compatible; Bantam/1.0)")
if cfg.APIKey != "" && cfg.APIKey != "-" { req.Header.Set("Authorization", "Bearer "+cfg.APIKey) }
@@ -673,12 +680,19 @@ func llm(cfg *Cfg, msgs []Message, tools []map[string]any) (Message, error) {
}
if err == nil { break }
if COL { fmt.Print("\r\033[K") }
if errors.Is(err, context.Canceled) || ctx.Err() != nil {
return Message{}, ctx.Err()
}
if is4xxClientErr {
return Message{}, err
}
if i < len(fib) {
fmt.Println(c(fmt.Sprintf("[network error: %v, retrying in %ds...]", err, fib[i]), 31))
time.Sleep(time.Duration(fib[i]) * time.Second)
select {
case <-ctx.Done():
return Message{}, ctx.Err()
case <-time.After(time.Duration(fib[i]) * time.Second):
}
}
}
if err != nil {
@@ -696,16 +710,19 @@ func llm(cfg *Cfg, msgs []Message, tools []map[string]any) (Message, error) {
} `json:"message"`
} `json:"choices"`
}
if err := json.NewDecoder(resp.Body).Decode(&cr); err != nil { return Message{}, err }
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 len(cr.Choices) == 0 { return Message{}, 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
}
return parseStream(resp.Body)
return parseStream(ctx, resp.Body)
}
func parseStream(r io.Reader) (Message, error) {
func parseStream(ctx context.Context, r io.Reader) (Message, error) {
var content, reas string
var rh, ch bool
var lineBuf string
@@ -724,6 +741,7 @@ func parseStream(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 }
ln := strings.TrimSpace(sc.Text())
if !strings.HasPrefix(ln, "data:") { continue }
data := strings.TrimSpace(ln[5:])
@@ -774,6 +792,7 @@ func parseStream(r io.Reader) (Message, error) {
if tc.Function.Arguments != "" { t.Function.Arguments += tc.Function.Arguments }
}
}
if err := ctx.Err(); err != nil { return Message{}, err }
flushTable()
if lineBuf != "" {
if !mdSt.inCode && isTableLine(lineBuf) {
@@ -796,12 +815,17 @@ func parseStream(r io.Reader) (Message, error) {
return m, sc.Err()
}
func shell(cmd string, timeout int) string {
ctx, cancel := context.WithTimeout(context.Background(), time.Duration(timeout)*time.Second)
func shell(ctx context.Context, cmd string, timeout int) string {
cmdCtx, cancel := context.WithTimeout(ctx, time.Duration(timeout)*time.Second)
defer cancel()
out, err := exec.CommandContext(ctx, "sh", "-c", cmd).CombinedOutput()
c := exec.CommandContext(cmdCtx, "sh", "-c", cmd)
c.WaitDelay = 100 * time.Millisecond
out, err := c.CombinedOutput()
res := strings.TrimSpace(string(out))
if ctx.Err() == context.DeadlineExceeded {
if ctx.Err() != nil {
return "[interrupted]\n\nexit: -1"
}
if cmdCtx.Err() == context.DeadlineExceeded {
return fmt.Sprintf("%s\n\n[shell timeout after %ds]\nexit: -1", res, timeout)
}
code := 0
@@ -824,11 +848,15 @@ func last(msgs []Message) string {
return ""
}
func AL(cfg *Cfg, msgs []Message, sp string, depth int) ([]Message, error) {
func AL(ctx context.Context, cfg *Cfg, msgs []Message, sp string, depth int) ([]Message, error) {
done := false
for i := 0; i < cfg.MaxALIterations && !done; i++ {
m, err := llm(cfg, msgs, TOOLS)
if err := ctx.Err(); err != nil { return msgs, err }
m, err := llm(ctx, cfg, msgs, TOOLS)
if err != nil {
if errors.Is(err, context.Canceled) || ctx.Err() != nil {
return msgs, err
}
if isInvalidAssistantErr(err) {
stripped := false
for j := len(msgs) - 1; j >= 0; j-- {
@@ -861,6 +889,7 @@ func AL(cfg *Cfg, msgs []Message, sp string, depth int) ([]Message, error) {
}
if len(m.ToolCalls) == 0 { done = true; break }
for _, tc := range m.ToolCalls {
if err := ctx.Err(); err != nil { return msgs, err }
fn, astr := tc.Function.Name, tc.Function.Arguments
fmt.Println(c(fmt.Sprintf("[tool call: %s(%s)]", fn, astr), 33))
res, sty := "", 2
@@ -871,18 +900,28 @@ func AL(cfg *Cfg, msgs []Message, sp string, depth int) ([]Message, error) {
switch fn {
case "shell_exec":
cmd, _ := a["command"].(string)
res = shell(cmd, cfg.ShellTimeout)
res = shell(ctx, cmd, cfg.ShellTimeout)
if err := ctx.Err(); err != nil { return msgs, err }
case "run_subagent":
pr, _ := a["prompt"].(string)
if depth >= MAX_DEPTH {
res, sty = fmt.Sprintf("[subagent depth limit (%d) reached, child not spawned]", MAX_DEPTH), 31
} else if subr, err := AL(cfg, []Message{{Role: "system", Content: strp(sp + "\n\nImportant: this is a child agent")}, {Role: "user", Content: strp(pr)}}, sp, depth+1); err != nil {
res, sty = "[subagent error: "+err.Error()+"]", 31
} else {
res, sty = last(subr), 2
subMsgs := []Message{
{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 errors.Is(err, context.Canceled) || ctx.Err() != nil {
return msgs, err
}
res, sty = "[subagent error: "+err.Error()+"]", 31
} else {
res, sty = last(subr), 2
}
}
default:
res, sty = "Unknown tool: "+fn, 31
res, sty = "Unknown tool: " + fn, 31
}
}
fmt.Println(c("[tool result: "+fn+"]", 32) + "\n" + c(res, sty) + "\n")
@@ -989,7 +1028,7 @@ func autosave(msgs []Message) {
os.WriteFile(filepath.Join(sdir(), "autosave.json"), b, 0644)
}
func summarize(cfg *Cfg, msgs []Message) (string, error) {
func summarize(ctx context.Context, cfg *Cfg, msgs []Message) (string, error) {
var sb strings.Builder
for _, m := range msgs {
if m.Role == "system" { continue }
@@ -1013,7 +1052,7 @@ func summarize(cfg *Cfg, msgs []Message) (string, error) {
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(&cc, []Message{{Role: "system", Content: strp(sys)}, {Role: "user", Content: strp("Summarize this conversation:\n\n" + joined)}}, nil)
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 }
@@ -1022,11 +1061,11 @@ func summarize(cfg *Cfg, msgs []Message) (string, error) {
return strings.TrimSpace(s), nil
}
func compact(cfg *Cfg, msgs []Message) ([]Message, string, error) {
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(cfg, msgs)
s, err := summarize(ctx, cfg, msgs)
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
}
@@ -1189,8 +1228,8 @@ func readLine(prompt string) (string, bool) {
e.buf[e.pos] = '\n'
e.pos++
case 0x03:
fmt.Print("\r\n")
return "", false
fmt.Print("^C\r\n")
return "", true
case 0x04:
if len(e.buf) == 0 {
fmt.Print("\r\n")
@@ -1247,14 +1286,23 @@ func main() {
if cmd != "" {
astr, _ := json.Marshal(map[string]string{"command": cmd})
fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33))
res := shell(cmd, cfg.ShellTimeout)
sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
res := shell(sigCtx, cmd, cfg.ShellTimeout)
cancel()
fmt.Println(c("[tool result: shell_exec]", 32) + "\n" + c(res, 2))
}
return
}
msgs = append(msgs, Message{Role: "user", Content: strp(u)})
if msgs, err = AL(&cfg, msgs, sp, 0); err != nil {
fmt.Println(c("[error: "+err.Error()+"]", 31))
sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
msgs, err = AL(sigCtx, &cfg, msgs, sp, 0)
cancel()
if err != nil {
if errors.Is(err, context.Canceled) || sigCtx.Err() != nil {
fmt.Println(c("\n[interrupted]", 33))
} else {
fmt.Println(c("[error: "+err.Error()+"]", 31))
}
os.Exit(1)
}
autosave(msgs)
@@ -1317,9 +1365,15 @@ func main() {
continue
}
fmt.Println(c("[compacting conversation...]", 33))
nm, sm, err := compact(&cfg, msgs)
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))
if errors.Is(err, context.Canceled) || sigCtx.Err() != nil {
fmt.Println(c("\n[interrupted]", 33))
} else {
fmt.Println(c("[compact failed: "+err.Error()+"]", 31))
}
continue
}
msgs = nm
@@ -1354,7 +1408,9 @@ func main() {
if cmd != "" {
astr, _ := json.Marshal(map[string]string{"command": cmd})
fmt.Println(c(fmt.Sprintf("[tool call: shell_exec(%s)]", string(astr)), 33))
res := shell(cmd, cfg.ShellTimeout)
sigCtx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
res := shell(sigCtx, cmd, cfg.ShellTimeout)
cancel()
fmt.Println(c("[tool result: shell_exec]", 32) + "\n" + c(res, 2))
}
continue
@@ -1365,11 +1421,20 @@ func main() {
}
continue
}
msgs = append(msgs, Message{Role: "user", Content: strp(u)})
var err error
if msgs, err = AL(&cfg, msgs, sp, 0); err != nil {
fmt.Println(c("[error: "+err.Error()+"]", 31))
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)
cancel()
if err != nil {
if errors.Is(err, context.Canceled) || sigCtx.Err() != nil {
fmt.Println(c("\n[interrupted]", 33))
} else {
fmt.Println(c("[error: "+err.Error()+"]", 31))
}
continue
}
msgs = resMsgs
autosave(msgs)
}
done: