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         // clipCharacters is the most characters of a tool output: its head and
  23         // its tail, half each.
  24         clipCharacters = 30000
  25
  26         // clipBytes is the most bytes that a tool keeps of each end of an
  27         // output, enough for clipCharacters/2 characters of 4 bytes.
  28         clipBytes = 4 * clipCharacters / 2
  29
  30         // shellWaitDelay is how long a command can keep its output open after it
  31         // ends or after its kill.
  32         shellWaitDelay = 2 * time.Second
  33 )
  34
  35 // chatLog is the log that the tools write their calls and results to.
  36 type chatLog interface {
  37         Append(kind, text string) (Message, error)
  38         Flush(ctx context.Context) error
  39 }
  40
  41 // chatReader gives the messages and the nodes that zoom and date read.
  42 type chatReader 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
  51 // an action outside the agent. Commands and files run as the user in the
  52 // working directory, with no sandbox.
  53 type toolbox struct {
  54         log     chatLog
  55         chat    chatReader
  56         workdir string
  57 }
  58
  59 // shellArguments are the arguments of shell.
  60 type shellArguments struct {
  61         Command string `json:"command"`
  62 }
  63
  64 // pathArguments are the arguments of read.
  65 type pathArguments struct {
  66         Path string `json:"path"`
  67 }
  68
  69 // writeArguments are the arguments of write.
  70 type writeArguments struct {
  71         Path string `json:"path"`
  72         Text string `json:"text"`
  73 }
  74
  75 // editArguments are the arguments of edit.
  76 type editArguments struct {
  77         Path string `json:"path"`
  78         Old  string `json:"old"`
  79         New  string `json:"new"`
  80 }
  81
  82 // zoomArguments are the arguments of zoom.
  83 type zoomArguments struct {
  84         ID int `json:"id"`
  85         N  int `json:"n"`
  86 }
  87
  88 // dateArguments are the arguments of date.
  89 type dateArguments struct {
  90         ID int `json:"id"`
  91 }
  92
  93 // Call logs the call of the tool name with its JSON arguments, makes the log
  94 // durable, runs the tool, and logs its clipped output, which it returns. A
  95 // failed tool gives its error as the output. An error of Call itself is an
  96 // error of the log. When ctx ends before the flush, the tool does not run.
  97 func (b *toolbox) Call(ctx context.Context, name, arguments string) (string, error) {
  98         _, err := b.log.Append("tool", callText(name, arguments))
  99         if err != nil {
 100                 return "", err
 101         }
 102
 103         err = b.log.Flush(ctx)
 104         if err != nil && ctx.Err() == nil {
 105                 return "", err
 106         }
 107
 108         output := "not run: the turn was interrupted"
 109         if err == nil {
 110                 output = clip(b.run(ctx, name, arguments))
 111         }
 112
 113         _, err = b.log.Append("echo", output)
 114         if err != nil {
 115                 return "", err
 116         }
 117
 118         return output, nil
 119 }
 120
 121 // run runs one tool and returns its output, or its error as text.
 122 func (b *toolbox) run(ctx context.Context, name, arguments string) string {
 123         var output string
 124         var err error
 125         switch name {
 126         case "zoom":
 127                 var args zoomArguments
 128                 err = decodeArguments(arguments, &args)
 129                 if err == nil {
 130                         output, err = b.zoom(args.ID, args.N)
 131                 }
 132         case "date":
 133                 var args dateArguments
 134                 err = decodeArguments(arguments, &args)
 135                 if err == nil {
 136                         output, err = b.date(args.ID)
 137                 }
 138         case "shell":
 139                 var args shellArguments
 140                 err = decodeArguments(arguments, &args)
 141                 if err == nil {
 142                         output = b.shell(ctx, args.Command)
 143                 }
 144         case "read":
 145                 var args pathArguments
 146                 err = decodeArguments(arguments, &args)
 147                 if err == nil {
 148                         output, err = b.read(args.Path)
 149                 }
 150         case "write":
 151                 var args writeArguments
 152                 err = decodeArguments(arguments, &args)
 153                 if err == nil {
 154                         output, err = b.write(args.Path, args.Text)
 155                 }
 156         case "edit":
 157                 var args editArguments
 158                 err = decodeArguments(arguments, &args)
 159                 if err == nil {
 160                         output, err = b.edit(args.Path, args.Old, args.New)
 161                 }
 162         default:
 163                 err = fmt.Errorf("unknown tool %q", name)
 164         }
 165
 166         if err != nil {
 167                 return "error: " + err.Error()
 168         }
 169
 170         return output
 171 }
 172
 173 // zoom returns message id whole for n = 1, or else the two lines under the line
 174 // id+n.
 175 func (b *toolbox) zoom(id, n int) (string, error) {
 176         if n < 1 || bits.OnesCount(uint(n)) != 1 || id < 0 || id%n != 0 {
 177                 return "", fmt.Errorf("line %d+%d does not exist: n must be a power of 2 and id a multiple of n", id, n)
 178         }
 179
 180         if id+n > b.chat.Len() {
 181                 return "", fmt.Errorf("line %d+%d goes past the last message, %d", id, n, b.chat.Len()-1)
 182         }
 183
 184         if n == 1 {
 185                 message, err := b.chat.Message(id)
 186                 if err != nil {
 187                         return "", err
 188                 }
 189
 190                 return leafSource(message), nil
 191         }
 192
 193         level := bits.TrailingZeros(uint(n)) - 1
 194         var lines []string
 195         for _, index := range []int{2 * id / n, 2*id/n + 1} {
 196                 node, built, err := b.chat.Node(level, index)
 197                 if err != nil {
 198                         return "", err
 199                 }
 200
 201                 text := pending
 202                 if built {
 203                         text = strings.ReplaceAll(node.Text, "\n", " ")
 204                 }
 205
 206                 lines = append(lines, fmt.Sprintf("%d+%d|%s", index<<level, 1<<level, text))
 207         }
 208
 209         return strings.Join(lines, "\n"), nil
 210 }
 211
 212 // date returns the date and time of message id, with its weekday.
 213 func (b *toolbox) date(id int) (string, error) {
 214         if id < 0 || id >= b.chat.Len() {
 215                 return "", fmt.Errorf("message %d does not exist", id)
 216         }
 217
 218         message, err := b.chat.Message(id)
 219         if err != nil {
 220                 return "", err
 221         }
 222
 223         date, err := time.Parse(time.RFC3339Nano, message.Date)
 224         if err != nil {
 225                 return "", err
 226         }
 227
 228         return date.Local().Format("Monday 2006-01-02 15:04:05 -07:00"), nil
 229 }
 230
 231 // shell runs command with bash in the working directory and returns its output
 232 // and error output. A failed command adds its exit status. When ctx ends, the
 233 // whole process group of the command is killed.
 234 func (b *toolbox) shell(ctx context.Context, command string) string {
 235         cmd := exec.CommandContext(ctx, "bash", "-c", command)
 236         cmd.Dir = b.workdir
 237         cmd.SysProcAttr = &syscall.SysProcAttr{
 238                 Setpgid: true,
 239         }
 240         cmd.Cancel = func() error {
 241                 return syscall.Kill(-cmd.Process.Pid, syscall.SIGKILL)
 242         }
 243         cmd.WaitDelay = shellWaitDelay
 244
 245         output := &headTail{}
 246         cmd.Stdout = output
 247         cmd.Stderr = output
 248         err := cmd.Run()
 249
 250         text := output.String()
 251         var exit *exec.ExitError
 252         switch {
 253         case ctx.Err() != nil:
 254                 text += "\n[interrupted]"
 255         case errors.As(err, &exit):
 256                 text += "\n[" + exit.String() + "]"
 257         case err != nil:
 258                 text += "\n[error: " + err.Error() + "]"
 259         case text == "":
 260                 text = "(no output)"
 261         }
 262
 263         return text
 264 }
 265
 266 // read returns the text of the file at path.
 267 func (b *toolbox) read(path string) (text string, err error) {
 268         file, err := os.Open(b.path(path))
 269         if err != nil {
 270                 return "", err
 271         }
 272
 273         defer func() {
 274                 err = errors.Join(err, file.Close())
 275         }()
 276
 277         output := &headTail{}
 278         _, err = io.Copy(output, file)
 279         if err != nil {
 280                 return "", err
 281         }
 282
 283         return output.String(), nil
 284 }
 285
 286 // write replaces the file at path with text, and makes its missing
 287 // directories.
 288 func (b *toolbox) write(path, text string) (string, error) {
 289         path = b.path(path)
 290         err := os.MkdirAll(filepath.Dir(path), 0o755)
 291         if err != nil {
 292                 return "", err
 293         }
 294
 295         err = os.WriteFile(path, []byte(text), 0o644)
 296         if err != nil {
 297                 return "", err
 298         }
 299
 300         return fmt.Sprintf("wrote %d bytes to %s", len(text), path), nil
 301 }
 302
 303 // edit replaces the one occurrence of old in the file at path with
 304 // replacement. It keeps the mode of the file.
 305 func (b *toolbox) edit(path, old, replacement string) (string, error) {
 306         path = b.path(path)
 307         data, err := os.ReadFile(path)
 308         if err != nil {
 309                 return "", err
 310         }
 311
 312         if old == "" {
 313                 return "", errors.New("old is empty")
 314         }
 315
 316         count := strings.Count(string(data), old)
 317         if count != 1 {
 318                 return "", fmt.Errorf("old occurs %d times in %s, not once", count, path)
 319         }
 320
 321         info, err := os.Stat(path)
 322         if err != nil {
 323                 return "", err
 324         }
 325
 326         err = os.WriteFile(path, []byte(strings.Replace(string(data), old, replacement, 1)), info.Mode().Perm())
 327         if err != nil {
 328                 return "", err
 329         }
 330
 331         return "edited " + path, nil
 332 }
 333
 334 // path returns path, relative to the working directory when it is not
 335 // absolute.
 336 func (b *toolbox) path(path string) string {
 337         if filepath.IsAbs(path) {
 338                 return path
 339         }
 340
 341         return filepath.Join(b.workdir, path)
 342 }
 343
 344 // headTail keeps the first and the last clipBytes bytes written to it, and
 345 // counts the bytes between them that it drops.
 346 type headTail struct {
 347         head    []byte
 348         tail    []byte
 349         dropped int
 350 }
 351
 352 // Write keeps the bytes of p at the ends of the output.
 353 func (h *headTail) Write(p []byte) (int, error) {
 354         n := len(p)
 355         room := min(clipBytes-len(h.head), len(p))
 356         h.head = append(h.head, p[:room]...)
 357         p = p[room:]
 358
 359         h.tail = append(h.tail, p...)
 360         if len(h.tail) > clipBytes {
 361                 cut := len(h.tail) - clipBytes
 362                 h.dropped += cut
 363                 h.tail = bytes.Clone(h.tail[cut:])
 364         }
 365
 366         return n, nil
 367 }
 368
 369 // String returns the output, clipped when it is too long. Invalid UTF-8 becomes
 370 // the replacement character.
 371 func (h *headTail) String() string {
 372         replacement := string(utf8.RuneError)
 373         if h.dropped == 0 {
 374                 return clip(strings.ToValidUTF8(string(h.head)+string(h.tail), replacement))
 375         }
 376
 377         head := firstCharacters(strings.ToValidUTF8(string(h.head), replacement), clipCharacters/2)
 378         tail := lastCharacters(strings.ToValidUTF8(string(h.tail), replacement), clipCharacters/2)
 379         cut := len(h.head) + h.dropped + len(h.tail) - len(head) - len(tail)
 380
 381         return joinClipped(head, tail, cut)
 382 }
 383
 384 // callText returns the log text of a tool call: the name and the JSON
 385 // arguments.
 386 func callText(name, arguments string) string {
 387         return name + " " + arguments
 388 }
 389
 390 // decodeArguments decodes the JSON arguments of a call into args. An unknown
 391 // field is an error.
 392 func decodeArguments(arguments string, args any) error {
 393         decoder := json.NewDecoder(strings.NewReader(arguments))
 394         decoder.DisallowUnknownFields()
 395         err := decoder.Decode(args)
 396         if err != nil {
 397                 return fmt.Errorf("arguments: %w", err)
 398         }
 399
 400         return nil
 401 }
 402
 403 // clip returns text when it has at most clipCharacters characters. Otherwise it
 404 // returns the first and the last clipCharacters/2 characters, with a mark of
 405 // the cut between them.
 406 func clip(text string) string {
 407         if utf8.RuneCountInString(text) <= clipCharacters {
 408                 return text
 409         }
 410
 411         head := firstCharacters(text, clipCharacters/2)
 412         tail := lastCharacters(text, clipCharacters/2)
 413
 414         return joinClipped(head, tail, len(text)-len(head)-len(tail))
 415 }
 416
 417 // joinClipped joins the head and the tail of a clipped output with a mark of
 418 // the cut bytes.
 419 func joinClipped(head, tail string, cut int) string {
 420         return head + "\n[... " + strconv.Itoa(cut) + " bytes cut ...]\n" + tail
 421 }
 422
 423 // firstCharacters returns the first count characters of text.
 424 func firstCharacters(text string, count int) string {
 425         end := 0
 426         for range count {
 427                 if end >= len(text) {
 428                         break
 429                 }
 430
 431                 _, size := utf8.DecodeRuneInString(text[end:])
 432                 end += size
 433         }
 434
 435         return text[:end]
 436 }
 437
 438 // lastCharacters returns the last count characters of text.
 439 func lastCharacters(text string, count int) string {
 440         start := len(text)
 441         for range count {
 442                 if start <= 0 {
 443                         break
 444                 }
 445
 446                 _, size := utf8.DecodeLastRuneInString(text[:start])
 447                 start -= size
 448         }
 449
 450         return text[start:]
 451 }