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