1package modes
2
3import (
4 "context"
5 "fmt"
6 "sort"
7 "strings"
8 "sync"
9 "time"
10
11 "github.com/mattn/go-runewidth"
12 "github.com/patriceckhart/zot/packages/provider"
13 "github.com/patriceckhart/zot/packages/tui"
14)
15
16type llamaDialogStep int
17
18const (
19 llamaClosed llamaDialogStep = iota
20 llamaLoading
21 llamaModels
22 llamaConfirmUnload
23 llamaConfirmRemove
24 llamaSearch
25 llamaModelDetails
26 llamaQuantizations
27 llamaProgress
28 llamaError
29)
30
31// llamaDialog manages models exposed by a configured llama.cpp router. Network
32// operations run outside the input loop and update this state under mu.
33type llamaDialog struct {
34 mu sync.Mutex
35
36 step llamaDialogStep
37 client *provider.LlamaCPPClient
38 hf *provider.HuggingFaceClient
39 models []provider.LlamaCPPModel
40 results []provider.HuggingFaceModel
41 details provider.HuggingFaceModelDetails
42 cursor int
43 query string
44 selected string
45 message string
46 progress provider.LlamaCPPProgress
47 cancel context.CancelFunc
48 searchCancel context.CancelFunc
49 searchSeq uint64
50 invalidate func()
51}
52
53func newLlamaDialog() *llamaDialog { return &llamaDialog{} }
54
55func (d *llamaDialog) Active() bool {
56 d.mu.Lock()
57 defer d.mu.Unlock()
58 return d.step != llamaClosed
59}
60
61func (d *llamaDialog) Open(serverURL, apiKey string, invalidate func()) error {
62 if strings.TrimSpace(serverURL) == "" {
63 return fmt.Errorf("configure the llama.cpp router with /login first")
64 }
65 client, err := provider.NewLlamaCPPClient(serverURL, apiKey)
66 if err != nil {
67 return err
68 }
69 d.mu.Lock()
70 d.step = llamaLoading
71 d.client = client
72 d.hf = provider.NewHuggingFaceClient(provider.FindHuggingFaceToken())
73 d.cursor = 0
74 d.query = ""
75 d.message = "Connecting to " + client.ServerURL
76 d.invalidate = invalidate
77 d.mu.Unlock()
78 go d.refresh(llamaModels, "")
79 return nil
80}
81
82func (d *llamaDialog) Close() {
83 d.mu.Lock()
84 if d.cancel != nil {
85 d.cancel()
86 d.cancel = nil
87 }
88 if d.searchCancel != nil {
89 d.searchCancel()
90 d.searchCancel = nil
91 }
92 d.step = llamaClosed
93 d.mu.Unlock()
94}
95
96func (d *llamaDialog) notify() {
97 d.mu.Lock()
98 invalidate := d.invalidate
99 d.mu.Unlock()
100 if invalidate != nil {
101 invalidate()
102 }
103}
104
105func (d *llamaDialog) setError(err error) {
106 d.mu.Lock()
107 if d.step != llamaClosed {
108 d.step = llamaError
109 d.message = err.Error()
110 d.cancel = nil
111 }
112 d.mu.Unlock()
113 d.notify()
114}
115
116func (d *llamaDialog) refresh(next llamaDialogStep, message string) {
117 d.mu.Lock()
118 client := d.client
119 d.mu.Unlock()
120 models, err := client.List(context.Background(), true)
121 if err != nil {
122 d.setError(err)
123 return
124 }
125 sort.Slice(models, func(i, j int) bool {
126 il := models[i].Status.Value == "loaded" || models[i].Status.Value == "sleeping"
127 jl := models[j].Status.Value == "loaded" || models[j].Status.Value == "sleeping"
128 if il != jl {
129 return il
130 }
131 return models[i].ID < models[j].ID
132 })
133 provider.SetManagedModels(provider.LlamaCPPModels(models, client.ServerURL))
134 d.mu.Lock()
135 if d.step != llamaClosed {
136 d.models = models
137 d.cursor = min(d.cursor, len(models))
138 d.step = next
139 d.message = message
140 d.cancel = nil
141 }
142 d.mu.Unlock()
143 d.notify()
144}
145
146func (d *llamaDialog) beginProgress(title, model string, run func(context.Context, func(provider.LlamaCPPProgress)) error) {
147 ctx, cancel := context.WithCancel(context.Background())
148 d.mu.Lock()
149 d.step = llamaProgress
150 d.selected = model
151 d.message = title
152 d.progress = provider.LlamaCPPProgress{Message: "Starting"}
153 d.cancel = cancel
154 d.mu.Unlock()
155 d.notify()
156 go func() {
157 err := run(ctx, func(progress provider.LlamaCPPProgress) {
158 d.mu.Lock()
159 if d.step == llamaProgress {
160 d.progress = progress
161 }
162 d.mu.Unlock()
163 d.notify()
164 })
165 if err != nil {
166 if ctx.Err() != nil {
167 d.refresh(llamaModels, "operation cancelled")
168 return
169 }
170 d.setError(err)
171 return
172 }
173 d.refresh(llamaModels, title+" complete")
174 }()
175}
176
177func (d *llamaDialog) load(model string) {
178 d.mu.Lock()
179 client := d.client
180 d.mu.Unlock()
181 d.beginProgress("Loading model", model, func(ctx context.Context, update func(provider.LlamaCPPProgress)) error {
182 _, err := client.LoadAndWait(ctx, model, update)
183 return err
184 })
185}
186
187func (d *llamaDialog) unload(model string) {
188 d.mu.Lock()
189 client := d.client
190 d.mu.Unlock()
191 d.beginProgress("Unloading model", model, func(ctx context.Context, update func(provider.LlamaCPPProgress)) error {
192 update(provider.LlamaCPPProgress{Message: "Unloading model"})
193 return client.UnloadAndWait(ctx, model)
194 })
195}
196
197func (d *llamaDialog) download(model string) {
198 d.mu.Lock()
199 client := d.client
200 d.mu.Unlock()
201 d.beginProgress("Downloading model", model, func(ctx context.Context, update func(provider.LlamaCPPProgress)) error {
202 _, err := client.DownloadAndWait(ctx, model, update)
203 return err
204 })
205}
206
207func (d *llamaDialog) remove(model string) {
208 d.mu.Lock()
209 client := d.client
210 d.mu.Unlock()
211 d.beginProgress("Removing model", model, func(ctx context.Context, update func(provider.LlamaCPPProgress)) error {
212 update(provider.LlamaCPPProgress{Message: "Removing cached files"})
213 return client.Remove(ctx, model)
214 })
215}
216
217func llamaModelLoaded(model provider.LlamaCPPModel) bool {
218 return model.Status.Value == "loaded" || model.Status.Value == "sleeping"
219}
220
221func (d *llamaDialog) scheduleSearchLocked() {
222 if d.searchCancel != nil {
223 d.searchCancel()
224 d.searchCancel = nil
225 }
226 d.searchSeq++
227 d.results = nil
228 d.cursor = 0
229 query := strings.TrimSpace(d.query)
230 if len(query) < 2 {
231 d.message = "type at least 2 characters"
232 return
233 }
234 d.message = "searching Hugging Face..."
235 ctx, cancel := context.WithCancel(context.Background())
236 d.searchCancel = cancel
237 seq := d.searchSeq
238 hf := d.hf
239 go func() {
240 timer := time.NewTimer(400 * time.Millisecond)
241 defer timer.Stop()
242 select {
243 case <-ctx.Done():
244 return
245 case <-timer.C:
246 }
247 results, err := hf.Search(ctx, query)
248 if ctx.Err() != nil {
249 return
250 }
251 d.mu.Lock()
252 if d.step != llamaSearch || d.searchSeq != seq || strings.TrimSpace(d.query) != query {
253 d.mu.Unlock()
254 return
255 }
256 d.searchCancel = nil
257 if err != nil {
258 d.results = nil
259 d.message = err.Error()
260 } else {
261 d.results = results
262 d.message = ""
263 if len(results) == 0 {
264 d.message = "no GGUF models found"
265 }
266 }
267 d.mu.Unlock()
268 d.notify()
269 }()
270}
271
272func (d *llamaDialog) HandleKey(k tui.Key) {
273 d.mu.Lock()
274 step := d.step
275 if k.Kind == tui.KeyCtrlC {
276 cancel := d.cancel
277 searchCancel := d.searchCancel
278 client := d.client
279 model := d.selected
280 d.step = llamaClosed
281 d.mu.Unlock()
282 if cancel != nil {
283 cancel()
284 go func() { _ = client.Unload(context.Background(), model) }()
285 }
286 if searchCancel != nil {
287 searchCancel()
288 }
289 return
290 }
291 if k.Kind == tui.KeyEsc {
292 switch step {
293 case llamaProgress:
294 cancel := d.cancel
295 client := d.client
296 model := d.selected
297 d.mu.Unlock()
298 if cancel != nil {
299 cancel()
300 go func() { _ = client.Unload(context.Background(), model) }()
301 }
302 return
303 case llamaModelDetails, llamaQuantizations:
304 d.step = llamaSearch
305 d.cursor = 0
306 d.mu.Unlock()
307 return
308 case llamaSearch:
309 if d.searchCancel != nil {
310 d.searchCancel()
311 d.searchCancel = nil
312 }
313 d.step = llamaModels
314 d.cursor = 0
315 d.mu.Unlock()
316 return
317 case llamaConfirmUnload, llamaConfirmRemove:
318 d.step = llamaModels
319 d.cursor = 0
320 d.mu.Unlock()
321 return
322 case llamaError:
323 d.step = llamaClosed
324 d.mu.Unlock()
325 return
326 default:
327 d.step = llamaClosed
328 d.mu.Unlock()
329 return
330 }
331 }
332
333 switch step {
334 case llamaModels:
335 count := len(d.models) + 1
336 switch k.Kind {
337 case tui.KeyUp:
338 if d.cursor > 0 {
339 d.cursor--
340 }
341 case tui.KeyDown:
342 if d.cursor+1 < count {
343 d.cursor++
344 }
345 case tui.KeyEnter:
346 if d.cursor == len(d.models) {
347 d.step, d.cursor, d.query, d.message = llamaSearch, 0, "", "type at least 2 characters"
348 d.results = nil
349 d.mu.Unlock()
350 return
351 }
352 if len(d.models) > 0 {
353 model := d.models[d.cursor]
354 if llamaModelLoaded(model) {
355 d.selected = model.ID
356 d.step = llamaConfirmUnload
357 d.cursor = 0
358 d.mu.Unlock()
359 return
360 }
361 if model.Status.Value == "unloaded" {
362 id := model.ID
363 d.mu.Unlock()
364 d.load(id)
365 return
366 }
367 d.message = model.ID + " is " + model.Status.Value
368 }
369 case tui.KeyRune:
370 if (k.Rune == 'd' || k.Rune == 'D') && !k.Ctrl && !k.Alt && d.cursor < len(d.models) {
371 model := d.models[d.cursor]
372 if !model.CanRemove {
373 d.message = "only router-downloaded cache models can be removed"
374 break
375 }
376 d.selected = model.ID
377 d.step = llamaConfirmRemove
378 d.cursor = 0
379 }
380 }
381 case llamaConfirmUnload:
382 switch k.Kind {
383 case tui.KeyUp, tui.KeyDown:
384 d.cursor = 1 - d.cursor
385 case tui.KeyEnter:
386 if d.cursor == 0 {
387 id := d.selected
388 d.mu.Unlock()
389 d.unload(id)
390 return
391 }
392 d.step = llamaModels
393 }
394 case llamaConfirmRemove:
395 switch k.Kind {
396 case tui.KeyUp, tui.KeyDown:
397 d.cursor = 1 - d.cursor
398 case tui.KeyEnter:
399 if d.cursor == 0 {
400 id := d.selected
401 d.mu.Unlock()
402 d.remove(id)
403 return
404 }
405 d.step = llamaModels
406 }
407 case llamaSearch:
408 switch k.Kind {
409 case tui.KeyBackspace:
410 runes := []rune(d.query)
411 if len(runes) > 0 {
412 d.query = string(runes[:len(runes)-1])
413 d.scheduleSearchLocked()
414 }
415 case tui.KeyPaste:
416 d.query += strings.Join(strings.Fields(k.Paste), " ")
417 d.scheduleSearchLocked()
418 case tui.KeyRune:
419 if !k.Ctrl && !k.Alt && k.Rune >= 0x20 && k.Rune < 0x7f {
420 d.query += string(k.Rune)
421 d.scheduleSearchLocked()
422 }
423 case tui.KeyUp:
424 if d.cursor > 0 {
425 d.cursor--
426 }
427 case tui.KeyDown:
428 if d.cursor+1 < len(d.results) {
429 d.cursor++
430 }
431 case tui.KeyEnter:
432 if len(d.results) == 0 {
433 break
434 }
435 if d.searchCancel != nil {
436 d.searchCancel()
437 d.searchCancel = nil
438 }
439 id := d.results[d.cursor].ID
440 hf := d.hf
441 d.step, d.message = llamaLoading, "loading model details"
442 d.mu.Unlock()
443 go func() {
444 details, err := hf.Details(context.Background(), id)
445 if err != nil {
446 d.setError(err)
447 return
448 }
449 d.mu.Lock()
450 if d.step != llamaClosed {
451 d.details, d.cursor = details, 0
452 if len(details.Quantizations) > 0 {
453 d.step = llamaQuantizations
454 } else {
455 d.step = llamaModelDetails
456 }
457 }
458 d.mu.Unlock()
459 d.notify()
460 }()
461 return
462 }
463 case llamaModelDetails:
464 if k.Kind == tui.KeyEnter {
465 id := d.details.ID
466 d.mu.Unlock()
467 d.download(id)
468 return
469 }
470 case llamaQuantizations:
471 switch k.Kind {
472 case tui.KeyUp:
473 if d.cursor > 0 {
474 d.cursor--
475 }
476 case tui.KeyDown:
477 if d.cursor+1 < len(d.details.Quantizations) {
478 d.cursor++
479 }
480 case tui.KeyEnter:
481 if len(d.details.Quantizations) > 0 {
482 id := d.details.ID + ":" + d.details.Quantizations[d.cursor].Name
483 d.mu.Unlock()
484 d.download(id)
485 return
486 }
487 }
488 }
489 d.mu.Unlock()
490}
491
492func (d *llamaDialog) Render(th tui.Theme, width int) []string {
493 d.mu.Lock()
494 defer d.mu.Unlock()
495 if d.step == llamaClosed {
496 return nil
497 }
498 lines := []string{frameHeader(th, "llama.cpp models", width)}
499 if d.client != nil {
500 lines = append(lines, th.FG256(th.Muted, d.client.ServerURL))
501 }
502 if d.message != "" {
503 lines = append(lines, th.FG256(th.Muted, d.message))
504 }
505
506 switch d.step {
507 case llamaLoading:
508 lines = append(lines, th.FG256(th.Muted, " working..."))
509 case llamaModels:
510 const visible = 12
511 start, end := dialogWindow(d.cursor, len(d.models)+1, visible)
512 for index := start; index < end; index++ {
513 label := "Download model from Hugging Face..."
514 if index < len(d.models) {
515 model := d.models[index]
516 status := model.Status.Value
517 if status == "loaded" || status == "sleeping" {
518 status = "loaded"
519 }
520 label = fmt.Sprintf("%-10s %s", status, model.ID)
521 }
522 line := " " + label
523 if index == d.cursor {
524 lines = append(lines, th.PadHighlight(line, width))
525 } else {
526 lines = append(lines, th.FG256(th.Muted, line))
527 }
528 }
529 lines = append(lines, th.FG256(th.Muted, "enter loads/unloads, d removes cached model, esc closes"))
530 case llamaConfirmUnload:
531 lines = append(lines, "", " Unload "+d.selected+"?")
532 for index, label := range []string{"Yes", "No"} {
533 line := " " + label
534 if index == d.cursor {
535 lines = append(lines, th.PadHighlight(line, width))
536 } else {
537 lines = append(lines, th.FG256(th.Muted, line))
538 }
539 }
540 case llamaConfirmRemove:
541 lines = append(lines, "", " Remove cached model "+d.selected+"?")
542 for index, label := range []string{"Yes", "No"} {
543 line := " " + label
544 if index == d.cursor {
545 lines = append(lines, th.PadHighlight(line, width))
546 } else {
547 lines = append(lines, th.FG256(th.Muted, line))
548 }
549 }
550 case llamaSearch:
551 lines = append(lines, "", th.FG256(th.Muted, " Model name or owner/repository"), th.FG256(th.FG, " "+d.query), "")
552 start, end := dialogWindow(d.cursor, len(d.results), 10)
553 for index := start; index < end; index++ {
554 result := d.results[index]
555 line := fmt.Sprintf(" %s %s downloads", result.ID, compactNumber(result.Downloads))
556 if index == d.cursor {
557 lines = append(lines, th.PadHighlight(line, width))
558 } else {
559 lines = append(lines, th.FG256(th.Muted, line))
560 }
561 }
562 lines = append(lines, th.FG256(th.Muted, "enter selects, esc goes back"))
563 case llamaModelDetails:
564 lines = append(lines, "", " "+d.details.ID)
565 if d.details.Gated != "" {
566 lines = append(lines, th.FG256(th.Error, " access approval and HF_TOKEN may be required"))
567 }
568 lines = append(lines, th.FG256(th.Muted, " enter downloads, esc goes back"))
569 case llamaQuantizations:
570 lines = append(lines, "", " Select quantization for "+d.details.ID)
571 if d.details.Gated != "" {
572 lines = append(lines, th.FG256(th.Error, " gated model: approve access on huggingface.co and configure HF_TOKEN on the router"))
573 }
574 start, end := dialogWindow(d.cursor, len(d.details.Quantizations), 10)
575 for index := start; index < end; index++ {
576 quant := d.details.Quantizations[index]
577 detail := ""
578 if quant.HasSize {
579 detail = " " + provider.FormatBytes(quant.Size)
580 }
581 if quant.Name == "Q4_K_M" {
582 detail += " recommended"
583 }
584 line := " " + quant.Name + detail
585 if index == d.cursor {
586 lines = append(lines, th.PadHighlight(line, width))
587 } else {
588 lines = append(lines, th.FG256(th.Muted, line))
589 }
590 }
591 case llamaProgress:
592 lines = append(lines, "", " "+d.selected, th.FG256(th.Muted, " "+d.progress.Message))
593 if d.progress.HasRatio {
594 const cells = 32
595 filled := int(d.progress.Ratio*float64(cells) + .5)
596 filled = max(0, min(cells, filled))
597 bar := strings.Repeat("#", filled) + strings.Repeat("-", cells-filled)
598 lines = append(lines, th.FG256(th.Accent, fmt.Sprintf(" %s %d%%", bar, int(d.progress.Ratio*100+.5))))
599 }
600 if d.progress.Detail != "" {
601 lines = append(lines, th.FG256(th.Muted, " "+d.progress.Detail))
602 }
603 lines = append(lines, th.FG256(th.Muted, " esc stops"))
604 case llamaError:
605 lines = append(lines, th.FG256(th.Error, " "+d.message), th.FG256(th.Muted, " esc closes"))
606 }
607 return append(lines, frameRule(th, width))
608}
609
610// CursorPos returns the search input caret location inside the padded dialog.
611func (d *llamaDialog) CursorPos() (row, col int) {
612 d.mu.Lock()
613 defer d.mu.Unlock()
614 if d.step != llamaSearch {
615 return -1, -1
616 }
617 row = 5 // header, injected frame padding, server URL, blank, prompt
618 if d.message != "" {
619 row++
620 }
621 return row, 2 + runewidth.StringWidth(d.query)
622}
623
624func dialogWindow(cursor, count, visible int) (int, int) {
625 if count <= visible {
626 return 0, count
627 }
628 start := cursor - visible/2
629 if start < 0 {
630 start = 0
631 }
632 if start+visible > count {
633 start = count - visible
634 }
635 return start, start + visible
636}
637
638func compactNumber(value int64) string {
639 if value >= 1_000_000 {
640 return fmt.Sprintf("%.1fM", float64(value)/1_000_000)
641 }
642 if value >= 1_000 {
643 return fmt.Sprintf("%.1fk", float64(value)/1_000)
644 }
645 return fmt.Sprintf("%d", value)
646}