git.giacomolli.uk - onechat/blob - turn.go

   1 package main
   2
   3 import (
   4         "context"
   5         "encoding/json"
   6         "errors"
   7         "fmt"
   8         "path/filepath"
   9         "strings"
  10         "sync"
  11         "time"
  12 )
  13
  14 // modelClient sends model calls.
  15 type modelClient interface {
  16         Create(ctx context.Context, request Request) (Response, error)
  17 }
  18
  19 // Event tells the clients what the agent does as it happens. Kind is the kind
  20 // of a logged message (user, reply, tool, echo), or one of busy, idle, and
  21 // error. ID is the ID of a logged message.
  22 type Event struct {
  23         Kind string `json:"kind"`
  24         ID   int    `json:"id,omitempty"`
  25         Text string `json:"text,omitempty"`
  26 }
  27
  28 // loop runs the turns of one agent. A user message goes to the log when it
  29 // comes, and it wakes the loop. A turn waits until every message before it is
  30 // summarized. Then its first request holds the prompt, the tools, and the
  31 // input of turnInput. The loop runs the tool calls of each response and sends
  32 // their results back, until a response has no tool call: the final answer.
  33 //
  34 // The view goes out as the messages of a viewLayout. After the turns, when no
  35 // message waits, a warm-up request sends only the prompt, the tools, and the
  36 // view messages, and stops at its first output item. It writes a cache entry
  37 // at the end of the view, where the next turn has a message end, so the first
  38 // request of that turn reads the view from the cache. The warm-up is not
  39 // logged, and its output is not used.
  40 //
  41 // Each later request of a turn begins with the whole input of the request
  42 // before it, then its output items, then the tool results. The output items
  43 // hold the reasoning with its encrypted content, because the calls do not
  44 // store it. So each request begins with all of the previous one and reads it
  45 // from the prompt cache. User messages that come during a turn join the next
  46 // request after the tool results. The replies, the tool calls, and the results
  47 // go to the log as they happen. The thoughts never do.
  48 type loop struct {
  49         agent  *Agent
  50         memory *Memory
  51         client modelClient
  52         tools  *toolbox
  53         layout *viewLayout
  54         warmUp bool
  55         clock  func() time.Time
  56         notify func(Event)
  57         wake   chan struct{}
  58
  59         mutex   sync.Mutex
  60         pending []Message
  61         cancel  context.CancelFunc
  62 }
  63
  64 // functionCall holds the fields of a function call output item.
  65 type functionCall struct {
  66         Type      string `json:"type"`
  67         Name      string `json:"name"`
  68         Arguments string `json:"arguments"`
  69         CallID    string `json:"call_id"`
  70 }
  71
  72 // functionCallOutput is the input item with the result of a function call.
  73 type functionCallOutput struct {
  74         Type   string `json:"type"`
  75         CallID string `json:"call_id"`
  76         Output string `json:"output"`
  77 }
  78
  79 // newLoop returns the loop of agent with its memory, with the warm-up on. Its
  80 // calls go through client, and notify gets each event. The clock gives the
  81 // date of a turn.
  82 func newLoop(agent *Agent, memory *Memory, client modelClient, clock func() time.Time, notify func(Event)) (*loop, error) {
  83         layout, err := loadLayout(filepath.Join(agent.dir, layoutFileName))
  84         if err != nil {
  85                 return nil, err
  86         }
  87
  88         l := &loop{
  89                 agent:  agent,
  90                 memory: memory,
  91                 client: client,
  92                 layout: layout,
  93                 warmUp: true,
  94                 clock:  clock,
  95                 notify: notify,
  96                 wake:   make(chan struct{}, 1),
  97         }
  98
  99         l.tools = &toolbox{
 100                 log:     l,
 101                 chat:    memory,
 102                 workdir: agent.config.Workdir,
 103         }
 104
 105         return l, nil
 106 }
 107
 108 // Submit logs a user message and wakes the loop. The message joins the next
 109 // request: the first one of a new turn, or the next one of a running turn.
 110 func (l *loop) Submit(text string) (Message, error) {
 111         message, err := l.Append("user", text)
 112         if err != nil {
 113                 return Message{}, err
 114         }
 115
 116         l.mutex.Lock()
 117         l.pending = append(l.pending, message)
 118         l.mutex.Unlock()
 119
 120         select {
 121         case l.wake <- struct{}{}:
 122         default:
 123         }
 124
 125         return message, nil
 126 }
 127
 128 // Interrupt ends the running turn. A running tool is killed, and the turn sends
 129 // no further request.
 130 func (l *loop) Interrupt() bool {
 131         l.mutex.Lock()
 132         defer l.mutex.Unlock()
 133
 134         if l.cancel == nil {
 135                 return false
 136         }
 137
 138         l.cancel()
 139
 140         return true
 141 }
 142
 143 // Append logs a message and tells the clients. The toolbox logs its calls and
 144 // results through it.
 145 func (l *loop) Append(kind, text string) (Message, error) {
 146         message, err := l.memory.Append(kind, text)
 147         if err != nil {
 148                 return Message{}, err
 149         }
 150
 151         l.notify(Event{
 152                 Kind: kind,
 153                 ID:   message.ID,
 154                 Text: text,
 155         })
 156
 157         return message, nil
 158 }
 159
 160 // Flush makes the log durable.
 161 func (l *loop) Flush(ctx context.Context) error {
 162         return l.memory.Flush(ctx)
 163 }
 164
 165 // Run runs a turn each time user messages wait, and a warm-up after the turns,
 166 // until ctx ends. It returns the error of a failed log write, which stops the
 167 // agent. A failed warm-up only costs the cache of the next turn, so its error
 168 // stays in the usage file.
 169 func (l *loop) Run(ctx context.Context) error {
 170         for {
 171                 turned := false
 172                 for l.waiting() {
 173                         err := l.turn(ctx)
 174                         if err != nil {
 175                                 return err
 176                         }
 177
 178                         turned = true
 179                 }
 180
 181                 if turned && l.warmUp {
 182                         _ = l.warm(ctx)
 183                 }
 184
 185                 select {
 186                 case <-l.wake:
 187                 case <-ctx.Done():
 188                         return ctx.Err()
 189                 }
 190         }
 191 }
 192
 193 // waiting tells if user messages wait for a request.
 194 func (l *loop) waiting() bool {
 195         l.mutex.Lock()
 196         defer l.mutex.Unlock()
 197
 198         return len(l.pending) > 0
 199 }
 200
 201 // take returns the user messages that wait for a request, and removes them.
 202 func (l *loop) take() []Message {
 203         l.mutex.Lock()
 204         defer l.mutex.Unlock()
 205
 206         messages := l.pending
 207         l.pending = nil
 208
 209         return messages
 210 }
 211
 212 // turn runs one turn. A failed or interrupted model call ends the turn and
 213 // tells the clients. The messages of the turn stay in the log and the view, and
 214 // the next turn starts only with a new message, so a lasting failure such as a
 215 // needed login does not repeat. Only a failed log write is an error of turn.
 216 func (l *loop) turn(ctx context.Context) error {
 217         ctx, cancel := context.WithCancel(ctx)
 218         defer cancel()
 219
 220         l.mutex.Lock()
 221         l.cancel = cancel
 222         l.mutex.Unlock()
 223
 224         defer func() {
 225                 l.mutex.Lock()
 226                 l.cancel = nil
 227                 l.mutex.Unlock()
 228         }()
 229
 230         l.notify(Event{
 231                 Kind: "busy",
 232         })
 233
 234         err := l.chain(ctx)
 235         if err != nil && !errors.Is(err, errLogWrite) {
 236                 l.notify(Event{
 237                         Kind: "error",
 238                         Text: turnError(ctx, err),
 239                 })
 240
 241                 err = nil
 242         }
 243
 244         l.notify(Event{
 245                 Kind: "idle",
 246         })
 247
 248         return err
 249 }
 250
 251 // errLogWrite marks an error of a log write, which ends the loop.
 252 var errLogWrite = errors.New("log write failed")
 253
 254 // chain runs the request chain of a turn until the final answer.
 255 func (l *loop) chain(ctx context.Context) error {
 256         messages := l.take()
 257         if len(messages) == 0 {
 258                 return nil
 259         }
 260
 261         err := l.memory.Wait(ctx)
 262         if err != nil {
 263                 return err
 264         }
 265
 266         instructions, err := l.agent.Prompt()
 267         if err != nil {
 268                 return err
 269         }
 270
 271         skip := make(map[int]bool, len(messages))
 272         for _, message := range messages {
 273                 skip[message.ID] = true
 274         }
 275
 276         view, err := l.layout.Split(l.memory.ViewLines(skip))
 277         if err != nil {
 278                 return err
 279         }
 280
 281         input := turnInput(view, l.state(), messages)
 282         for turn := "first"; ; turn = "next" {
 283                 response, err := l.client.Create(ctx, Request{
 284                         Model:        l.agent.config.Model,
 285                         Effort:       l.agent.config.Effort,
 286                         Instructions: instructions,
 287                         Tools:        agentTools(),
 288                         Input:        input,
 289                         Turn:         turn,
 290                 })
 291                 if err != nil {
 292                         return err
 293                 }
 294
 295                 input, err = l.handle(ctx, input, response)
 296                 if err != nil {
 297                         return err
 298                 }
 299
 300                 if len(input) == 0 {
 301                         return nil
 302                 }
 303         }
 304 }
 305
 306 // handle logs the replies of response and runs its tool calls. It returns the
 307 // input of the next request: input, the output items, the tool results, and
 308 // the user messages that came in the meantime. After the final answer it
 309 // returns no input.
 310 func (l *loop) handle(ctx context.Context, input []any, response Response) ([]any, error) {
 311         var outputs []any
 312         for _, raw := range response.Output {
 313                 item, err := inputItem(raw)
 314                 if err != nil {
 315                         return nil, err
 316                 }
 317
 318                 input = append(input, item)
 319
 320                 var call functionCall
 321                 err = json.Unmarshal(raw, &call)
 322                 if err != nil {
 323                         return nil, fmt.Errorf("decode output item: %w", err)
 324                 }
 325
 326                 switch call.Type {
 327                 case "message":
 328                         text, err := outputText([]json.RawMessage{raw})
 329                         if err != nil {
 330                                 return nil, err
 331                         }
 332
 333                         text = strings.TrimSpace(text)
 334                         if text == "" {
 335                                 continue
 336                         }
 337
 338                         _, err = l.Append("reply", text)
 339                         if err != nil {
 340                                 return nil, fmt.Errorf("%w: %w", errLogWrite, err)
 341                         }
 342                 case "function_call":
 343                         output, err := l.tools.Call(ctx, call.Name, call.Arguments)
 344                         if err != nil {
 345                                 return nil, fmt.Errorf("%w: %w", errLogWrite, err)
 346                         }
 347
 348                         outputs = append(outputs, functionCallOutput{
 349                                 Type:   "function_call_output",
 350                                 CallID: call.CallID,
 351                                 Output: output,
 352                         })
 353                 }
 354         }
 355
 356         if len(outputs) == 0 {
 357                 return nil, nil
 358         }
 359
 360         input = append(input, outputs...)
 361         for _, message := range l.take() {
 362                 input = append(input, userMessage(message.Text))
 363         }
 364
 365         return input, nil
 366 }
 367
 368 // warm sends the warm-up request, when the builds are done and no message
 369 // waits.
 370 func (l *loop) warm(ctx context.Context) error {
 371         err := l.memory.Wait(ctx)
 372         if err != nil || l.waiting() {
 373                 return err
 374         }
 375
 376         instructions, err := l.agent.Prompt()
 377         if err != nil {
 378                 return err
 379         }
 380
 381         view, err := l.layout.Split(l.memory.ViewLines(nil))
 382         if err != nil {
 383                 return err
 384         }
 385
 386         _, err = l.client.Create(ctx, Request{
 387                 Model:        l.agent.config.Model,
 388                 Effort:       l.agent.config.Effort,
 389                 Instructions: instructions,
 390                 Tools:        agentTools(),
 391                 Input:        warmUpInput(view),
 392                 Turn:         "warmup",
 393                 StopAtOutput: true,
 394         })
 395
 396         return err
 397 }
 398
 399 // state returns the per-turn state: the end of the view, and the date and time
 400 // of the turn.
 401 func (l *loop) state() string {
 402         return "</chat>\n\nNow: " + l.clock().Format("Monday 2006-01-02 15:04 -07:00")
 403 }
 404
 405 // turnInput returns the input of the first request of a turn, stable parts
 406 // first: a user message for each view message, one with the per-turn state,
 407 // and one for each new message, whole. The view leaves out the leaf lines of
 408 // the new messages.
 409 func turnInput(view []string, state string, messages []Message) []any {
 410         input := warmUpInput(view)
 411         input = append(input, userMessage(state))
 412         for _, message := range messages {
 413                 input = append(input, userMessage(message.Text))
 414         }
 415
 416         return input
 417 }
 418
 419 // warmUpInput returns a user message for each view message.
 420 func warmUpInput(view []string) []any {
 421         input := make([]any, len(view))
 422         for n, text := range view {
 423                 input[n] = userMessage(text)
 424         }
 425
 426         return input
 427 }
 428
 429 // inputItem returns an output item as an input item for the next request. The
 430 // calls do not store their items, so the item goes back whole, without its id.
 431 func inputItem(raw json.RawMessage) (map[string]any, error) {
 432         var item map[string]any
 433         err := json.Unmarshal(raw, &item)
 434         if err != nil {
 435                 return nil, fmt.Errorf("decode output item: %w", err)
 436         }
 437
 438         delete(item, "id")
 439
 440         return item, nil
 441 }
 442
 443 // turnError returns the text of the error that ended a turn.
 444 func turnError(ctx context.Context, err error) string {
 445         if ctx.Err() != nil {
 446                 return "interrupted"
 447         }
 448
 449         return err.Error()
 450 }