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