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