1package provider
2
3import (
4 "bufio"
5 "context"
6 "encoding/json"
7 "errors"
8 "fmt"
9 "io"
10 "net/http"
11 "net/url"
12 "os"
13 "path/filepath"
14 "regexp"
15 "sort"
16 "strconv"
17 "strings"
18 "time"
19)
20
21const (
22 LlamaCPPProviderID = "llama.cpp"
23 defaultHuggingFaceURL = "https://huggingface.co"
24 llamaCPPRequestTimeout = 15 * time.Second
25 llamaCPPProgressPollTime = 250 * time.Millisecond
26)
27
28// LlamaCPPModel is one entry returned by a llama.cpp router's /models endpoint.
29type LlamaCPPModel struct {
30 ID string `json:"id"`
31 Aliases []string `json:"aliases,omitempty"`
32 Source string `json:"source,omitempty"`
33 CanRemove bool `json:"can_remove,omitempty"`
34 Status struct {
35 Value string `json:"value"`
36 Args []string `json:"args,omitempty"`
37 Failed bool `json:"failed,omitempty"`
38 ExitCode *int `json:"exit_code,omitempty"`
39 Progress map[string]LlamaCPPBytes `json:"progress,omitempty"`
40 } `json:"status"`
41 Architecture struct {
42 InputModalities []string `json:"input_modalities,omitempty"`
43 } `json:"architecture,omitempty"`
44 Meta struct {
45 Context int `json:"n_ctx,omitempty"`
46 TrainContext int `json:"n_ctx_train,omitempty"`
47 Size int64 `json:"size,omitempty"`
48 FileType string `json:"ftype,omitempty"`
49 } `json:"meta,omitempty"`
50 // Progress is emitted at the model's top level by current llama.cpp
51 // routers. Status.Progress supports versions that follow the documented
52 // nested shape.
53 Progress map[string]LlamaCPPBytes `json:"progress,omitempty"`
54}
55
56// LlamaCPPBytes describes byte progress for one downloaded file.
57type LlamaCPPBytes struct {
58 Done int64 `json:"done"`
59 Total int64 `json:"total"`
60}
61
62// LlamaCPPProgress is a user-facing load or download progress update.
63type LlamaCPPProgress struct {
64 Message string
65 Ratio float64
66 HasRatio bool
67 Detail string
68}
69
70// LlamaCPPClient talks to the model-management API exposed by llama-server
71// when it is running in router mode.
72type LlamaCPPClient struct {
73 ServerURL string
74 APIKey string
75 HTTP *http.Client
76}
77
78// NormalizeLlamaCPPURL validates a router URL and removes a trailing /v1,
79// which belongs to the inference API rather than the management API.
80func NormalizeLlamaCPPURL(value string) (string, error) {
81 u, err := url.Parse(strings.TrimSpace(value))
82 if err != nil {
83 return "", fmt.Errorf("parse llama.cpp server URL: %w", err)
84 }
85 if u.Scheme != "http" && u.Scheme != "https" {
86 return "", errors.New("llama.cpp server URL must use http or https")
87 }
88 if u.Host == "" {
89 return "", errors.New("llama.cpp server URL is missing a host")
90 }
91 u.RawQuery = ""
92 u.Fragment = ""
93 u.Path = strings.TrimSuffix(strings.TrimRight(u.Path, "/"), "/v1")
94 return strings.TrimRight(u.String(), "/"), nil
95}
96
97// LlamaCPPInferenceURL returns the OpenAI-compatible inference base URL.
98func LlamaCPPInferenceURL(serverURL string) (string, error) {
99 normalized, err := NormalizeLlamaCPPURL(serverURL)
100 if err != nil {
101 return "", err
102 }
103 return normalized + "/v1", nil
104}
105
106func NewLlamaCPPClient(serverURL, apiKey string) (*LlamaCPPClient, error) {
107 normalized, err := NormalizeLlamaCPPURL(serverURL)
108 if err != nil {
109 return nil, err
110 }
111 return &LlamaCPPClient{ServerURL: normalized, APIKey: apiKey, HTTP: &http.Client{Timeout: llamaCPPRequestTimeout}}, nil
112}
113
114func (c *LlamaCPPClient) request(ctx context.Context, method, path string, body io.Reader, out any) error {
115 req, err := http.NewRequestWithContext(ctx, method, c.ServerURL+path, body)
116 if err != nil {
117 return err
118 }
119 if body != nil {
120 req.Header.Set("content-type", "application/json")
121 }
122 if c.APIKey != "" {
123 req.Header.Set("authorization", "Bearer "+c.APIKey)
124 }
125 resp, err := c.HTTP.Do(req)
126 if err != nil {
127 return err
128 }
129 defer resp.Body.Close()
130 data, err := io.ReadAll(resp.Body)
131 if err != nil {
132 return err
133 }
134 if resp.StatusCode < 200 || resp.StatusCode >= 300 {
135 var payload struct {
136 Error struct {
137 Message string `json:"message"`
138 } `json:"error"`
139 }
140 _ = json.Unmarshal(data, &payload)
141 if payload.Error.Message != "" {
142 return errors.New(payload.Error.Message)
143 }
144 return fmt.Errorf("llama.cpp returned HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(data)))
145 }
146 if out != nil && len(data) > 0 {
147 if err := json.Unmarshal(data, out); err != nil {
148 return fmt.Errorf("parse llama.cpp response: %w", err)
149 }
150 }
151 return nil
152}
153
154func (c *LlamaCPPClient) List(ctx context.Context, reload bool) ([]LlamaCPPModel, error) {
155 path := "/models"
156 if reload {
157 path += "?reload=1"
158 }
159 var response struct {
160 Data []LlamaCPPModel `json:"data"`
161 }
162 if err := c.request(ctx, http.MethodGet, path, nil, &response); err != nil {
163 return nil, err
164 }
165 for _, model := range response.Data {
166 if model.ID == "" || model.Status.Value == "" {
167 return nil, errors.New("server is not running in llama.cpp router mode")
168 }
169 }
170 return response.Data, nil
171}
172
173func (c *LlamaCPPClient) modelAction(ctx context.Context, path, model string) error {
174 body := strings.NewReader(`{"model":` + strconv.Quote(model) + `}`)
175 return c.request(ctx, http.MethodPost, path, body, nil)
176}
177
178func (c *LlamaCPPClient) Load(ctx context.Context, model string) error {
179 return c.modelAction(ctx, "/models/load", model)
180}
181
182func (c *LlamaCPPClient) Unload(ctx context.Context, model string) error {
183 return c.modelAction(ctx, "/models/unload", model)
184}
185
186func (c *LlamaCPPClient) Download(ctx context.Context, model string) error {
187 return c.modelAction(ctx, "/models", model)
188}
189
190// Remove deletes a router-managed cache model from disk. Models discovered
191// from --models-dir or presets are not removable through the router API.
192func (c *LlamaCPPClient) Remove(ctx context.Context, model string) error {
193 path := "/models?" + url.Values{"model": {model}}.Encode()
194 return c.request(ctx, http.MethodDelete, path, nil, nil)
195}
196
197type llamaCPPEvent struct {
198 Model string `json:"model"`
199 Event string `json:"event"`
200 Data json.RawMessage `json:"data"`
201}
202
203func (c *LlamaCPPClient) watch(ctx context.Context, events chan<- llamaCPPEvent) error {
204 req, err := http.NewRequestWithContext(ctx, http.MethodGet, c.ServerURL+"/models/sse", nil)
205 if err != nil {
206 return err
207 }
208 if c.APIKey != "" {
209 req.Header.Set("authorization", "Bearer "+c.APIKey)
210 }
211 // http.Client.Timeout covers the complete response body lifetime, which is
212 // appropriate for ordinary management requests but would terminate this
213 // long-lived SSE stream. Preserve the configured transport and other client
214 // behavior while letting ctx control the stream lifetime.
215 streamClient := *c.HTTP
216 streamClient.Timeout = 0
217 resp, err := streamClient.Do(req)
218 if err != nil {
219 return err
220 }
221 defer resp.Body.Close()
222 if resp.StatusCode != http.StatusOK {
223 return fmt.Errorf("llama.cpp SSE returned HTTP %d", resp.StatusCode)
224 }
225 scanner := bufio.NewScanner(resp.Body)
226 for scanner.Scan() {
227 line := scanner.Text()
228 if !strings.HasPrefix(line, "data:") {
229 continue
230 }
231 var event llamaCPPEvent
232 if json.Unmarshal([]byte(strings.TrimSpace(strings.TrimPrefix(line, "data:"))), &event) == nil {
233 select {
234 case events <- event:
235 case <-ctx.Done():
236 return ctx.Err()
237 }
238 }
239 }
240 return scanner.Err()
241}
242
243func loadProgress(data json.RawMessage) (LlamaCPPProgress, bool) {
244 var payload struct {
245 Progress struct {
246 Stages []string `json:"stages"`
247 Current string `json:"current"`
248 Stage string `json:"stage"`
249 Value *float64 `json:"value"`
250 } `json:"progress"`
251 }
252 if json.Unmarshal(data, &payload) != nil {
253 return LlamaCPPProgress{}, false
254 }
255 stage := payload.Progress.Current
256 if stage == "" {
257 stage = payload.Progress.Stage
258 }
259 if stage == "" && payload.Progress.Value == nil {
260 return LlamaCPPProgress{}, false
261 }
262 progress := LlamaCPPProgress{Message: "Loading model"}
263 if stage != "" {
264 progress.Message = "Loading " + strings.ReplaceAll(stage, "_", " ")
265 }
266 if payload.Progress.Value != nil {
267 ratio := max(0, min(1, *payload.Progress.Value))
268 if stage != "" && len(payload.Progress.Stages) > 0 {
269 for index, candidate := range payload.Progress.Stages {
270 if candidate == stage {
271 ratio = (float64(index) + ratio) / float64(len(payload.Progress.Stages))
272 break
273 }
274 }
275 }
276 progress.Ratio, progress.HasRatio = ratio, true
277 }
278 return progress, true
279}
280
281func downloadProgress(files map[string]LlamaCPPBytes) (LlamaCPPProgress, bool) {
282 var done, total int64
283 for _, file := range files {
284 done += file.Done
285 total += file.Total
286 }
287 if total <= 0 {
288 return LlamaCPPProgress{}, false
289 }
290 return LlamaCPPProgress{
291 Message: "Downloading model",
292 Ratio: float64(done) / float64(total), HasRatio: true,
293 Detail: FormatBytes(done) + " / " + FormatBytes(total),
294 }, true
295}
296
297func (c *LlamaCPPClient) LoadAndWait(ctx context.Context, model string, update func(LlamaCPPProgress)) (LlamaCPPModel, error) {
298 watchCtx, cancel := context.WithCancel(ctx)
299 defer cancel()
300 events := make(chan llamaCPPEvent, 8)
301 go func() { _ = c.watch(watchCtx, events) }()
302 if err := c.Load(ctx, model); err != nil {
303 return LlamaCPPModel{}, err
304 }
305 update(LlamaCPPProgress{Message: "Loading model"})
306 ticker := time.NewTicker(llamaCPPProgressPollTime)
307 defer ticker.Stop()
308 for {
309 select {
310 case <-ctx.Done():
311 return LlamaCPPModel{}, ctx.Err()
312 case event := <-events:
313 if event.Model == model {
314 if progress, ok := loadProgress(event.Data); ok {
315 update(progress)
316 }
317 }
318 case <-ticker.C:
319 models, err := c.List(ctx, false)
320 if err != nil {
321 return LlamaCPPModel{}, err
322 }
323 for _, entry := range models {
324 if entry.ID != model {
325 continue
326 }
327 if entry.Status.Value == "loaded" || entry.Status.Value == "sleeping" {
328 return entry, nil
329 }
330 if entry.Status.Failed {
331 if entry.Status.ExitCode != nil {
332 return LlamaCPPModel{}, fmt.Errorf("model exited with code %d", *entry.Status.ExitCode)
333 }
334 return LlamaCPPModel{}, errors.New("model failed to load")
335 }
336 }
337 }
338 }
339}
340
341func (c *LlamaCPPClient) UnloadAndWait(ctx context.Context, model string) error {
342 if err := c.Unload(ctx, model); err != nil {
343 return err
344 }
345 ticker := time.NewTicker(100 * time.Millisecond)
346 defer ticker.Stop()
347 for {
348 models, err := c.List(ctx, false)
349 if err != nil {
350 return err
351 }
352 found := false
353 for _, entry := range models {
354 if entry.ID == model {
355 found = true
356 if entry.Status.Value == "unloaded" {
357 return nil
358 }
359 }
360 }
361 if !found {
362 return nil
363 }
364 select {
365 case <-ctx.Done():
366 return ctx.Err()
367 case <-ticker.C:
368 }
369 }
370}
371
372func (c *LlamaCPPClient) DownloadAndWait(ctx context.Context, model string, update func(LlamaCPPProgress)) ([]LlamaCPPModel, error) {
373 watchCtx, cancel := context.WithCancel(ctx)
374 defer cancel()
375 events := make(chan llamaCPPEvent, 8)
376 go func() { _ = c.watch(watchCtx, events) }()
377 if err := c.Download(ctx, model); err != nil {
378 return nil, err
379 }
380 update(LlamaCPPProgress{Message: "Downloading model"})
381 ticker := time.NewTicker(500 * time.Millisecond)
382 defer ticker.Stop()
383 sawDownload := false
384 polls := 0
385 for {
386 select {
387 case <-ctx.Done():
388 return nil, ctx.Err()
389 case event := <-events:
390 if event.Model != model {
391 continue
392 }
393 switch event.Event {
394 case "download_progress":
395 sawDownload = true
396 var payload struct {
397 Progress map[string]LlamaCPPBytes `json:"progress"`
398 }
399 if json.Unmarshal(event.Data, &payload) == nil {
400 if progress, ok := downloadProgress(payload.Progress); ok {
401 update(progress)
402 }
403 }
404 case "download_failed":
405 return nil, errors.New("model download failed")
406 case "download_finished":
407 return c.List(ctx, true)
408 }
409 case <-ticker.C:
410 models, err := c.List(ctx, false)
411 if err != nil {
412 return nil, err
413 }
414 polls++
415 for _, entry := range models {
416 if entry.ID != model {
417 continue
418 }
419 if entry.Status.Value == "downloading" {
420 sawDownload = true
421 files := entry.Status.Progress
422 if len(files) == 0 {
423 files = entry.Progress
424 }
425 if progress, ok := downloadProgress(files); ok {
426 update(progress)
427 }
428 } else if sawDownload || polls >= 2 {
429 return c.List(ctx, true)
430 }
431 }
432 }
433 }
434}
435
436// LlamaCPPModels converts loaded router entries into zut model metadata.
437func LlamaCPPModels(models []LlamaCPPModel, serverURL string) []Model {
438 baseURL, err := LlamaCPPInferenceURL(serverURL)
439 if err != nil {
440 return nil
441 }
442 var out []Model
443 for _, entry := range models {
444 if entry.Status.Value != "loaded" && entry.Status.Value != "sleeping" {
445 continue
446 }
447 contextWindow := entry.Meta.Context
448 if contextWindow <= 0 {
449 contextWindow = entry.Meta.TrainContext
450 }
451 if contextWindow <= 0 {
452 contextWindow = 128000
453 }
454 maxOutput := min(16384, contextWindow)
455 out = append(out, Model{
456 Provider: LlamaCPPProviderID, ID: entry.ID, DisplayName: entry.ID,
457 ContextWindow: contextWindow, MaxOutput: maxOutput,
458 BaseURL: baseURL, Source: "live",
459 })
460 }
461 return out
462}
463
464// FormatBytes renders byte counts using binary units.
465func FormatBytes(bytes int64) string {
466 if bytes < 1024 {
467 return fmt.Sprintf("%d B", bytes)
468 }
469 units := []string{"KiB", "MiB", "GiB", "TiB"}
470 value := float64(bytes) / 1024
471 unit := units[0]
472 for index := 1; index < len(units) && value >= 1024; index++ {
473 value /= 1024
474 unit = units[index]
475 }
476 if value >= 10 {
477 return fmt.Sprintf("%.1f %s", value, unit)
478 }
479 return fmt.Sprintf("%.2f %s", value, unit)
480}
481
482// HuggingFaceModel is one GGUF repository search result.
483type HuggingFaceModel struct {
484 ID string `json:"id"`
485 Downloads int64 `json:"downloads"`
486}
487
488type HuggingFaceQuantization struct {
489 Name string
490 Size int64
491 HasSize bool
492}
493
494type HuggingFaceModelDetails struct {
495 ID string
496 Gated string
497 Quantizations []HuggingFaceQuantization
498}
499
500type HuggingFaceClient struct {
501 Token string
502 BaseURL string
503 HTTP *http.Client
504}
505
506func NewHuggingFaceClient(token string) *HuggingFaceClient {
507 return &HuggingFaceClient{Token: token, BaseURL: defaultHuggingFaceURL, HTTP: &http.Client{Timeout: llamaCPPRequestTimeout}}
508}
509
510func (c *HuggingFaceClient) request(ctx context.Context, path string, out any) error {
511 req, err := http.NewRequestWithContext(ctx, http.MethodGet, strings.TrimRight(c.BaseURL, "/")+path, nil)
512 if err != nil {
513 return err
514 }
515 if c.Token != "" {
516 req.Header.Set("authorization", "Bearer "+c.Token)
517 }
518 resp, err := c.HTTP.Do(req)
519 if err != nil {
520 return err
521 }
522 defer resp.Body.Close()
523 if resp.StatusCode < 200 || resp.StatusCode >= 300 {
524 if resp.StatusCode == http.StatusTooManyRequests {
525 return errors.New("hugging face rate limit reached")
526 }
527 return fmt.Errorf("hugging face returned HTTP %d", resp.StatusCode)
528 }
529 if err := json.NewDecoder(resp.Body).Decode(out); err != nil {
530 return fmt.Errorf("parse Hugging Face response: %w", err)
531 }
532 return nil
533}
534
535func (c *HuggingFaceClient) Search(ctx context.Context, query string) ([]HuggingFaceModel, error) {
536 params := url.Values{"search": {query}, "filter": {"gguf"}, "sort": {"downloads"}, "direction": {"-1"}, "limit": {"20"}}
537 var models []HuggingFaceModel
538 if err := c.request(ctx, "/api/models?"+params.Encode(), &models); err != nil {
539 return nil, err
540 }
541 return models, nil
542}
543
544var quantizationPattern = regexp.MustCompile(`(?i)(?:^|[-_.])((?:UD-)?(?:IQ[0-9](?:_[A-Z0-9]+)+|Q[0-9](?:_[A-Z0-9]+)+|BF16|F16|F32|MXFP[0-9](?:_[A-Z0-9]+)*))$`)
545var shardSuffixPattern = regexp.MustCompile(`-[0-9]{5}-of-[0-9]{5}$`)
546
547func (c *HuggingFaceClient) Details(ctx context.Context, id string) (HuggingFaceModelDetails, error) {
548 parts := strings.Split(id, "/")
549 for i := range parts {
550 parts[i] = url.PathEscape(parts[i])
551 }
552 encoded := strings.Join(parts, "/")
553 var payload struct {
554 ID string `json:"id"`
555 Gated json.RawMessage `json:"gated"`
556 Siblings []struct {
557 Filename string `json:"rfilename"`
558 Size *int64 `json:"size"`
559 } `json:"siblings"`
560 }
561 if err := c.request(ctx, "/api/models/"+encoded+"?blobs=true", &payload); err != nil {
562 return HuggingFaceModelDetails{}, err
563 }
564 type quantSize struct {
565 total int64
566 complete bool
567 }
568 sizes := map[string]quantSize{}
569 for _, file := range payload.Siblings {
570 name := filepath.Base(file.Filename)
571 if !strings.HasSuffix(strings.ToLower(name), ".gguf") || strings.HasPrefix(strings.ToLower(name), "mmproj") {
572 continue
573 }
574 stem := shardSuffixPattern.ReplaceAllString(name[:len(name)-5], "")
575 match := quantizationPattern.FindStringSubmatch(stem)
576 if len(match) < 2 {
577 continue
578 }
579 quant := strings.ToUpper(match[1])
580 current, ok := sizes[quant]
581 if !ok {
582 current.complete = true
583 }
584 if file.Size == nil {
585 current.complete = false
586 } else {
587 current.total += *file.Size
588 }
589 sizes[quant] = current
590 }
591 var quantizations []HuggingFaceQuantization
592 for name, size := range sizes {
593 quantizations = append(quantizations, HuggingFaceQuantization{Name: name, Size: size.total, HasSize: size.complete})
594 }
595 sort.Slice(quantizations, func(i, j int) bool {
596 if quantizations[i].Name == "Q4_K_M" {
597 return true
598 }
599 if quantizations[j].Name == "Q4_K_M" {
600 return false
601 }
602 if quantizations[i].HasSize != quantizations[j].HasSize {
603 return quantizations[i].HasSize
604 }
605 if quantizations[i].Size != quantizations[j].Size {
606 return quantizations[i].Size < quantizations[j].Size
607 }
608 return quantizations[i].Name < quantizations[j].Name
609 })
610 gated := ""
611 _ = json.Unmarshal(payload.Gated, &gated)
612 if payload.ID == "" {
613 payload.ID = id
614 }
615 return HuggingFaceModelDetails{ID: payload.ID, Gated: gated, Quantizations: quantizations}, nil
616}
617
618// FindHuggingFaceToken follows the standard Hugging Face token locations.
619func FindHuggingFaceToken() string {
620 if token := strings.TrimSpace(os.Getenv("HF_TOKEN")); token != "" {
621 return token
622 }
623 home, _ := os.UserHomeDir()
624 paths := []string{os.Getenv("HF_TOKEN_PATH")}
625 if value := os.Getenv("HF_HOME"); value != "" {
626 paths = append(paths, filepath.Join(value, "token"))
627 }
628 if value := os.Getenv("XDG_CACHE_HOME"); value != "" {
629 paths = append(paths, filepath.Join(value, "huggingface", "token"))
630 }
631 if home != "" {
632 paths = append(paths, filepath.Join(home, ".cache", "huggingface", "token"))
633 }
634 for _, path := range paths {
635 if path == "" {
636 continue
637 }
638 if data, err := os.ReadFile(path); err == nil {
639 if token := strings.TrimSpace(string(data)); token != "" {
640 return token
641 }
642 }
643 }
644 return ""
645}