mirror of
https://github.com/mudler/LocalAI.git
synced 2026-06-22 15:49:12 -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>
65 lines
1.6 KiB
Go
65 lines
1.6 KiB
Go
package oci
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"os"
|
|
|
|
"github.com/mudler/LocalAI/pkg/xio"
|
|
ocispec "github.com/opencontainers/image-spec/specs-go/v1"
|
|
|
|
oras "oras.land/oras-go/v2"
|
|
"oras.land/oras-go/v2/registry/remote"
|
|
"oras.land/oras-go/v2/registry/remote/auth"
|
|
"oras.land/oras-go/v2/registry/remote/retry"
|
|
)
|
|
|
|
func FetchImageBlob(ctx context.Context, r, reference, dst string, statusReader func(ocispec.Descriptor) io.Writer) error {
|
|
// 0. Create a file store for the output
|
|
fs, err := os.Create(dst)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer fs.Close()
|
|
|
|
// 1. Connect to a remote repository
|
|
repo, err := remote.NewRepository(r)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to create repository: %v", err)
|
|
}
|
|
repo.SkipReferrersGC = true
|
|
|
|
// Identify LocalAI to the registry. This mirrors oras' auth.DefaultClient
|
|
// (same retry policy) but advertises a LocalAI User-Agent instead of the
|
|
// library default.
|
|
client := &auth.Client{
|
|
Client: retry.DefaultClient,
|
|
Cache: auth.NewCache(),
|
|
}
|
|
client.SetUserAgent(UserAgent())
|
|
repo.Client = client
|
|
|
|
// https://github.com/oras-project/oras/blob/main/cmd/oras/internal/option/remote.go#L364
|
|
// https://github.com/oras-project/oras/blob/main/cmd/oras/root/blob/fetch.go#L136
|
|
desc, reader, err := oras.Fetch(ctx, repo.Blobs(), reference, oras.DefaultFetchOptions)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to fetch image: %v", err)
|
|
}
|
|
|
|
if statusReader != nil {
|
|
// 3. Write the file to the file store
|
|
_, err = xio.Copy(ctx, io.MultiWriter(fs, statusReader(desc)), reader)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
} else {
|
|
_, err = xio.Copy(ctx, fs, reader)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|