reworked the ollamahttp client

brevdev · May 17, 2024 · daf6a61 · daf6a61
1 parent d264af3
commit daf6a61
Show file tree

Hide file tree

Showing 5 changed files with 35 additions and 40 deletions.
diff --git a/pkg/cmd/cmd.go b/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)
 

diff --git a/pkg/cmd/ollama/ollama.go b/pkg/cmd/ollama/ollama.go
@@ -47,24 +47,24 @@ func validateModelType(input string, ollamaStore OllamaStore) (bool, error) {
 	var tag string
 
 	split := strings.Split(input, ":")
-	if len(split) > 2 {
-		return false, fmt.Errorf("invalid model type: %s", input)
-	} else if len(split) == 2 {
-		model = strings.Split(input, ":")[0]
-		tag = strings.Split(input, ":")[1]
-	} else {
+	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)
 	}
-	// use the validateollamamodel function to check if the model is valid
 	valid, err := ollamaStore.ValidateOllamaModel(model, tag)
 	if err != nil {
 		return false, fmt.Errorf("error validating model: %s", err)
 	}
 	return valid, nil
 }
 
-func oNewCmdOllama(t *terminal.Terminal, ollamaStore OllamaStore) *cobra.Command {
+func NewCmdOllama(t *terminal.Terminal, ollamaStore OllamaStore) *cobra.Command {
 	var model string
 
 	cmd := &cobra.Command{

diff --git a/pkg/store/http.go b/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}
 }
 
@@ -217,25 +230,3 @@ func IsNetworkErrorWithStatus(err error, statusCodes []int) bool {
 		return false
 	}
 }
-
-type OllamaHTTPClient struct {
-	restyClient *resty.Client
-}
-
-type OllamaHTTPStore struct {
-	ollamaHTTPClient *OllamaHTTPClient
-}
-
-func NewOllamaHTTPClient(ollamaAPIURL string) *OllamaHTTPClient {
-	restyClient := resty.New().SetBaseURL(ollamaAPIURL)
-
-	return &OllamaHTTPClient{
-		restyClient: restyClient,
-	}
-}
-
-func WithOllamaHTTPClient(ollamaHTTPClient *OllamaHTTPClient) *OllamaHTTPStore {
-	return &OllamaHTTPStore{
-		ollamaHTTPClient: ollamaHTTPClient,
-	}
-}
diff --git a/pkg/store/http_test.go b/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
 }
 

diff --git a/pkg/store/workspace.go b/pkg/store/workspace.go
@@ -693,16 +693,18 @@ var (
 	ollamaModelPath        = fmt.Sprintf(ollamaModelPathPattern, fmt.Sprintf("{%s}", modelNameParamName), fmt.Sprintf("{%s}", tagNameParamName))
 )
 
-func (o OllamaHTTPStore) ValidateOllamaModel(req OllamaModelRequest) (bool, error) {
-	res, err := o.ollamaHTTPClient.restyClient.R().
+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, req.Model).
-		SetPathParam(tagNameParamName, req.Tag).
+		SetPathParam(modelNameParamName, model).
+		SetPathParam(tagNameParamName, tag).
 		Get(ollamaModelPath)
 	if err != nil {
 		return false, breverrors.WrapAndTrace(err)
 	}
-	if res.StatusCode() == 200 {
+	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)
 		}