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