Ollama: allow any ollama model

pkg/cmd/cmd.go

@@ -96,6 +96,7 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin
 
 	loginCmdStore := fsStore.WithNoAuthHTTPClient(
 		store.NewNoAuthHTTPClient(conf.GetBrevAPIURl()),
+		store.NewOllamaHTTPClient(conf.GetOllamaAPIURL()),
 	).
 		WithAuth(loginAuth, store.WithDebug(conf.GetDebugHTTP()))
 
@@ -111,6 +112,7 @@ func NewBrevCommand() *cobra.Command { //nolint:funlen,gocognit,gocyclo // defin
 	}
 	noAuthCmdStore := fsStore.WithNoAuthHTTPClient(
 		store.NewNoAuthHTTPClient(conf.GetBrevAPIURl()),
+		store.NewOllamaHTTPClient(conf.GetOllamaAPIURL()),
 	)
 	noLoginCmdStore := noAuthCmdStore.WithAuth(noLoginAuth)
 

pkg/cmd/ollama/ollama.go

@@ -4,6 +4,7 @@ package ollama
 import (
 	_ "embed"
 	"fmt"
+	"strings"
 	"time"
 
 	"github.com/google/uuid"
@@ -25,7 +26,6 @@ var (
 	ollamaExample = `
   brev ollama --model llama3
 	`
-	modelTypes = []string{"llama3"}
 )
 
 //go:embed ollamaverb.yaml
@@ -39,15 +39,29 @@ type OllamaStore interface {
 	GetWorkspace(workspaceID string) (*entity.Workspace, error)
 	BuildVerbContainer(workspaceID string, verbYaml string) (*store.BuildVerbRes, error)
 	ModifyPublicity(workspace *entity.Workspace, applicationName string, publicity bool) (*entity.Tunnel, error)
+	ValidateOllamaModel(model string, tag string) (bool, error)
 }
 
-func validateModelType(modelType string) bool {
-	for _, v := range modelTypes {
-		if modelType == v {
-			return true
-		}
+func validateModelType(input string, ollamaStore OllamaStore) (bool, error) {
+	var model string
+	var tag string
+
+	split := strings.Split(input, ":")
+	switch len(split) {
+	case 2:
+		model = split[0]
+		tag = split[1]
+	case 1:
+		model = input
+		tag = "latest"
+	default:
+		return false, fmt.Errorf("invalid model type: %s", input)
 	}
-	return false
+	valid, err := ollamaStore.ValidateOllamaModel(model, tag)
+	if err != nil {
+		return false, fmt.Errorf("error validating model: %s", err)
+	}
+	return valid, nil
 }
 
 func NewCmdOllama(t *terminal.Terminal, ollamaStore OllamaStore) *cobra.Command {
@@ -67,7 +81,10 @@ func NewCmdOllama(t *terminal.Terminal, ollamaStore OllamaStore) *cobra.Command
 				return fmt.Errorf("model type must be specified")
 			}
 
-			isValid := validateModelType(model)
+			isValid, valErr := validateModelType(model, ollamaStore)
+			if valErr != nil {
+				return valErr
+			}
 			if !isValid {
 				return fmt.Errorf("invalid model type: %s", model)
 			}

pkg/config/config.go

@@ -15,6 +15,7 @@ const (
 	defaultWorkspaceTemplate EnvVarName = "DEFAULT_WORKSPACE_TEMPLATE"
 	sentryURL                EnvVarName = "DEFAULT_SENTRY_URL"
 	debugHTTP                EnvVarName = "DEBUG_HTTP"
+	ollamaAPIURL             EnvVarName = "OLLAMA_API_URL"
 )
 
 type ConstantsConfig struct{}
@@ -27,6 +28,10 @@ func (c ConstantsConfig) GetBrevAPIURl() string {
 	return getEnvOrDefault(brevAPIURL, "https://brevapi.us-west-2-prod.control-plane.brev.dev")
 }
 
+func (c ConstantsConfig) GetOllamaAPIURL() string {
+	return getEnvOrDefault(ollamaAPIURL, "https://registry.ollama.ai")
+}
+
 func (c ConstantsConfig) GetServiceMeshCoordServerURL() string {
 	return getEnvOrDefault(coordURL, "")
 }

pkg/store/http.go

@@ -17,21 +17,34 @@ type NoAuthHTTPStore struct {
 	FileStore
 	noAuthHTTPClient *NoAuthHTTPClient
 	BasicStore
+	ollamaHTTPClient *OllamaHTTPClient
 }
 
-func (f *FileStore) WithNoAuthHTTPClient(c *NoAuthHTTPClient) *NoAuthHTTPStore {
-	return &NoAuthHTTPStore{*f, c, f.b}
+func (f *FileStore) WithNoAuthHTTPClient(c *NoAuthHTTPClient, o *OllamaHTTPClient) *NoAuthHTTPStore {
+	return &NoAuthHTTPStore{*f, c, f.b, o}
 }
 
 // Used if need new instance to customize settings
 func (n NoAuthHTTPStore) NewNoAuthHTTPStore() *NoAuthHTTPStore {
-	return n.WithNoAuthHTTPClient(NewNoAuthHTTPClient(n.noAuthHTTPClient.restyClient.BaseURL))
+	return n.WithNoAuthHTTPClient(NewNoAuthHTTPClient(n.noAuthHTTPClient.restyClient.BaseURL), NewOllamaHTTPClient(n.ollamaHTTPClient.restyClient.BaseURL))
 }
 
 type NoAuthHTTPClient struct {
 	restyClient *resty.Client
 }
 
+type OllamaHTTPClient struct {
+	restyClient *resty.Client
+}
+
+func NewOllamaHTTPClient(ollamaAPIURL string) *OllamaHTTPClient {
+	restyClient := resty.New().SetBaseURL(ollamaAPIURL)
+
+	return &OllamaHTTPClient{
+		restyClient: restyClient,
+	}
+}
+
 func NewNoAuthHTTPClient(brevAPIURL string) *NoAuthHTTPClient {
 	restyClient := NewRestyClient(brevAPIURL)
 	return &NoAuthHTTPClient{restyClient}
@@ -67,7 +80,7 @@ func (f *FileStore) WithAuthHTTPClient(c *AuthHTTPClient) *AuthHTTPStore {
 	if id == "" {
 		c.restyClient.SetQueryParam("local", "true")
 	}
-	na := f.WithNoAuthHTTPClient(NewNoAuthHTTPClient(c.restyClient.BaseURL))
+	na := f.WithNoAuthHTTPClient(NewNoAuthHTTPClient(c.restyClient.BaseURL), NewOllamaHTTPClient(c.restyClient.BaseURL))
 	return &AuthHTTPStore{NoAuthHTTPStore: *na, authHTTPClient: c}
 }
 

pkg/store/http_test.go

@@ -11,7 +11,7 @@ import (
 
 func MakeMockNoHTTPStore() *NoAuthHTTPStore {
 	fs := MakeMockFileStore()
-	nh := fs.WithNoAuthHTTPClient(NewNoAuthHTTPClient(""))
+	nh := fs.WithNoAuthHTTPClient(NewNoAuthHTTPClient(""), NewOllamaHTTPClient(""))
 	return nh
 }
 

pkg/store/workspace.go

@@ -647,3 +647,74 @@ func (s AuthHTTPStore) ModifyPublicity(workspace *entity.Workspace, applicationN
 	}
 	return &result, nil
 }
+
+type OllamaRegistrySuccessResponse struct {
+	SchemaVersion int           `json:"schemaVersion"`
+	MediaType     string        `json:"mediaType"`
+	Config        OllamaConfig  `json:"config"`
+	Layers        []OllamaLayer `json:"layers"`
+}
+
+type OllamaConfig struct {
+	MediaType string `json:"mediaType"`
+	Size      int    `json:"size"`
+	Digest    string `json:"digest"`
+}
+
+type OllamaLayer struct {
+	MediaType string `json:"mediaType"`
+	Size      int    `json:"size"`
+	Digest    string `json:"digest"`
+}
+
+type OllamaRegistryFailureResponse struct {
+	Errors []OllamaRegistryError `json:"errors"`
+}
+
+type OllamaRegistryError struct {
+	Code    string                    `json:"code"`
+	Message string                    `json:"message"`
+	Detail  OllamaRegistryErrorDetail `json:"detail"`
+}
+
+type OllamaRegistryErrorDetail struct {
+	Tag string `json:"Tag"`
+}
+
+type OllamaModelRequest struct {
+	Model string
+	Tag   string
+}
+
+var (
+	modelNameParamName     = "modelName"
+	tagNameParamName       = "tagName"
+	ollamaModelPathPattern = "v2/library/%s/manifests/%s"
+	ollamaModelPath        = fmt.Sprintf(ollamaModelPathPattern, fmt.Sprintf("{%s}", modelNameParamName), fmt.Sprintf("{%s}", tagNameParamName))
+)
+
+func (s *AuthHTTPStore) ValidateOllamaModel(model string, tag string) (bool, error) {
+	// TODO: building a resty client here because we use this endpoint once. Should there be an OllamaHTTPClient as a part of the AuthHTTPStore or NoAuthHTTPStore?
+	// create a resty client
+	res, err := s.ollamaHTTPClient.restyClient.R().
+		SetHeader("Accept", "application/vnd.docker.distribution.manifest.v2+json").
+		SetPathParam(modelNameParamName, model).
+		SetPathParam(tagNameParamName, tag).
+		Get(ollamaModelPath)
+	if err != nil {
+		return false, breverrors.WrapAndTrace(err)
+	}
+	if res.StatusCode() == 200 { //nolint:gocritic // 200 is a valid status code
+		if err := json.Unmarshal(res.Body(), &OllamaRegistrySuccessResponse{}); err != nil {
+			return false, breverrors.WrapAndTrace(err)
+		}
+		return true, nil
+	} else if res.StatusCode() == 404 {
+		if err := json.Unmarshal(res.Body(), &OllamaRegistryFailureResponse{}); err != nil {
+			return false, breverrors.WrapAndTrace(err)
+		}
+		return false, nil
+	} else {
+		return false, breverrors.New("invalid response from ollama registry")
+	}
+}

tools/go.mod

@@ -1,6 +1,11 @@
 module github.com/brevdev/brev-cli/tools
 
 go 1.21
+<<<<<<< HEAD
+
+toolchain go1.22.2
+=======
+>>>>>>> main
 
 require (
 	github.com/golangci/golangci-lint v1.57.2
@@ -92,6 +97,10 @@ require (
 	github.com/bkielbasa/cyclop v1.2.1 // indirect
 	github.com/blakesmith/ar v0.0.0-20190502131153-809d4375e1fb // indirect
 	github.com/blizzy78/varnamelen v0.8.0 // indirect
+<<<<<<< HEAD
+	github.com/bombsimon/wsl/v3 v3.4.0 // indirect
+=======
+>>>>>>> main
 	github.com/bombsimon/wsl/v4 v4.2.1 // indirect
 	github.com/breml/bidichk v0.2.7 // indirect
 	github.com/breml/errchkjson v0.3.6 // indirect
@@ -138,6 +147,10 @@ require (
 	github.com/docker/go-units v0.5.0 // indirect
 	github.com/elliotchance/orderedmap/v2 v2.2.0 // indirect
 	github.com/emirpasic/gods v1.18.1 // indirect
+<<<<<<< HEAD
+	github.com/esimonov/ifshort v1.0.4 // indirect
+=======
+>>>>>>> main
 	github.com/ettle/strcase v0.2.0 // indirect
 	github.com/evanphx/json-patch/v5 v5.6.0 // indirect
 	github.com/fatih/color v1.16.0 // indirect
@@ -226,6 +239,10 @@ require (
 	github.com/karamaru-alpha/copyloopvar v1.0.10 // indirect
 	github.com/kevinburke/ssh_config v1.2.0 // indirect
 	github.com/kisielk/errcheck v1.7.0 // indirect
+<<<<<<< HEAD
+	github.com/kisielk/gotool v1.0.0 // indirect
+=======
+>>>>>>> main
 	github.com/kkHAIKE/contextcheck v1.1.5 // indirect
 	github.com/klauspost/compress v1.17.0 // indirect
 	github.com/klauspost/pgzip v1.2.6 // indirect
@@ -250,6 +267,10 @@ require (
 	github.com/mattn/go-mastodon v0.0.6 // indirect
 	github.com/mattn/go-runewidth v0.0.14 // indirect
 	github.com/matttproud/golang_protobuf_extensions v1.0.4 // indirect
+<<<<<<< HEAD
+	github.com/mbilski/exhaustivestruct v1.2.0 // indirect
+=======
+>>>>>>> main
 	github.com/mgechev/revive v1.3.7 // indirect
 	github.com/mitchellh/copystructure v1.2.0 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
@@ -382,6 +403,11 @@ require (
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	honnef.co/go/tools v0.4.7 // indirect
+<<<<<<< HEAD
+	mvdan.cc/interfacer v0.0.0-20180901003855-c20040233aed // indirect
+	mvdan.cc/lint v0.0.0-20170908181259-adc824a0674b // indirect
+=======
+>>>>>>> main
 	mvdan.cc/unparam v0.0.0-20240104100049-c549a3470d14 // indirect
 	sigs.k8s.io/kind v0.20.0 // indirect
 	sigs.k8s.io/yaml v1.3.0 // indirect