diff --git a/docs/content/getting-started/models.md b/docs/content/getting-started/models.md index cf949f715..b05f05728 100644 --- a/docs/content/getting-started/models.md +++ b/docs/content/getting-started/models.md @@ -131,6 +131,10 @@ local-ai run ollama://gemma:2b local-ai run oci://localai/phi-2:latest ``` +{{% notice note %}} +When pulling models from Ollama or OCI registries, LocalAI identifies itself with a `LocalAI/` `User-Agent` header so registry operators can attribute usage to LocalAI. +{{% /notice %}} + ### Run Models via URI To run models via URI, specify a URI to a model file or a configuration file when starting LocalAI. Valid syntax includes: diff --git a/pkg/oci/blob.go b/pkg/oci/blob.go index 0f5a2cf66..e034c4162 100644 --- a/pkg/oci/blob.go +++ b/pkg/oci/blob.go @@ -11,6 +11,8 @@ import ( 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 { @@ -28,6 +30,16 @@ func FetchImageBlob(ctx context.Context, r, reference, dst string, statusReader } 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) diff --git a/pkg/oci/image.go b/pkg/oci/image.go index 2d00c3479..4dad02c7d 100644 --- a/pkg/oci/image.go +++ b/pkg/oci/image.go @@ -176,6 +176,7 @@ func GetImage(targetImage, targetPlatform string, auth *registrytypes.AuthConfig opts := []remote.Option{ remote.WithTransport(tr), remote.WithPlatform(*platform), + remote.WithUserAgent(UserAgent()), } if auth != nil { opts = append(opts, remote.WithAuth(staticAuth{auth})) @@ -223,6 +224,7 @@ func GetImageDigest(targetImage, targetPlatform string, auth *registrytypes.Auth opts := []remote.Option{ remote.WithTransport(tr), remote.WithPlatform(*platform), + remote.WithUserAgent(UserAgent()), } if auth != nil { opts = append(opts, remote.WithAuth(staticAuth{auth})) diff --git a/pkg/oci/ollama.go b/pkg/oci/ollama.go index 2fb928281..f0a874013 100644 --- a/pkg/oci/ollama.go +++ b/pkg/oci/ollama.go @@ -47,6 +47,7 @@ func OllamaModelManifest(image string) (*Manifest, error) { 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 { diff --git a/pkg/oci/useragent.go b/pkg/oci/useragent.go new file mode 100644 index 000000000..82277c70d --- /dev/null +++ b/pkg/oci/useragent.go @@ -0,0 +1,19 @@ +package oci + +import ( + "fmt" + + "github.com/mudler/LocalAI/internal" +) + +// UserAgent returns the User-Agent string LocalAI sends on outbound registry +// requests (OCI registries and Ollama). It identifies the client as LocalAI +// and, when the binary was built with a version stamp, appends it so registries +// can attribute client-side usage to LocalAI rather than to the generic +// User-Agent of the underlying transport library. +func UserAgent() string { + if internal.Version == "" { + return "LocalAI" + } + return fmt.Sprintf("LocalAI/%s", internal.Version) +} diff --git a/pkg/oci/useragent_test.go b/pkg/oci/useragent_test.go new file mode 100644 index 000000000..14a10534c --- /dev/null +++ b/pkg/oci/useragent_test.go @@ -0,0 +1,32 @@ +package oci_test + +import ( + "github.com/mudler/LocalAI/internal" + . "github.com/mudler/LocalAI/pkg/oci" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("OCI", func() { + Context("UserAgent", func() { + var savedVersion string + + BeforeEach(func() { + savedVersion = internal.Version + }) + + AfterEach(func() { + internal.Version = savedVersion + }) + + It("identifies as LocalAI when no version is stamped", func() { + internal.Version = "" + Expect(UserAgent()).To(Equal("LocalAI")) + }) + + It("appends the build version when one is stamped", func() { + internal.Version = "v3.2.1" + Expect(UserAgent()).To(Equal("LocalAI/v3.2.1")) + }) + }) +})