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