Add and remove models from collections (#396)

Author: marccampbell

cli/cmd/collection_addmodel.go

@@ -0,0 +1,38 @@
+package cmd
+
+import (
+	"fmt"
+
+	"github.com/pkg/errors"
+	"github.com/spf13/cobra"
+)
+
+func (r *runners) InitCollectionAddModel(parent *cobra.Command) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:          "add-model",
+		Short:        "add a model to a collection",
+		Long:         `add a model to a collection`,
+		RunE:         r.addModelToCollection,
+		SilenceUsage: true,
+	}
+	parent.AddCommand(cmd)
+
+	cmd.Flags().StringVar(&r.args.modelCollectionAddModelName, "model-name", "", "The name of the model")
+	cmd.MarkFlagRequired("model-id")
+	cmd.Flags().StringVar(&r.args.modelCollectionAddModelCollectionID, "collection-id", "", "The ID of the collection")
+	cmd.MarkFlagRequired("collection-id")
+
+	cmd.Flags().StringVar(&r.outputFormat, "output", "table", "The output format to use. One of: json|table (default: table)")
+
+	return cmd
+}
+
+func (r *runners) addModelToCollection(_ *cobra.Command, args []string) error {
+	err := r.kotsAPI.UpdateModelsInCollection(r.args.modelCollectionAddModelCollectionID, []string{r.args.modelCollectionAddModelName}, nil)
+	if err != nil {
+		return errors.Wrap(err, "list model collections")
+	}
+
+	fmt.Fprintf(r.w, "Model %s has been added to collection %s\n", r.args.modelCollectionAddModelName, r.args.modelCollectionAddModelCollectionID)
+	return nil
+}

cli/cmd/collection_removemodel.go

@@ -0,0 +1,38 @@
+package cmd
+
+import (
+	"fmt"
+
+	"github.com/pkg/errors"
+	"github.com/spf13/cobra"
+)
+
+func (r *runners) InitCollectionRemoveModel(parent *cobra.Command) *cobra.Command {
+	cmd := &cobra.Command{
+		Use:          "remove-model",
+		Short:        "remove a model from a collection",
+		Long:         `remove a model from a collection`,
+		RunE:         r.removeModelFromCollection,
+		SilenceUsage: true,
+	}
+	parent.AddCommand(cmd)
+
+	cmd.Flags().StringVar(&r.args.modelCollectionRmModelName, "model-name", "", "The name of the model")
+	cmd.MarkFlagRequired("model-id")
+	cmd.Flags().StringVar(&r.args.modelCollectionRmModelCollectionID, "collection-id", "", "The ID of the collection")
+	cmd.MarkFlagRequired("collection-id")
+
+	cmd.Flags().StringVar(&r.outputFormat, "output", "table", "The output format to use. One of: json|table (default: table)")
+
+	return cmd
+}
+
+func (r *runners) removeModelFromCollection(_ *cobra.Command, args []string) error {
+	err := r.kotsAPI.UpdateModelsInCollection(r.args.modelCollectionRmModelCollectionID, nil, []string{r.args.modelCollectionRmModelName})
+	if err != nil {
+		return errors.Wrap(err, "list model collections")
+	}
+
+	fmt.Fprintf(r.w, "Model %s has been removed from collection %s\n", r.args.modelCollectionRmModelName, r.args.modelCollectionRmModelCollectionID)
+	return nil
+}

cli/cmd/model_ls.go

@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"github.com/replicatedhq/replicated/cli/print"
 	"github.com/spf13/cobra"
 )
 
@@ -19,5 +20,10 @@ func (r *runners) InitModelList(parent *cobra.Command) *cobra.Command {
 }
 
 func (r *runners) listModels(_ *cobra.Command, args []string) error {
-	return nil
+	models, err := r.kotsAPI.ListModels()
+	if err != nil {
+		return err
+	}
+
+	return print.Models(r.outputFormat, r.w, models)
 }

cli/cmd/root.go

