pulling works now

Signed-off-by: Ettore Di Giacinto <[email protected]>

pkg/oci/blob.go

@@ -10,8 +10,8 @@ import (
 	"oras.land/oras-go/v2/registry/remote"
 )
 
-func FetchImage(r, reference, dst string) error {
-	// 0. Create a file store
+func FetchImageBlob(r, reference, dst string) error {
+	// 0. Create a file store for the output
 	fs, err := os.Create(dst)
 	if err != nil {
 		return err
@@ -24,23 +24,11 @@ func FetchImage(r, reference, dst string) error {
 	if err != nil {
 		return fmt.Errorf("failed to create repository: %v", err)
 	}
-	// https://github.com/sozercan/aikit/commit/6f7cbd8744f00971d05dc46ca5f20deefd569a49
+	repo.SkipReferrersGC = true
+
 	// 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
-	fmt.Println(repo.Reference.Registry)
-	// Note: The below code can be omitted if authentication is not required
-	//	repo.Client = &auth.Client{
-	//		Client: retry.DefaultClient,
-	//		Cache:  auth.NewCache(),
-	//	Credential: auth.StaticCredential(reg, auth.Credential{
-	//		Username: "username",
-	//		Password: "password",
-	//	}),
-	//}
-
-	// 2. Copy from the remote repository to the file store
-
-	desc, reader, err := oras.Fetch(ctx, repo, reference, oras.DefaultFetchOptions)
+	_, reader, err := oras.Fetch(ctx, repo.Blobs(), reference, oras.DefaultFetchOptions)
 	if err != nil {
 		return fmt.Errorf("failed to fetch image: %v", err)
 	}
@@ -50,8 +38,5 @@ func FetchImage(r, reference, dst string) error {
 	if err != nil {
 		return err
 	}
-
-	fmt.Println("manifest descriptor:", desc)
-
 	return nil
 }

pkg/oci/blob_test.go

@@ -7,16 +7,10 @@ import (
 )
 
 var _ = Describe("OCI", func() {
-
-	Context("when template is loaded successfully", func() {
-		FIt("should evaluate the template correctly", func() {
-
-			// https://registry.ollama.ai/v2/library/gemma/manifests/2b
-			// "application/vnd.ollama.image.model"
-			err := FetchImage("registry.ollama.ai/library/gemma", "sha256:c1864a5eb19305c40519da12cc543519e48a0697ecd30e15d5ac228644957d12", "/tmp/foo")
+	Context("pulling images", func() {
+		It("should fetch blobs correctly", func() {
+			err := FetchImageBlob("registry.ollama.ai/library/gemma", "sha256:c1864a5eb19305c40519da12cc543519e48a0697ecd30e15d5ac228644957d12", "/tmp/foo")
 			Expect(err).NotTo(HaveOccurred())
-
 		})
 	})
-
 })

pkg/oci/image.go

@@ -68,6 +68,23 @@ func ExtractOCIImage(img v1.Image, targetDestination string) error {
 	return err
 }
 
+func ParseImageParts(image string) (tag, repository, dstimage string) {
+	tag = "latest"
+	repository = "library"
+	if strings.Contains(image, ":") {
+		parts := strings.Split(image, ":")
+		image = parts[0]
+		tag = parts[1]
+	}
+	if strings.Contains("/", image) {
+		parts := strings.Split(image, "/")
+		repository = parts[0]
+		image = parts[1]
+	}
+	dstimage = image
+	return tag, repository, image
+}
+
 // GetImage if returns the proper image to pull with transport and auth
 // tries local daemon first and then fallbacks into remote
 // if auth is nil, it will try to use the default keychain https://github.com/google/go-containerregistry/tree/main/pkg/authn#tldr-for-consumers-of-this-package

pkg/oci/ollama.go

@@ -0,0 +1,88 @@
+package oci
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+)
+
+// 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")
+	client := &http.Client{}
+	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(image string, output string) error {
+	_, repository, imageNoTag := ParseImageParts(image)
+
+	blobID, err := OllamaModelBlob(image)
+	if err != nil {
+		return err
+	}
+	fmt.Println("Downloading blob", blobID)
+	fmt.Println("Downloading to", output)
+	fmt.Println("Downloading from", fmt.Sprintf("registry.ollama.ai/%s/%s", repository, imageNoTag))
+
+	return FetchImageBlob(fmt.Sprintf("registry.ollama.ai/%s/%s", repository, imageNoTag), blobID, output)
+}

pkg/oci/ollama_test.go

@@ -0,0 +1,16 @@
+package oci_test
+
+import (
+	. "github.com/go-skynet/LocalAI/pkg/oci" // Update with your module path
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("OCI", func() {
+	Context("ollama", func() {
+		FIt("pulls model files", func() {
+			err := OllamaFetchModel("gemma:2b", "/tmp/foo")
+			Expect(err).NotTo(HaveOccurred())
+		})
+	})
+})