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