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