diff --git a/pkg/cmd/cmd.go b/pkg/cmd/cmd.go index 99c65767..c04a58fc 100644 --- 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 index 0bcba20d..6d7ea77a 100644 --- a/pkg/cmd/ollama/ollama.go +++ b/pkg/cmd/ollama/ollama.go @@ -47,16 +47,16 @@ 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) @@ -64,7 +64,7 @@ func validateModelType(input string, ollamaStore OllamaStore) (bool, error) { 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 index 9428ed27..922b9b14 100644 --- 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 index 961ee2c4..9c7d36d9 100644 --- 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 index d27333e5..40c480b9 100644 --- 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) }