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