refactor: consolidate usage of GetURI (#674)

Signed-off-by: mudler <mudler@localai.io>

pkg/utils/uri.go

@@ -1,12 +1,34 @@
 package utils
 
 import (
+	"fmt"
 	"io/ioutil"
 	"net/http"
 	"strings"
 )
 
-func GetURI(url string, f func(i []byte) error) error {
+const (
+	githubURI = "github:"
+)
+
+func GetURI(url string, f func(url string, i []byte) error) error {
+	if strings.HasPrefix(url, githubURI) {
+		parts := strings.Split(url, ":")
+		repoParts := strings.Split(parts[1], "@")
+		branch := "main"
+
+		if len(repoParts) > 1 {
+			branch = repoParts[1]
+		}
+
+		repoPath := strings.Split(repoParts[0], "/")
+		org := repoPath[0]
+		project := repoPath[1]
+		projectPath := strings.Join(repoPath[2:], "/")
+
+		url = fmt.Sprintf("https://raw.githubusercontent.com/%s/%s/%s/%s", org, project, branch, projectPath)
+	}
+
 	if strings.HasPrefix(url, "file://") {
 		rawURL := strings.TrimPrefix(url, "file://")
 		// Read the response body
@@ -16,7 +38,7 @@ func GetURI(url string, f func(i []byte) error) error {
 		}
 
 		// Unmarshal YAML data into a struct
-		return f(body)
+		return f(url, body)
 	}
 
 	// Send a GET request to the URL
@@ -33,5 +55,5 @@ func GetURI(url string, f func(i []byte) error) error {
 	}
 
 	// Unmarshal YAML data into a struct
-	return f(body)
+	return f(url, body)
 }

pkg/utils/uri_test.go

@@ -0,0 +1,36 @@
+package utils_test
+
+import (
+	. "github.com/go-skynet/LocalAI/pkg/utils"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("Gallery API tests", func() {
+	Context("URI", func() {
+		It("parses github with a branch", func() {
+			Expect(
+				GetURI("github:go-skynet/model-gallery/gpt4all-j.yaml", func(url string, i []byte) error {
+					Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml"))
+					return nil
+				}),
+			).ToNot(HaveOccurred())
+		})
+		It("parses github without a branch", func() {
+			Expect(
+				GetURI("github:go-skynet/model-gallery/gpt4all-j.yaml@main", func(url string, i []byte) error {
+					Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml"))
+					return nil
+				}),
+			).ToNot(HaveOccurred())
+		})
+		It("parses github with urls", func() {
+			Expect(
+				GetURI("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml", func(url string, i []byte) error {
+					Expect(url).To(Equal("https://raw.githubusercontent.com/go-skynet/model-gallery/main/gpt4all-j.yaml"))
+					return nil
+				}),
+			).ToNot(HaveOccurred())
+		})
+	})
+})

pkg/utils/utils_suite_test.go

@@ -0,0 +1,13 @@
+package utils_test
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestUtils(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Utils test suite")
+}