@@ -220,6 +220,8 @@ func Execute(rootCmd *cobra.Command, stdin io.Reader, stdout io.Writer, stderr i
 	runCmds.InitCollectionList(collectionCmd)
 	runCmds.InitCollectionCreate(collectionCmd)
 	runCmds.InitCollectionRemove(collectionCmd)
+	runCmds.InitCollectionAddModel(collectionCmd)
+	runCmds.InitCollectionRemoveModel(collectionCmd)
 
 	clusterCmd := runCmds.InitClusterCommand(runCmds.rootCmd)
 	runCmds.InitClusterCreate(clusterCmd)

cli/cmd/runner.go

@@ -242,8 +242,12 @@ type runnerArgs struct {
 	removeClusterNames  []string
 	removeClusterDryRun bool
 
-	modelCollectionCreateName string
-	modelPushName             string
+	modelCollectionCreateName           string
+	modelPushName                       string
+	modelCollectionAddModelName         string
+	modelCollectionAddModelCollectionID string
+	modelCollectionRmModelName          string
+	modelCollectionRmModelCollectionID  string
 
 	lsAppVersion                            string
 	lsVersionsClusterKubernetesDistribution string

cli/print/ai-models.go

@@ -30,3 +30,25 @@ func ModelCollections(outputFormat string, w *tabwriter.Writer, collections []ty
 
 	return w.Flush()
 }
+
+var modelsTmplSrc = `NAME
+{{ range . -}}
+{{ .Name }}
+{{ end }}`
+
+var modelsTmpl = template.Must(template.New("registries").Funcs(funcs).Parse(modelsTmplSrc))
+
+func Models(outputFormat string, w *tabwriter.Writer, models []types.Model) error {
+	if outputFormat == "table" {
+		if err := modelsTmpl.Execute(w, models); err != nil {
+			return err
+		}
+	} else if outputFormat == "json" {
+		mAsByte, _ := json.MarshalIndent(models, "", "  ")
+		if _, err := fmt.Fprintln(w, string(mAsByte)); err != nil {
+			return err
+		}
+	}
+
+	return w.Flush()
+}

pkg/kotsclient/collection_update_models.go

@@ -0,0 +1,27 @@
+package kotsclient
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/pkg/errors"
+)
+
+type updateModelsInCollectionRequest struct {
+	AddModels    []string `json:"add_models"`
+	RemoveModels []string `json:"remove_models"`
+}
+
+func (c *VendorV3Client) UpdateModelsInCollection(collectionID string, modelsToAdd []string, modelsToRemove []string) error {
+	var reqBody = updateModelsInCollectionRequest{
+		AddModels:    modelsToAdd,
+		RemoveModels: modelsToRemove,
+	}
+
+	err := c.DoJSON("PATCH", fmt.Sprintf("/v3/models/collection/%s/models", collectionID), http.StatusOK, reqBody, nil)
+	if err != nil {
+		return errors.Wrap(err, "update models in collection")
+	}
+
+	return nil
+}

pkg/kotsclient/model_list.go

@@ -0,0 +1,26 @@
+package kotsclient
+
+import (
+	"net/http"
+
+	"github.com/pkg/errors"
+	"github.com/replicatedhq/replicated/pkg/types"
+)
+
+type modelsResponse struct {
+	Models []types.Model `json:"models"`
+}
+
+func (c *VendorV3Client) ListModels() ([]types.Model, error) {
+	var response = modelsResponse{}
+
+	err := c.DoJSON("GET", "/v3/models", http.StatusOK, nil, &response)
+	if err != nil {
+		return nil, errors.Wrap(err, "list models")
+	}
+
+	results := make([]types.Model, 0)
+	results = append(results, response.Models...)
+
+	return results, nil
+}

pkg/types/ai-models.go

@@ -8,3 +8,9 @@ type ModelCollection struct {
 	CreatedAt time.Time `json:"createdAt"`
 	UpdatedAt time.Time `json:"updatedAt"`
 }
+
+type Model struct {
+	Name      string    `json:"name"`
+	CreatedAt time.Time `json:"createdAt"`
+	UpdatedAt time.Time `json:"updatedAt"`
+}