bots-garden/mini-mepublic Fork 0
main
Commits
Clone
git clone https://git.rickub.com/bots-garden/mini-me.git
git clone ssh://git@rickub.com/bots-garden/mini-me.git

Host key fingerprint (ed25519): SHA256:iycHnxEyq0Q7uyVpB7JlznP0G7JrTPXLYRcAU5CSLhc — verify it before your first connect.

engine.go · 622 lines · 23.5 KBGo Blame HistoryRaw
💾 Saved. d722711 k33g 4h ago1// Package engine wraps the connection to the LLM engine — Docker Model Runner,
2// llama-server, or any other Provider of the registry — and the generation of
3// answers. Everything here is engine-agnostic: the only lines that know WHICH
4// server is behind the URL live in the Provider implementations (see
5// provider.go and openai_compat.go). This package used to be called `dmr`;
6// measured before the split, exactly one function (Init) and two string
7// constants were specific to Docker Model Runner. The rest — the whitespace
8// gap, the watchdog, the retry, the command counter — already behaved the same
9// against the fake engine, which is not DMR either.
10package engine
11
12import (
13 "context"
14 "encoding/json"
15 "fmt"
16 "maps"
17 "slices"
18 "strings"
19 "sync"
20 "sync/atomic"
21 "time"
22
23 "mm/internal/config"
24 "mm/internal/spinner"
25 "mm/internal/ui"
26
27 "github.com/firebase/genkit/go/ai"
28 "github.com/firebase/genkit/go/genkit"
29)
30
31// Engine is what main.go hands to the agent and to the tools: an initialised
32// Genkit, the FULL model reference Genkit expects ("dmr/<id>",
33// "llamacpp/<id>"), and the provider that built it — kept around for the one
34// thing that stays provider-specific after start-up: explaining an error in one
35// line (see Explain).
36type Engine struct {
37 G *genkit.Genkit
38 Model string
39 Provider Provider
40 Backend Backend
41
42 // ContextWindow is the served context size as far as the agent knows it:
43 // the config's hint, or what Probe learned from the server (llama-server's
44 // /props). 0 = unknown. The context compression (internal/compact) reads
45 // it; the banner shows it with its origin, because a number without an
46 // origin never gets corrected.
47 ContextWindow int
48
49 // lastInput is the input-token count the server reported on the most
50 // recent call to the model, captured in Generate's WrapModel hook. It is
51 // the exact figure the compression trigger prefers over its estimate.
52 lastInput atomic.Int64
53}
54
55// New reads the provider named in the config, resolves its backend (URL, key,
56// model) and opens it. Unknown providers are refused HERE, with the list of
57// known names: a typo in `provider:` must not turn into a "model not found"
58// twenty seconds later, once the model has tried to load.
59func New(ctx context.Context) (*Engine, error) {
60 p, err := Lookup(config.Cfg.Provider)
61 if err != nil {
62 return nil, err
63 }
64 b, err := p.Resolve(config.Cfg)
65 if err != nil {
66 return nil, err
67 }
68 g, model, err := p.Open(ctx, b)
69 if err != nil {
70 return nil, err
71 }
72 return &Engine{G: g, Model: model, Provider: p, Backend: b, ContextWindow: b.ContextWindow}, nil
73}
74
75// Probe asks the provider what it can learn about the server BEFORE the first
76// question: reachability, the context window it serves, a model that is not
77// pulled. Never fatal — the user may well start the server after the agent.
78//
79// The context window it learns is remembered on the Engine. The provider has
80// already applied the precedence — the config's value wins over the server's,
81// because the operator knows what they started the server with — so what is
82// stored here is the one number the whole agent reasons with.
83func (e *Engine) Probe(ctx context.Context) Info {
84 info := e.Provider.Probe(ctx, e.Backend)
85 if info.ContextWindow > 0 {
86 e.ContextWindow = info.ContextWindow
87 }
88 return info
89}
90
91// EnsureContextWindow re-runs the provider's probe when the window is still
92// unknown, and reports what it learned: the window and its origin ("/props"),
93// or 0 and "" when the server still tells nothing.
94//
95// Why: the start-up probe runs once (main.go), and on a demo machine the
96// server is often started AFTER the agent — observed: bob started before
97// llama-server left `ctx: unknown` for the whole session, so the compression
98// could only ever trigger on maxMessages. Asking again, lazily, fixes that
99// without a background poller. Cheap by construction: the caller only asks
100// while the window is unknown, and a known window is returned as-is, without
101// a request — a `/props` hit per question would be noise the server does not
102// deserve. Warnings are deliberately dropped here: the start-up probe already
103// printed them, and repeating "nothing answers" at every question is nagging.
104func (e *Engine) EnsureContextWindow(ctx context.Context) (window int, source string) {
105 if e.ContextWindow > 0 {
106 return e.ContextWindow, ""
107 }
108 info := e.Provider.Probe(ctx, e.Backend)
109 if info.ContextWindow > 0 {
110 e.ContextWindow = info.ContextWindow
111 return info.ContextWindow, info.ContextSource
112 }
113 return 0, ""
114}
115
116// LastInputTokens returns the server's own count of the context it last read,
117// or 0 when it reports no usage (the fake engine of part 03, some proxies).
118func (e *Engine) LastInputTokens() int { return int(e.lastInput.Load()) }
119
120// ForgetInputTokens clears that count. Called after a compression: the measure
121// described the history that was just replaced, and keeping it would trigger a
122// second compression on the very next question.
123func (e *Engine) ForgetInputTokens() { e.lastInput.Store(0) }
124
125// Explain turns a transport error into the one-line, actionable message the
126// REPL prints between brackets. See Provider.Explain for the rationale.
127func (e *Engine) Explain(err error) string {
128 return e.Provider.Explain(err, e.Backend)
129}
130
131// blanks is the set of characters the model adds around its text without any of
132// it showing on screen.
133const blanks = " \t\r\n"
134
135// gap holds back the TRAILING whitespace of the streamed text: it is only
136// printed if visible text follows. Without that, the line breaks the model puts
137// after its sentence are printed as-is, and the screen hollows out.
138//
139// The model produces MORE AND MORE of it as the turns go by — it imitates its
140// own messages, which it re-reads in the history every turn. Measured: 2 blank
141// lines before the first 🛠️, 4 before the second, 7, then 11. Holding the
142// whitespace back is enough to remove all of it, since that whitespace is never
143// followed by text: it is a tool that writes next.
144type gap struct {
145 // The `bash` tool may call drop() from SEVERAL goroutines — Genkit runs the
146 // calls of a single turn in parallel (see printMu in tools).
147 mu sync.Mutex
148 held string // whitespace waiting to be followed by text
149 started bool // visible text has already been printed
150}
151
152// next returns what should be printed for this chunk, "" when there is nothing
153// to print right now.
154func (g *gap) next(text string) string {
155 g.mu.Lock()
156 defer g.mu.Unlock()
157
158 body := strings.TrimRight(text, blanks)
159 if body == "" { // an all-whitespace chunk: set aside, not printed
160 if g.started {
161 g.held += text
162 }
163 return ""
164 }
165
166 // The chunk splits in three: leading whitespace, text, trailing whitespace.
167 tail := text[len(body):]
168 visible := strings.TrimLeft(body, blanks)
169 lead := body[:len(body)-len(visible)]
170
171 out := ""
172 if g.started { // before the first word, all whitespace is dropped
173 out = squeeze(g.held + lead)
174 }
175 g.held, g.started = tail, true
176 return out + visible
177}
178
179// drop forgets the pending whitespace. Called when someone ELSE is about to
180// write — a tool announcing its command: that whitespace belonged to the text
181// before it, and printing it now would reopen the hole we just closed.
182func (g *gap) drop() {
183 g.mu.Lock()
184 g.held = ""
185 g.mu.Unlock()
186}
187
188// squeeze brings any run of line breaks down to two, that is ONE blank line at
189// most. Spaces are left alone: they carry the indentation of a code block, and
190// cutting it would shift the displayed code.
191func squeeze(s string) string {
192 var b strings.Builder
193 run := 0
194 for _, r := range s {
195 if r == '\n' {
196 run++
197 if run > 2 {
198 continue
199 }
200 } else {
201 run = 0
202 }
203 b.WriteRune(r)
204 }
205 return b.String()
206}
207
208// Generate queries the model, printing the answer as it comes (streaming). Some
209// OpenAI-compatible servers sometimes cut the SSE stream abruptly ("unexpected
210// end of JSON input"): if nothing has been printed AND no command has run yet,
211// we retry once WITHOUT streaming — which is more robust. The "no command"
212// condition matters: see the comment on the retry below.
213//
214// It returns the FULL conversation alongside the response — every intermediate
215// turn included — because Genkit only hands back the last message. The caller
216// needs the rest, or the agent forgets the commands it just ran.
217//
218// A spinner covers every moment where the program is waiting rather than
219// printing. On a local engine the first of those is long — the model may still
220// be loading, and the whole prompt has to be processed before a single token
221// comes back — and without it the agent looks hung.
222//
223// Watchdog: if no tokens arrive for a certain duration, we assume the connection
224// is hung and we cancel the context. This is the "silent loop" of
225// STALL_DETECTION_PROBLEM.md — a liveness failure, not a behavioural one.
226func (e *Engine) Generate(ctx context.Context, messages []*ai.Message, tools []ai.ToolRef) (*ai.ModelResponse, []*ai.Message, error) {
227 // Genkit only hands us back the LAST message. The intermediate turns — "I
228 // am calling bash", "here is the output" — exist only in the request it
229 // builds for the next turn. This middleware captures that request on its way
230 // through: the last turn seen therefore holds the whole conversation.
231 //
232 // The `WrapModel` hook wraps every call to the model — one per turn — so the
233 // last request seen is the most complete one. `MiddlewareFunc` adapts a
234 // closure to the `ai.Middleware` contract without going through a plugin:
235 // `New` is called once per `Generate`, which keeps `full` captured at the
236 // right level.
237 var full []*ai.Message
238 var gaps gap
239 capture := ai.MiddlewareFunc(func(context.Context) (*ai.Hooks, error) {
240 return &ai.Hooks{
241 WrapModel: func(ctx context.Context, params *ai.ModelParams, next ai.ModelNext) (*ai.ModelResponse, error) {
242 full = slices.Clone(params.Request.Messages)
243 resp, err := next(ctx, params)
244 // The server's own count of what it just read: the one exact
245 // measure of the context size, kept for the compression
246 // trigger. Zero means the server reported nothing.
247 if resp != nil && resp.Usage != nil && resp.Usage.InputTokens > 0 {
248 e.lastInput.Store(int64(resp.Usage.InputTokens))
249 }
250 return resp, err
251 },
252 // The tool is about to write to the screen: we forget the
253 // whitespace the previous text left pending. The hook lands at
254 // exactly the right moment — just before the tool prints its 🛠️
255 // line.
256 WrapTool: func(ctx context.Context, params *ai.ToolParams, next ai.ToolNext) (*ai.MultipartToolResponse, error) {
257 gaps.drop()
258 return next(ctx, params)
259 },
260 }, nil
261 })
262
263 opts := []ai.GenerateOption{
264 ai.WithModelName(e.Model),
265 ai.WithMessages(messages...),
266 ai.WithTools(tools...),
267 ai.WithMaxTurns(config.Cfg.MaxTurns), // max model ↔ tools round trips
268 ai.WithConfig(config.Cfg.Sampling),
269 ai.WithUse(capture),
270 }
271
272 // The prefix is printed by the first token rather than up front, so the
273 // spinner has a line to itself and vanishes without taking anything with it.
274 printed := false
275 emit := func(text string) {
276 // An event-consuming front end (ACP) takes the chunk raw: the editor
277 // renders Markdown and owns the layout, so the gap/squeeze machinery —
278 // which fights the terminal screen, not the text — stays out of the
279 // way. `printed` is still maintained: the no-streaming retry below must
280 // know whether the client already saw part of an answer.
281 if s := ui.ActiveSink(); s != nil {
282 if text == "" {
283 return
284 }
285 printed = true
286 s.Text(text)
287 return
288 }
289 // Trailing whitespace is held back, not printed: see the gap type.
290 out := gaps.next(text)
291 if out == "" {
292 return
293 }
294 // Stop BEFORE every print, not only the first one.
295 //
296 // A question can take several turns, and the tool restarts the spinner
297 // on its way back to the model. The text of the next turn then landed on
298 // the spinner's line, which erased it on the following frame
299 // (`\r\033[2K`): "Hello 👋 to Sam" became "⠦ Thinking… 0s 👋 to Sam".
300 // Worse, the last sentence of an answer vanished entirely, erased by the
301 // final Stop().
302 //
303 // Stop() is idempotent and returns immediately when nothing is running,
304 // so calling it on every chunk costs nothing.
305 spinner.Stop()
306 if !printed {
307 fmt.Print("🤖 ")
308 printed = true
309 }
310 fmt.Print(out)
311 }
312
313 // --- Watchdog ---
314 // A sub-context the watchdog cancels when no token has arrived for
315 // WatchdogTimeout; the streaming callback beats it on every chunk.
316 genCtx, beat, cancelGen := watchdog(ctx, config.Cfg.WatchdogTimeout)
317 defer cancelGen()
318
319 spinner.Start("Thinking")
320 defer spinner.Stop()
321
322 resp, err := genkit.Generate(genCtx, e.G, append(opts,
323 ai.WithStreaming(func(_ context.Context, chunk *ai.ModelResponseChunk) error {
324 beat()
325 emit(chunk.Text())
326 return nil
327 }),
328 )...)
329
330 // `full` grows as soon as a tool turn has happened: it is the witness that
331 // says whether commands ran, and it is used twice below.
332 toolsRan := len(full) > len(messages)
333
334 // The watchdog cut `genCtx` while the PARENT context is intact: this is a
335 // silence of the connection, not a cancellation by the user.
336 //
337 // We question the contexts rather than the error. Testing
338 // `err == context.Canceled` does not work: Genkit wraps the error, the
339 // comparison fails, and the turn then went off into the non-streamed retry —
340 // that is, it called the MUTE engine again. Measured: 5 s of watchdog then
341 // 20 s of "Retrying without streaming…" against a silent server, and the
342 // "stalled" message never shown.
343 stalled := genCtx.Err() != nil && ctx.Err() == nil
344 if err != nil && stalled {
345 // We return `full` despite the failure: if commands ran before the
346 // silence, their work must not be lost.
347 return nil, full, fmt.Errorf("generation timed out (stalled)")
348 }
349
350 // The retry is only worth it for an SSE stream cut before the first byte.
351 // Any other error (maxTurns exceeded, a failing tool, invalid arguments)
352 // happens AFTER commands have run, and replaying the turn would replay them:
353 // measured, an `echo >>` wrote itself twice. `full` grows as soon as a tool
354 // turn has happened, which is enough to tell the two apart.
355 if err != nil && !printed && !toolsRan {
356 spinner.Start("Retrying without streaming")
357 if resp, err = genkit.Generate(ctx, e.G, opts...); err == nil {
358 emit(resp.Text())
359 }
360 }
361
362 // `full` stops at the last call to the model; its answer is appended after.
363 if resp != nil && resp.Message != nil {
364 full = append(full, resp.Message)
365 }
366 return resp, full, err
367}
368
369// watchdog returns a child of ctx that is cancelled when beat() has not been
370// called for `timeout`. This is the "silent loop" of STALL_DETECTION_PROBLEM.md
371// — a liveness failure, not a behavioural one — and it is shared by Generate
372// and Summarize so that both fail the same way.
373//
374// The timestamp is atomic because beat() runs on the streaming goroutine while
375// the ticker reads it from another.
376func watchdog(ctx context.Context, timeout time.Duration) (genCtx context.Context, beat func(), cancel context.CancelFunc) {
377 var last atomic.Int64 // UnixNano
378 last.Store(time.Now().UnixNano())
379 genCtx, cancel = context.WithCancel(ctx)
380
381 go func() {
382 ticker := time.NewTicker(1 * time.Second)
383 defer ticker.Stop()
384 for {
385 select {
386 case <-genCtx.Done():
387 return
388 case <-ticker.C:
389 if time.Since(time.Unix(0, last.Load())) > timeout {
390 // Stop() before printing, like everything else: without it
391 // the watchdog message lands on the spinner's line and the
392 // next frame erases it — the diagnosis would vanish at the
393 // very moment it is needed.
394 spinner.Stop()
395 // ui.Out, not stdout: in ACP mode this diagnosis belongs to
396 // the logs (stderr), never to the JSON-RPC stream.
397 fmt.Fprintf(ui.Out, "\n[watchdog: connection hang detected after %v]\n", timeout)
398 cancel()
399 return
400 }
401 }
402 }
403 }()
404
405 return genCtx, func() { last.Store(time.Now().UnixNano()) }, cancel
406}
407
408// Summarize asks the model for a plain-text answer to `messages` — no tools,
409// no screen output. It exists for the context compression (internal/compact):
410// the summary must not stream onto the screen as if it were an answer, and a
411// request that offers tools invites the model to call them instead of writing.
412// It goes through the same Genkit and the same model reference as Generate,
413// so it lands on the same provider and the same server.
414//
415// The watchdog is kept, with twice the usual patience: this is the longest
416// prefill of the session (the whole old history at once) and nothing streams
417// until it is done. Streaming stays on so that the summary's own tokens beat
418// the watchdog while it is being written.
419func (e *Engine) Summarize(ctx context.Context, messages []*ai.Message, maxTokens int) (string, error) {
420 // The agent's sampling, minus what does not apply here: `parallel_tool_calls`
421 // without tools is a 400 on the OpenAI API, and `max_tokens` is the summary's
422 // budget, not the answer's.
423 sampling := maps.Clone(config.Cfg.Sampling)
424 if sampling == nil {
425 sampling = map[string]any{}
426 }
427 delete(sampling, "parallel_tool_calls")
428 sampling["max_tokens"] = maxTokens
429
430 genCtx, beat, cancel := watchdog(ctx, 2*config.Cfg.WatchdogTimeout)
431 defer cancel()
432
433 resp, err := genkit.Generate(genCtx, e.G,
434 ai.WithModelName(e.Model),
435 ai.WithMessages(messages...),
436 ai.WithConfig(sampling),
437 ai.WithStreaming(func(_ context.Context, _ *ai.ModelResponseChunk) error {
438 beat()
439 return nil
440 }),
441 )
442 if err != nil {
443 // Same test as in Generate: the contexts, not the wrapped error.
444 if genCtx.Err() != nil && ctx.Err() == nil {
445 return "", fmt.Errorf("summary timed out (stalled)")
446 }
447 return "", err
448 }
449 text := strings.TrimSpace(resp.Text())
450 if text == "" {
451 return "", fmt.Errorf("the model returned an empty summary")
452 }
453 return text, nil
454}
455
456// Names of the tools whose responses we count. They live here rather than
457// scattered through the display code: the counter is a measurement, and a
458// measurement needs a single definition.
459const (
460 toolBash = "bash"
461 toolSkill = "read_skill"
462 toolRead = "read_file"
463 toolWrite = "write_file"
464 toolEdit = "edit_file"
465)
466
467// executed returns, in order, the text of the calls to a given tool that
468// ACTUALLY ran.
469//
470// We start from the RESPONSES: a call the model asked for but which never ran
471// (interrupted turn, maxTurns exceeded) produces none. The text of the call, on
472// the other hand, is ONLY in the request — hence the pairing by Ref, the
473// identifier the server gives to each call.
474func executed(history []*ai.Message, name string) []string {
475 // Calls with no Ref (some servers do not set one) are paired in order, BUT
476 // only against the same tool name: a turn mixes `bash` and `read_skill`, and
477 // the first pending call is not necessarily for the right tool.
478 type call struct{ ref, name, text string }
479 var pending []call
480 var out []string
481
482 for _, m := range history {
483 for _, p := range m.Content {
484 switch {
485 case p.IsToolRequest():
486 pending = append(pending, call{p.ToolRequest.Ref, p.ToolRequest.Name, commandText(p.ToolRequest)})
487
488 case p.IsToolResponse() && p.ToolResponse != nil && p.ToolResponse.Name == name:
489 ref := p.ToolResponse.Ref
490 i := slices.IndexFunc(pending, func(c call) bool {
491 if ref != "" {
492 return c.ref == ref
493 }
494 return c.name == name
495 })
496 if i < 0 { // lost call: the name beats nothing
497 out = append(out, name)
498 continue
499 }
500 out = append(out, pending[i].text)
501 pending = slices.Delete(pending, i, i+1)
502 }
503 }
504 }
505 return out
506}
507
508// commandText describes a tool call in one line. For `bash` it is the command
509// itself — the rest is only wrapping; for any other tool, its name followed by
510// its arguments.
511func commandText(r *ai.ToolRequest) string {
512 if in, ok := r.Input.(map[string]any); ok {
513 if c, ok := in["command"].(string); ok && strings.TrimSpace(c) != "" {
514 return strings.TrimSpace(c)
515 }
516 }
517 args, err := json.Marshal(r.Input)
518 if err != nil || string(args) == "null" {
519 return r.Name
520 }
521 return r.Name + " " + string(args)
522}
523
524// responses counts the responses of a given tool in a history.
525func responses(history []*ai.Message, name string) int {
526 return len(executed(history, name))
527}
528
529// Commands counts the commands actually run in a history — one per response of
530// the `bash` tool. It is the only reliable measure of the work done: a model
531// that tells stories without acting produces a fine answer and a count of zero.
532//
533// The other tools do NOT count. Loading a skill runs nothing, and counting those
534// reads would make the number say the opposite of what it is for: measured
535// before the fix, a turn with two reads and a single command displayed
536// "⚙ 3 command(s)".
537func Commands(history []*ai.Message) int {
538 return responses(history, toolBash)
539}
540
541// CommandList returns, in order, the text of the commands Commands counts — the
542// one shown in the recap under the count when `displayCommands` is on. The count
543// says HOW MANY, the list says WHAT.
544func CommandList(history []*ai.Message) []string {
545 return executed(history, toolBash)
546}
547
548// FileOps returns, in order, the file-tool calls that ran — read_file,
549// write_file, edit_file — as one line each ("edit_file hello/main.go (2 edits)").
550//
551// They are counted APART from the commands, not folded into them: this part
552// compares editing through the `edit` CLI (one bash command per edit) with
553// editing through built-in tools, and a recap that counted both as "commands"
554// would hide the very difference being measured. `⚙ 3 command(s) · 📝 2 file
555// op(s)` keeps the two columns readable side by side.
556func FileOps(history []*ai.Message) []string {
557 return executedAny(history, toolRead, toolWrite, toolEdit)
558}
559
560// executedAny is executed() across several tool names, keeping the order in
561// which the responses appear — a per-name pass would lose the interleaving of
562// a read followed by an edit.
563func executedAny(history []*ai.Message, names ...string) []string {
564 type call struct{ ref, name, text string }
565 var pending []call
566 var out []string
567 isOurs := func(n string) bool { return slices.Contains(names, n) }
568
569 for _, m := range history {
570 for _, p := range m.Content {
571 switch {
572 case p.IsToolRequest():
573 pending = append(pending, call{p.ToolRequest.Ref, p.ToolRequest.Name, fileOpText(p.ToolRequest)})
574 case p.IsToolResponse() && p.ToolResponse != nil && isOurs(p.ToolResponse.Name):
575 ref, name := p.ToolResponse.Ref, p.ToolResponse.Name
576 i := slices.IndexFunc(pending, func(c call) bool {
577 if ref != "" {
578 return c.ref == ref
579 }
580 return c.name == name
581 })
582 if i < 0 {
583 out = append(out, name)
584 continue
585 }
586 out = append(out, pending[i].text)
587 pending = slices.Delete(pending, i, i+1)
588 }
589 }
590 }
591 return out
592}
593
594// fileOpText describes a file-tool call in one line: the tool, the path, and
595// for edit_file the number of edits — never the texts themselves, which can
596// run to dozens of lines and belong in the diff, not in the recap.
597func fileOpText(r *ai.ToolRequest) string {
598 in, _ := r.Input.(map[string]any)
599 path, _ := in["path"].(string)
600 switch r.Name {
601 case toolEdit:
602 if edits, ok := in["edits"].([]any); ok {
603 return fmt.Sprintf("%s %s (%d edit(s))", r.Name, path, len(edits))
604 }
605 case toolRead:
606 if s, ok := in["start"].(float64); ok && s > 0 {
607 e, _ := in["end"].(float64)
608 return fmt.Sprintf("%s %s %d-%d", r.Name, path, int(s), int(e))
609 }
610 }
611 if path == "" {
612 return commandText(r)
613 }
614 return r.Name + " " + path
615}
616
617// Skills counts the skills loaded — one per response of `read_skill`. Shown
618// next to the command count: seeing which procedure the agent chose to follow is
619// half the point of this part.
620func Skills(history []*ai.Message) int {
621 return responses(history, toolSkill)
622}