]> git.giacomolli.uk - onechat/blob - turn.go
Give each subagent its own prompt cache key
[onechat] / turn.go
1 package main
2
3 import (
4         "context"
5         "errors"
6         "fmt"
7         "path/filepath"
8         "sync"
9         "time"
10 )
11
12 // newTurnLoop returns the loop of agent with its memory and its subagents. Its
13 // calls go through client, and notify gets each event. The clock gives the date
14 // of a turn.
15 func newTurnLoop(agent *Agent, memory *Memory, subagents *Subagents, client ChainClient, clock func() time.Time, notify func(Event)) (*TurnLoop, error) {
16         layout, err := loadViewLayout(filepath.Join(agent.dir, viewLayoutFileName))
17         if err != nil {
18                 return nil, err
19         }
20
21         l := &TurnLoop{
22                 agent:  agent,
23                 memory: memory,
24                 client: client,
25                 layout: layout,
26                 clock:  clock,
27                 notify: notify,
28                 wake:   make(chan struct{}, 1),
29         }
30
31         l.runs = &SubagentRuns{
32                 agent:     agent,
33                 memory:    memory,
34                 subagents: subagents,
35                 layout:    layout,
36                 client:    client,
37                 state:     l.state,
38                 report:    l.Report,
39                 note:      l.note,
40                 fail:      l.fail,
41                 notify:    notify,
42                 running:   make(map[string]*SubagentRun),
43         }
44
45         l.tools = &Toolbox{
46                 log:       l,
47                 chat:      memory,
48                 subagents: subagents,
49                 runs:      l.runs,
50                 workdir:   agent.config.Workdir,
51         }
52
53         return l, nil
54 }
55
56 // Event tells the clients what the agent does as it happens. Kind is the kind
57 // of a logged message (user, reply, tool, echo, note, work), or one of ready,
58 // busy, idle, error, and subagent, which tells that the subagent Text starts a
59 // run. ID is the ID of a logged message. Source names the client that sent a
60 // user message, such as terminal or telegram, so that each client can show the
61 // messages of the others, or the subagent of a report. Status tells how the run
62 // of a report ended: finished, stopped, or failed.
63 type Event struct {
64         Kind   string `json:"kind"`
65         ID     int    `json:"id,omitempty"`
66         Text   string `json:"text,omitempty"`
67         Source string `json:"source,omitempty"`
68         Status string `json:"status,omitempty"`
69 }
70
71 // TurnLoop runs the turns of one agent. A user message goes to the log when it
72 // comes, and it wakes the loop. A turn waits until every message before it is
73 // summarized. Then the turn runs a Chain whose first request holds the prompt,
74 // the tools, and the input of turnInput.
75 //
76 // The view goes out as the messages of a ViewLayout, so the first request of a
77 // turn reads the view of the turn before from the prompt cache. User messages
78 // that come during a turn join the next request of its chain after the tool
79 // results, and so do the reports of the subagents. The subagents run until the
80 // loop ends.
81 type TurnLoop struct {
82         agent  *Agent
83         memory *Memory
84         client ChainClient
85         tools  *Toolbox
86         runs   *SubagentRuns
87         layout *ViewLayout
88         clock  func() time.Time
89         notify func(Event)
90         wake   chan struct{}
91
92         mutex   sync.Mutex
93         pending []Message
94         cancel  context.CancelFunc
95         failure error
96 }
97
98 // Submit logs a user message from the client source and wakes the loop. The
99 // message joins the next request: the first one of a new turn, or the next one
100 // of a running turn.
101 func (l *TurnLoop) Submit(source, text string) (Message, error) {
102         message, err := l.record("user", source, text)
103         if err != nil {
104                 return Message{}, err
105         }
106
107         l.mutex.Lock()
108         l.pending = append(l.pending, message)
109         l.mutex.Unlock()
110
111         select {
112         case l.wake <- struct{}{}:
113         default:
114         }
115
116         return message, nil
117 }
118
119 // Report logs the report text of the subagent name as a work message and tells
120 // the clients, with the status of the run. Unless the subagent was stopped, the
121 // report joins the waiting messages and wakes the loop, as a user message of
122 // Submit does. A failed write stops the loop, because the log refuses all
123 // writes after it.
124 func (l *TurnLoop) Report(name, text, status string) {
125         message, err := l.memory.Append("work", text)
126         if err == nil {
127                 l.notify(Event{
128                         Kind:   "work",
129                         ID:     message.ID,
130                         Text:   text,
131                         Source: name,
132                         Status: status,
133                 })
134         }
135
136         l.mutex.Lock()
137         switch {
138         case err != nil:
139                 l.failure = fmt.Errorf("%w: %w", errChainLogWrite, err)
140         case status != subagentStopped:
141                 l.pending = append(l.pending, message)
142         }
143         l.mutex.Unlock()
144
145         select {
146         case l.wake <- struct{}{}:
147         default:
148         }
149 }
150
151 // StopSubagent ends the run of the subagent name and tells if it ran. Its
152 // report is logged, but it starts no turn.
153 func (l *TurnLoop) StopSubagent(name string) (bool, error) {
154         return l.runs.Stop(name)
155 }
156
157 // Interrupt ends the running turn. A running tool is killed, and the turn sends
158 // no further request.
159 func (l *TurnLoop) Interrupt() bool {
160         l.mutex.Lock()
161         defer l.mutex.Unlock()
162
163         if l.cancel == nil {
164                 return false
165         }
166
167         l.cancel()
168
169         return true
170 }
171
172 // Busy tells if a turn runs.
173 func (l *TurnLoop) Busy() bool {
174         l.mutex.Lock()
175         defer l.mutex.Unlock()
176
177         return l.cancel != nil
178 }
179
180 // Append logs a message and tells the clients. The toolbox logs its calls and
181 // results through it.
182 func (l *TurnLoop) Append(kind, text string) (Message, error) {
183         return l.record(kind, "", text)
184 }
185
186 // Flush makes the log durable.
187 func (l *TurnLoop) Flush(ctx context.Context) error {
188         return l.memory.Flush(ctx)
189 }
190
191 // Run runs a turn each time user messages or reports wait, until ctx ends. It
192 // returns the error of a failed log write, which stops the agent. At its start
193 // it resumes the subagents that a stop of the daemon cut, and at its end it
194 // ends the runs of the subagents and waits for them.
195 func (l *TurnLoop) Run(ctx context.Context) error {
196         defer l.runs.Close()
197
198         err := l.runs.Resume(ctx)
199         if err != nil {
200                 return fmt.Errorf("resume subagents: %w", err)
201         }
202
203         for {
204                 for l.waiting() {
205                         err := l.turn(ctx)
206                         if err != nil {
207                                 return err
208                         }
209                 }
210
211                 select {
212                 case <-l.wake:
213                 case <-ctx.Done():
214                         return ctx.Err()
215                 }
216
217                 err := l.failed()
218                 if err != nil {
219                         return err
220                 }
221         }
222 }
223
224 // note logs a note and tells the clients. A note starts no turn.
225 func (l *TurnLoop) note(text string) error {
226         _, err := l.record("note", "", text)
227
228         return err
229 }
230
231 // record logs a message and tells the clients, with the client source of a user
232 // message.
233 func (l *TurnLoop) record(kind, source, text string) (Message, error) {
234         message, err := l.memory.Append(kind, text)
235         if err != nil {
236                 return Message{}, err
237         }
238
239         l.notify(Event{
240                 Kind:   kind,
241                 ID:     message.ID,
242                 Text:   text,
243                 Source: source,
244         })
245
246         return message, nil
247 }
248
249 // waiting tells if user messages or reports wait for a request.
250 func (l *TurnLoop) waiting() bool {
251         l.mutex.Lock()
252         defer l.mutex.Unlock()
253
254         return len(l.pending) > 0
255 }
256
257 // fail stops the loop with err.
258 func (l *TurnLoop) fail(err error) {
259         l.mutex.Lock()
260         l.failure = errors.Join(l.failure, err)
261         l.mutex.Unlock()
262
263         select {
264         case l.wake <- struct{}{}:
265         default:
266         }
267 }
268
269 // failed returns the error of a failed write of a report.
270 func (l *TurnLoop) failed() error {
271         l.mutex.Lock()
272         defer l.mutex.Unlock()
273
274         return l.failure
275 }
276
277 // turn runs one turn. A failed or interrupted model call ends the turn and
278 // tells the clients. The messages of the turn stay in the log and the view, and
279 // the next turn starts only with a new message, so a lasting failure such as a
280 // needed login does not repeat. Only a failed log write is an error of turn.
281 func (l *TurnLoop) turn(ctx context.Context) error {
282         ctx, cancel := context.WithCancel(ctx)
283         defer cancel()
284
285         l.mutex.Lock()
286         l.cancel = cancel
287         l.mutex.Unlock()
288
289         defer func() {
290                 l.mutex.Lock()
291                 l.cancel = nil
292                 l.mutex.Unlock()
293         }()
294
295         l.notify(Event{
296                 Kind: "busy",
297         })
298
299         err := l.chain(ctx)
300         if err != nil && !errors.Is(err, errChainLogWrite) {
301                 l.notify(Event{
302                         Kind: "error",
303                         Text: turnError(ctx, err),
304                 })
305
306                 err = nil
307         }
308
309         l.notify(Event{
310                 Kind: "idle",
311         })
312
313         return err
314 }
315
316 // chain runs the request chain of a turn until the final answer.
317 func (l *TurnLoop) chain(ctx context.Context) error {
318         messages := l.take()
319         if len(messages) == 0 {
320                 return nil
321         }
322
323         err := l.memory.Wait(ctx)
324         if err != nil {
325                 return err
326         }
327
328         instructions, err := l.agent.Prompt()
329         if err != nil {
330                 return err
331         }
332
333         skip := make(map[int]bool, len(messages))
334         for _, message := range messages {
335                 skip[message.ID] = true
336         }
337
338         view, err := l.layout.Split(l.memory.ViewLines(skip))
339         if err != nil {
340                 return err
341         }
342
343         chain := &Chain{
344                 client: l.client,
345                 tools:  l.tools,
346                 request: ModelRequest{
347                         Model:        l.agent.config.Model,
348                         Effort:       l.agent.config.Effort,
349                         Instructions: instructions,
350                         Tools:        promptTools(),
351                 },
352                 take: l.take,
353         }
354
355         _, err = chain.Run(ctx, turnInput(view, l.state(), messages))
356
357         return err
358 }
359
360 // take returns the user messages that wait for a request, and removes them.
361 func (l *TurnLoop) take() []Message {
362         l.mutex.Lock()
363         defer l.mutex.Unlock()
364
365         messages := l.pending
366         l.pending = nil
367
368         return messages
369 }
370
371 // state returns the per-turn state: the end of the view, and the date and time
372 // of the turn.
373 func (l *TurnLoop) state() string {
374         return "</chat>\n\nNow: " + l.clock().Format("Monday 2006-01-02 15:04 -07:00")
375 }
376
377 // turnError returns the text of the error that ended a turn.
378 func turnError(ctx context.Context, err error) string {
379         if ctx.Err() != nil {
380                 return "interrupted"
381         }
382
383         return err.Error()
384 }
385
386 // turnInput returns the input of the first request of a turn, stable parts
387 // first: a user message for each view message, one with the per-turn state,
388 // and one for each new message, whole. The view leaves out the leaf lines of
389 // the new messages.
390 func turnInput(view []string, state string, messages []Message) []any {
391         input := make([]any, 0, len(view)+1+len(messages))
392         for _, text := range view {
393                 input = append(input, modelUserMessage(text))
394         }
395
396         input = append(input, modelUserMessage(state))
397         for _, message := range messages {
398                 input = append(input, modelUserMessage(message.Text))
399         }
400
401         return input
402 }