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

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