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