From 8bdd0e770368a4042ef1c293e8cf968ff9641f8c Mon Sep 17 00:00:00 2001
From: thxCode <thxcode0824@gmail.com>
Date: Wed, 21 Aug 2024 15:29:11 +0800
Subject: [PATCH] feat: ollama base url

Signed-off-by: thxCode <thxcode0824@gmail.com>
---
 cmd/gguf-parser/README.md |  5 ++-
 cmd/gguf-parser/main.go   | 11 +++++-
 ollama_model.go           | 21 ++++++++++-
 ollama_model_option.go    | 79 +++++++++++++++++++++++++++++++++++++++
 4 files changed, 112 insertions(+), 4 deletions(-)
 create mode 100644 ollama_model_option.go

diff --git a/cmd/gguf-parser/README.md b/cmd/gguf-parser/README.md
index 4830351..d084065 100644
--- a/cmd/gguf-parser/README.md
+++ b/cmd/gguf-parser/README.md
@@ -84,8 +84,9 @@ GLOBAL OPTIONS:
 
    Model/Remote/Ollama
 
-   --ol-model value  Model name of Ollama, e.g. gemma2.
-   --ol-usage        Specify respecting the extending layers introduced by Ollama, works with --ol-model, which affects the usage estimation. (default: false)
+   --ol-base-url value  Model base URL of Ollama, e.g. https://registry.ollama.ai. (default: "https://registry.ollama.ai")
+   --ol-model value     Model name of Ollama, e.g. gemma2.
+   --ol-usage           Specify respecting the extending layers introduced by Ollama, works with --ol-model, which affects the usage estimation. (default: false)
 
    Output
 
diff --git a/cmd/gguf-parser/main.go b/cmd/gguf-parser/main.go
index 798090a..e06bc0e 100644
--- a/cmd/gguf-parser/main.go
+++ b/cmd/gguf-parser/main.go
@@ -223,6 +223,14 @@ func main() {
 					"works with --ms-repo/--ms-file pair or --ms-draft-repo/--ms-draft-file pair. " +
 					"See https://modelscope.cn/my/myaccesstoken.",
 			},
+			&cli.StringFlag{
+				Destination: &olBaseURL,
+				Value:       olBaseURL,
+				Category:    "Model/Remote/Ollama",
+				Name:        "ol-base-url",
+				Usage: "Model base URL of Ollama, e.g. " +
+					"https://registry.ollama.ai.",
+			},
 			&cli.StringFlag{
 				Destination: &olModel,
 				Value:       olModel,
@@ -561,6 +569,7 @@ var (
 	msDraftRepo  string // for estimate
 	msDraftFile  string // for estimate
 	msToken      string
+	olBaseURL    = "https://registry.ollama.ai"
 	olModel      string
 	olUsage      bool
 	// load options
@@ -730,7 +739,7 @@ func mainAction(c *cli.Context) error {
 			}
 			gf, err = ParseGGUFFileFromModelScope(ctx, msRepo, msFile, ropts...)
 		case olModel != "":
-			om := ParseOllamaModel(olModel)
+			om := ParseOllamaModel(olModel, SetOllamaModelBaseURL(olBaseURL))
 			gf, err = ParseGGUFFileFromOllamaModel(ctx, om, ropts...)
 			if om != nil && olUsage {
 				// Parameters override.
diff --git a/ollama_model.go b/ollama_model.go
index c64fd63..3ffd950 100644
--- a/ollama_model.go
+++ b/ollama_model.go
@@ -69,17 +69,36 @@ type (
 
 // ParseOllamaModel parses the given Ollama model string,
 // and returns the OllamaModel, or nil if the model is invalid.
-func ParseOllamaModel(model string) *OllamaModel {
+func ParseOllamaModel(model string, opts ...OllamaModelOption) *OllamaModel {
 	if model == "" {
 		return nil
 	}
 
+	var o _OllamaModelOptions
+	for _, opt := range opts {
+		opt(&o)
+	}
+
 	om := OllamaModel{
 		Schema:    OllamaDefaultScheme,
 		Registry:  OllamaDefaultRegistry,
 		Namespace: OllamaDefaultNamespace,
 		Tag:       OllamaDefaultTag,
 	}
+	{
+		if o.DefaultScheme != "" {
+			om.Schema = o.DefaultScheme
+		}
+		if o.DefaultRegistry != "" {
+			om.Registry = o.DefaultRegistry
+		}
+		if o.DefaultNamespace != "" {
+			om.Namespace = o.DefaultNamespace
+		}
+		if o.DefaultTag != "" {
+			om.Tag = o.DefaultTag
+		}
+	}
 
 	m := model
 
diff --git a/ollama_model_option.go b/ollama_model_option.go
new file mode 100644
index 0000000..325d4eb
--- /dev/null
+++ b/ollama_model_option.go
@@ -0,0 +1,79 @@
+package gguf_parser
+
+import (
+	"net/url"
+	"strings"
+)
+
+type (
+	_OllamaModelOptions struct {
+		DefaultScheme    string
+		DefaultRegistry  string
+		DefaultNamespace string
+		DefaultTag       string
+	}
+	OllamaModelOption func(*_OllamaModelOptions)
+)
+
+// SetOllamaModelBaseURL parses the given base URL,
+// and sets default schema/registry for OllamaModel.
+func SetOllamaModelBaseURL(baseURL string) OllamaModelOption {
+	baseURL = strings.TrimSpace(baseURL)
+	return func(o *_OllamaModelOptions) {
+		if baseURL == "" {
+			return
+		}
+
+		if !strings.Contains(baseURL, "://") {
+			baseURL = "https://" + baseURL
+		}
+
+		u, err := url.Parse(baseURL)
+		if err != nil {
+			return
+		}
+
+		o.DefaultScheme = u.Scheme
+		o.DefaultRegistry = u.Host
+	}
+}
+
+// SetOllamaModelDefaultScheme sets the default scheme for OllamaModel.
+func SetOllamaModelDefaultScheme(scheme string) OllamaModelOption {
+	return func(o *_OllamaModelOptions) {
+		if scheme == "" {
+			return
+		}
+		o.DefaultScheme = scheme
+	}
+}
+
+// SetOllamaModelDefaultRegistry sets the default registry for OllamaModel.
+func SetOllamaModelDefaultRegistry(registry string) OllamaModelOption {
+	return func(o *_OllamaModelOptions) {
+		if registry == "" {
+			return
+		}
+		o.DefaultRegistry = registry
+	}
+}
+
+// SetOllamaModelDefaultNamespace sets the default namespace for OllamaModel.
+func SetOllamaModelDefaultNamespace(namespace string) OllamaModelOption {
+	return func(o *_OllamaModelOptions) {
+		if namespace == "" {
+			return
+		}
+		o.DefaultNamespace = namespace
+	}
+}
+
+// SetOllamaModelDefaultTag sets the default tag for OllamaModel.
+func SetOllamaModelDefaultTag(tag string) OllamaModelOption {
+	return func(o *_OllamaModelOptions) {
+		if tag == "" {
+			return
+		}
+		o.DefaultTag = tag
+	}
+}