git.giacomolli.uk - onechat/blob - tools.go

   1 package main
   2
   3 import (
   4         "bytes"
   5         "context"
   6         "encoding/json"
   7         "errors"
   8         "fmt"
   9         "io"
  10         "math/bits"
  11         "os"
  12         "os/exec"
  13         "path/filepath"
  14         "strconv"
  15         "strings"
  16         "syscall"
  17         "time"
  18         "unicode/utf8"
  19 )
  20
  21 const (
  22         // toolClipCharacters is the most characters of a tool output: its head
  23         // and its tail, half each.
  24         toolClipCharacters = 30000
  25
  26         // toolClipBytes is the most bytes that a tool keeps of each end of an
  27         // output, enough for toolClipCharacters/2 characters of 4 bytes.
  28         toolClipBytes = 4 * toolClipCharacters / 2
  29
  30         // toolShellWaitDelay is how long a command can keep its output open
  31         // after it ends or after its kill.
  32         toolShellWaitDelay = 2 * time.Second
  33 )
  34
  35 // ToolLog is the log that the tools write their calls and results to.
  36 type ToolLog interface {
  37         Append(kind, text string) (Message, error)
  38         Flush(ctx context.Context) error
  39 }
  40
  41 // ToolChat gives the messages and the nodes that zoom and date read.
  42 type ToolChat interface {
  43         Len() int
  44         Message(id int) (Message, error)
  45         Node(level, index int) (Node, bool, error)
  46 }
  47
  48 // Toolbox runs the tool calls of one agent. Each call and each result goes to
  49 // the log as it happens: the call as a tool message, the result as an echo
  50 // message. The call is durable before its tool runs, so the log never misses an
  51 // action outside the agent. Commands and files run as the user in the working
  52 // directory, with no sandbox.
  53 type Toolbox struct {
  54         log     ToolLog
  55         chat    ToolChat
  56         workdir string
  57 }
  58
  59 // Call logs the call of the tool name with its JSON arguments, makes the log
  60 // durable, runs the tool, and logs its clipped output, which it returns. A
  61 // failed tool gives its error as the output. An error of Call itself is an
  62 // error of the log. When ctx ends before the flush, the tool does not run.
  63 func (b *Toolbox) Call(ctx context.Context, name, arguments string) (string, error) {
  64         _, err := b.log.Append("tool", toolCallText(name, arguments))
  65         if err != nil {
  66                 return "", err
  67         }
  68
  69         err = b.log.Flush(ctx)
  70         if err != nil && ctx.Err() == nil {
  71                 return "", err
  72         }
  73
  74         output := "not run: the turn was interrupted"
  75         if err == nil {
  76                 output = toolClip(b.run(ctx, name, arguments))
  77         }
  78
  79         _, err = b.log.Append("echo", output)
  80         if err != nil {
  81                 return "", err
  82         }
  83
  84         return output, nil
  85 }
  86
  87 // run runs one tool and returns its output, or its error as text.
  88 func (b *Toolbox) run(ctx context.Context, name, arguments string) string {
  89         var output string
  90         var err error
  91         switch name {
  92         case "zoom":
  93                 var args ToolZoomArguments
  94                 err = toolDecodeArguments(arguments, &args)
  95                 if err == nil {
  96                         output, err = b.zoom(args.ID, args.N)
  97                 }
  98         case "date":
  99                 var args ToolDateArguments
 100                 err = toolDecodeArguments(arguments, &args)
 101                 if err == nil {
 102                         output, err = b.date(args.ID)
 103                 }
 104         case "shell":
 105                 var args ToolShellArguments
 106                 err = toolDecodeArguments(arguments, &args)
 107                 if err == nil {
 108                         output = b.shell(ctx, args.Command)
 109                 }
 110         case "read":
 111                 var args ToolPathArguments
 112                 err = toolDecodeArguments(arguments, &args)
 113                 if err == nil {
 114                         output, err = b.read(args.Path)
 115                 }
 116         case "write":
 117                 var args ToolWriteArguments
 118                 err = toolDecodeArguments(arguments, &args)
 119                 if err == nil {
 120                         output, err = b.write(args.Path, args.Text)
 121                 }
 122         case "edit":
 123                 var args ToolEditArguments
 124                 err = toolDecodeArguments(arguments, &args)
 125                 if err == nil {
 126                         output, err = b.edit(args.Path, args.Old, args.New)
 127                 }
 128         default:
 129                 err = fmt.Errorf("unknown tool %q", name)
 130         }
 131
 132         if err != nil {
 133                 return "error: " + err.Error()
 134         }
 135
 136         return output
 137 }
 138
 139 // zoom returns message id whole for n = 1, or else the two lines under the line
 140 // id+n.
 141 func (b *Toolbox) zoom(id, n int) (string, error) {
 142         if n < 1 || bits.OnesCount(uint(n)) != 1 || id < 0 || id%n != 0 {
 143                 return "", fmt.Errorf("line %d+%d does not exist: n must be a power of 2 and id a multiple of n", id, n)
 144         }
 145
 146         if id+n > b.chat.Len() {
 147                 return "", fmt.Errorf("line %d+%d goes past the last message, %d", id, n, b.chat.Len()-1)
 148         }
 149
 150         if n == 1 {
 151                 message, err := b.chat.Message(id)
 152                 if err != nil {
 153                         return "", err
 154                 }
 155
 156                 return leafSource(message), nil
 157         }
 158
 159         level := bits.TrailingZeros(uint(n)) - 1
 160         var lines []string
 161         for _, index := range []int{2 * id / n, 2*id/n + 1} {
 162                 node, built, err := b.chat.Node(level, index)
 163                 if err != nil {
 164                         return "", err
 165                 }
 166
 167                 text := viewPending
 168                 if built {
 169                         text = strings.ReplaceAll(node.Text, "\n", " ")
 170                 }
 171
 172                 lines = append(lines, fmt.Sprintf("%d+%d|%s", index<<level, 1<<level, text))
 173         }
 174
 175         return strings.Join(lines, "\n"), nil
 176 }
 177
 178 // date returns the date and time of message id, with its weekday.
 179 func (b *Toolbox) date(id int) (string, error) {
 180         if id < 0 || id >= b.chat.Len() {
 181                 return "", fmt.Errorf("message %d does not exist", id)
 182         }
 183
 184         message, err := b.chat.Message(id)
 185         if err != nil {
 186                 return "", err
 187         }
 188
 189         date, err := time.Parse(time.RFC3339Nano, message.Date)
 190         if err != nil {
 191                 return "", err
 192         }
 193
 194         return date.Local().Format("Monday 2006-01-02 15:04:05 -07:00"), nil
 195 }
 196
 197 // shell runs command with bash in the working directory and returns its output
 198 // and error output. A failed command adds its exit status. When ctx ends, the
 199 // whole process group of the command is killed.
 200 func (b *Toolbox) shell(ctx context.Context, command string) string {
 201         cmd := exec.CommandContext(ctx, "bash", "-c", command)
 202         cmd.Dir = b.workdir
 203         cmd.SysProcAttr = &syscall.SysProcAttr{
 204                 Setpgid: true,
 205         }
 206         cmd.Cancel = func() error {
 207                 return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
 208         }
 209         cmd.WaitDelay = toolShellWaitDelay
 210
 211         output := &ToolHeadTail{}
 212         cmd.Stdout = output
 213         cmd.Stderr = output
 214         err := cmd.Run()
 215
 216         text := output.String()
 217         var exit *exec.ExitError
 218         switch {
 219         case ctx.Err() != nil:
 220                 text += "\n[interrupted]"
 221         case errors.As(err, &exit):
 222                 text += "\n[" + exit.String() + "]"
 223         case err != nil:
 224                 text += "\n[error: " + err.Error() + "]"
 225         case text == "":
 226                 text = "(no output)"
 227         }
 228
 229         return text
 230 }
 231
 232 // read returns the text of the file at path.
 233 func (b *Toolbox) read(path string) (text string, err error) {
 234         file, err := os.Open(b.path(path))
 235         if err != nil {
 236                 return "", err
 237         }
 238
 239         defer func() {
 240                 err = errors.Join(err, file.Close())
 241         }()
 242
 243         output := &ToolHeadTail{}
 244         _, err = io.Copy(output, file)
 245         if err != nil {
 246                 return "", err
 247         }
 248
 249         return output.String(), nil
 250 }
 251
 252 // write replaces the file at path with text, and makes its missing
 253 // directories.
 254 func (b *Toolbox) write(path, text string) (string, error) {
 255         path = b.path(path)
 256         err := os.MkdirAll(filepath.Dir(path), 0o755)
 257         if err != nil {
 258                 return "", err
 259         }
 260
 261         err = os.WriteFile(path, []byte(text), 0o644)
 262         if err != nil {
 263                 return "", err
 264         }
 265
 266         return fmt.Sprintf("wrote %d bytes to %s", len(text), path), nil
 267 }
 268
 269 // edit replaces the one occurrence of old in the file at path with
 270 // replacement. It keeps the mode of the file.
 271 func (b *Toolbox) edit(path, old, replacement string) (string, error) {
 272         path = b.path(path)
 273         data, err := os.ReadFile(path)
 274         if err != nil {
 275                 return "", err
 276         }
 277
 278         if old == "" {
 279                 return "", errors.New("old is empty")
 280         }
 281
 282         count := strings.Count(string(data), old)
 283         if count != 1 {
 284                 return "", fmt.Errorf("old occurs %d times in %s, not once", count, path)
 285         }
 286
 287         info, err := os.Stat(path)
 288         if err != nil {
 289                 return "", err
 290         }
 291
 292         err = os.WriteFile(path, []byte(strings.Replace(string(data), old, replacement, 1)), info.Mode().Perm())
 293         if err != nil {
 294                 return "", err
 295         }
 296
 297         return "edited " + path, nil
 298 }
 299
 300 // path returns path, relative to the working directory when it is not
 301 // absolute.
 302 func (b *Toolbox) path(path string) string {
 303         if filepath.IsAbs(path) {
 304                 return path
 305         }
 306
 307         return filepath.Join(b.workdir, path)
 308 }
 309
 310 // ToolZoomArguments are the arguments of zoom.
 311 type ToolZoomArguments struct {
 312         ID int `json:"id"`
 313         N  int `json:"n"`
 314 }
 315
 316 // ToolDateArguments are the arguments of date.
 317 type ToolDateArguments struct {
 318         ID int `json:"id"`
 319 }
 320
 321 // ToolShellArguments are the arguments of shell.
 322 type ToolShellArguments struct {
 323         Command string `json:"command"`
 324 }
 325
 326 // ToolPathArguments are the arguments of read.
 327 type ToolPathArguments struct {
 328         Path string `json:"path"`
 329 }
 330
 331 // ToolWriteArguments are the arguments of write.
 332 type ToolWriteArguments struct {
 333         Path string `json:"path"`
 334         Text string `json:"text"`
 335 }
 336
 337 // ToolEditArguments are the arguments of edit.
 338 type ToolEditArguments struct {
 339         Path string `json:"path"`
 340         Old  string `json:"old"`
 341         New  string `json:"new"`
 342 }
 343
 344 // ToolHeadTail keeps the first and the last toolClipBytes bytes written to it,
 345 // and counts the bytes between them that it drops.
 346 type ToolHeadTail struct {
 347         head    []byte
 348         tail    []byte
 349         dropped int
 350 }
 351
 352 // String returns the output, clipped when it is too long. Invalid UTF-8 becomes
 353 // the replacement character.
 354 func (h *ToolHeadTail) String() string {
 355         replacement := string(utf8.RuneError)
 356         if h.dropped == 0 {
 357                 return toolClip(strings.ToValidUTF8(string(h.head)+string(h.tail), replacement))
 358         }
 359
 360         head := toolFirstCharacters(strings.ToValidUTF8(string(h.head), replacement), toolClipCharacters/2)
 361         tail := toolLastCharacters(strings.ToValidUTF8(string(h.tail), replacement), toolClipCharacters/2)
 362         cut := len(h.head) + h.dropped + len(h.tail) - len(head) - len(tail)
 363
 364         return toolJoinClipped(head, tail, cut)
 365 }
 366
 367 // Write keeps the bytes of p at the ends of the output.
 368 func (h *ToolHeadTail) Write(p []byte) (int, error) {
 369         n := len(p)
 370         room := min(toolClipBytes-len(h.head), len(p))
 371         h.head = append(h.head, p[:room]...)
 372         p = p[room:]
 373
 374         h.tail = append(h.tail, p...)
 375         if len(h.tail) > toolClipBytes {
 376                 cut := len(h.tail) - toolClipBytes
 377                 h.dropped += cut
 378                 h.tail = bytes.Clone(h.tail[cut:])
 379         }
 380
 381         return n, nil
 382 }
 383
 384 // toolCallText returns the log text of a tool call: the name and the JSON
 385 // arguments.
 386 func toolCallText(name, arguments string) string {
 387         return name + " " + arguments
 388 }
 389
 390 // toolClip returns text when it has at most toolClipCharacters characters.
 391 // Otherwise it returns the first and the last toolClipCharacters/2 characters,
 392 // with a mark of the cut between them.
 393 func toolClip(text string) string {
 394         if utf8.RuneCountInString(text) <= toolClipCharacters {
 395                 return text
 396         }
 397
 398         head := toolFirstCharacters(text, toolClipCharacters/2)
 399         tail := toolLastCharacters(text, toolClipCharacters/2)
 400
 401         return toolJoinClipped(head, tail, len(text)-len(head)-len(tail))
 402 }
 403
 404 // toolDecodeArguments decodes the JSON arguments of a call into args. An
 405 // unknown field is an error.
 406 func toolDecodeArguments(arguments string, args any) error {
 407         decoder := json.NewDecoder(strings.NewReader(arguments))
 408         decoder.DisallowUnknownFields()
 409         err := decoder.Decode(args)
 410         if err != nil {
 411                 return fmt.Errorf("arguments: %w", err)
 412         }
 413
 414         return nil
 415 }
 416
 417 // toolFirstCharacters returns the first count characters of text.
 418 func toolFirstCharacters(text string, count int) string {
 419         end := 0
 420         for range count {
 421                 if end >= len(text) {
 422                         break
 423                 }
 424
 425                 _, size := utf8.DecodeRuneInString(text[end:])
 426                 end += size
 427         }
 428
 429         return text[:end]
 430 }
 431
 432 // toolLastCharacters returns the last count characters of text.
 433 func toolLastCharacters(text string, count int) string {
 434         start := len(text)
 435         for range count {
 436                 if start <= 0 {
 437                         break
 438                 }
 439
 440                 _, size := utf8.DecodeLastRuneInString(text[:start])
 441                 start -= size
 442         }
 443
 444         return text[start:]
 445 }
 446
 447 // toolJoinClipped joins the head and the tail of a clipped output with a mark
 448 // of the cut bytes.
 449 func toolJoinClipped(head, tail string, cut int) string {
 450         return head + "\n[... " + strconv.Itoa(cut) + " bytes cut ...]\n" + tail
 451 }