mirror of
https://github.com/mudler/LocalAI.git
synced 2026-06-26 09:26:55 -04:00
64a4351f3a
LocalAI pulls models from OCI registries (via go-containerregistry), the Ollama registry, and OCI blob stores (via oras), but every request went out with the underlying library's generic User-Agent, so registry operators had no way to attribute traffic to LocalAI. Add an oci.UserAgent() helper that returns "LocalAI" (or "LocalAI/<version>" when the binary is built with a version stamp via internal.Version) and wire it into all three pull paths: - pkg/oci/image.go: remote.WithUserAgent on the go-containerregistry image and digest requests - pkg/oci/ollama.go: a User-Agent header on the Ollama manifest request - pkg/oci/blob.go: a LocalAI User-Agent on the oras blob client. This mirrors oras' auth.DefaultClient (same retry.DefaultClient policy); only the advertised User-Agent changes. Implements #6258. Assisted-by: Claude:claude-opus-4-8 golangci-lint Signed-off-by: Vijay Sai <vijaysaijnv@gmail.com>
93 lines
2.5 KiB
Go
93 lines
2.5 KiB
Go
package oci
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
|
|
ocispec "github.com/opencontainers/image-spec/specs-go/v1"
|
|
|
|
"github.com/mudler/LocalAI/pkg/httpclient"
|
|
)
|
|
|
|
// Define the main struct for the JSON data
|
|
type Manifest struct {
|
|
SchemaVersion int `json:"schemaVersion"`
|
|
MediaType string `json:"mediaType"`
|
|
Config Config `json:"config"`
|
|
Layers []LayerDetail `json:"layers"`
|
|
}
|
|
|
|
// Define the struct for the "config" section
|
|
type Config struct {
|
|
Digest string `json:"digest"`
|
|
MediaType string `json:"mediaType"`
|
|
Size int `json:"size"`
|
|
}
|
|
|
|
// Define the struct for each item in the "layers" array
|
|
type LayerDetail struct {
|
|
Digest string `json:"digest"`
|
|
MediaType string `json:"mediaType"`
|
|
Size int `json:"size"`
|
|
}
|
|
|
|
func OllamaModelManifest(image string) (*Manifest, error) {
|
|
// parse the repository and tag from `image`. `image` should be for e.g. gemma:2b, or foobar/gemma:2b
|
|
|
|
// if there is a : in the image, then split it
|
|
// if there is no : in the image, then assume it is the latest tag
|
|
tag, repository, image := ParseImageParts(image)
|
|
|
|
// get e.g. https://registry.ollama.ai/v2/library/llama3/manifests/latest
|
|
req, err := http.NewRequest("GET", "https://registry.ollama.ai/v2/"+repository+"/"+image+"/manifests/"+tag, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
req.Header.Set("Accept", "application/vnd.docker.distribution.manifest.v2+json")
|
|
req.Header.Set("User-Agent", UserAgent())
|
|
client := httpclient.New(httpclient.WithFollowRedirects())
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
// parse the JSON response
|
|
var manifest Manifest
|
|
err = json.NewDecoder(resp.Body).Decode(&manifest)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return &manifest, nil
|
|
}
|
|
|
|
func OllamaModelBlob(image string) (string, error) {
|
|
manifest, err := OllamaModelManifest(image)
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
// find a application/vnd.ollama.image.model in the mediaType
|
|
|
|
for _, layer := range manifest.Layers {
|
|
if layer.MediaType == "application/vnd.ollama.image.model" {
|
|
return layer.Digest, nil
|
|
}
|
|
}
|
|
|
|
return "", nil
|
|
}
|
|
|
|
func OllamaFetchModel(ctx context.Context, image string, output string, statusWriter func(ocispec.Descriptor) io.Writer) error {
|
|
_, repository, imageNoTag := ParseImageParts(image)
|
|
|
|
blobID, err := OllamaModelBlob(image)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
return FetchImageBlob(ctx, fmt.Sprintf("registry.ollama.ai/%s/%s", repository, imageNoTag), blobID, output, statusWriter)
|
|
}
|