rm multiple models

commands/rm.go

@@ -9,26 +9,27 @@ import (
 
 func newRemoveCmd(desktopClient *desktop.Client) *cobra.Command {
 	c := &cobra.Command{
-		Use:   "rm MODEL",
-		Short: "Remove a model downloaded from Docker Hub",
+		Use:   "rm [MODEL...]",
+		Short: "Remove models downloaded from Docker Hub",
 		Args: func(cmd *cobra.Command, args []string) error {
-			if len(args) != 1 {
+			if len(args) < 1 {
 				return fmt.Errorf(
-					"'docker model rm' requires 1 argument.\n\n" +
-						"Usage:  docker model rm MODEL\n\n" +
+					"'docker model rm' requires at least 1 argument.\n\n" +
+						"Usage:  docker model rm [MODEL...]\n\n" +
 						"See 'docker model rm --help' for more information",
 				)
 			}
 			return nil
 		},
 		RunE: func(cmd *cobra.Command, args []string) error {
-			model := args[0]
-			response, err := desktopClient.Remove(model)
+			response, err := desktopClient.Remove(args)
+			if response != "" {
+				cmd.Println(response)
+			}
 			if err != nil {
 				err = handleClientError(err, "Failed to remove model")
 				return handleNotRunningError(err)
 			}
-			cmd.Println(response)
 			return nil
 		},
 	}

desktop/desktop.go

@@ -121,28 +121,18 @@ func (c *Client) List(jsonFormat, openai bool, model string) (string, error) {
 	}
 	if model != "" {
 		if !strings.Contains(strings.Trim(model, "/"), "/") {
-			// We assume a model name is invalid if it does not contain a "/".
-			return "", fmt.Errorf("invalid model name: %s", model)
+			// Do an extra API call to check if the model parameter isn't a model ID.
+			var err error
+			if model, err = c.modelNameFromID(model); err != nil {
+				return "", fmt.Errorf("invalid model name: %s", model)
+			}
 		}
 		modelsRoute += "/" + model
 	}
 
-	resp, err := c.doRequest(http.MethodGet, modelsRoute, nil)
-	if err != nil {
-		return "", c.handleQueryError(err, modelsRoute)
-	}
-	defer resp.Body.Close()
-
-	if resp.StatusCode != http.StatusOK {
-		if model != "" && resp.StatusCode == http.StatusNotFound {
-			return "", errors.Wrap(ErrNotFound, model)
-		}
-		return "", fmt.Errorf("failed to list models: %s", resp.Status)
-	}
-
-	body, err := io.ReadAll(resp.Body)
+	body, err := c.listRaw(modelsRoute, model)
 	if err != nil {
-		return "", fmt.Errorf("failed to read response body: %w", err)
+		return "", err
 	}
 
 	if openai {
@@ -182,6 +172,48 @@ func (c *Client) List(jsonFormat, openai bool, model string) (string, error) {
 	return prettyPrintModels(modelsJson), nil
 }
 
+func (c *Client) listRaw(route string, model string) ([]byte, error) {
+	resp, err := c.doRequest(http.MethodGet, route, nil)
+	if err != nil {
+		return nil, c.handleQueryError(err, route)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		if model != "" && resp.StatusCode == http.StatusNotFound {
+			return nil, errors.Wrap(ErrNotFound, model)
+		}
+		return nil, fmt.Errorf("failed to list models: %s", resp.Status)
+	}
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response body: %w", err)
+	}
+	return body, nil
+
+}
+
+func (c *Client) modelNameFromID(id string) (string, error) {
+	bodyResponse, err := c.listRaw(inference.ModelsPrefix, "")
+	if err != nil {
+		return "", err
+	}
+
+	var modelsJson []Model
+	if err := json.Unmarshal(bodyResponse, &modelsJson); err != nil {
+		return "", fmt.Errorf("failed to unmarshal response body: %w", err)
+	}
+
+	for _, m := range modelsJson {
+		if m.ID[7:19] == id || strings.TrimPrefix(m.ID, "sha256:") == id || m.ID == id {
+			return m.Tags[0], nil
+		}
+	}
+
+	return "", fmt.Errorf("model with ID %s not found", id)
+}
+
 func (c *Client) Chat(model, prompt string) error {
 	reqBody := OpenAIChatRequest{
 		Model: model,
@@ -250,26 +282,38 @@ func (c *Client) Chat(model, prompt string) error {
 	return nil
 }
 
-func (c *Client) Remove(model string) (string, error) {
-	removePath := inference.ModelsPrefix + "/" + model
-	resp, err := c.doRequest(http.MethodDelete, removePath, nil)
-	if err != nil {
-		return "", c.handleQueryError(err, removePath)
-	}
-	defer resp.Body.Close()
+func (c *Client) Remove(models []string) (string, error) {
+	modelRemoved := ""
+	for _, model := range models {
+		// Check if not a model ID passed as parameter.
+		if !strings.Contains(model, "/") {
+			var err error
+			modelID := model
+			if model, err = c.modelNameFromID(model); err != nil {
+				return modelRemoved, fmt.Errorf("invalid model name: %s", modelID)
+			}
+		}
 
-	if resp.StatusCode != http.StatusOK {
-		var bodyStr string
-		body, err := io.ReadAll(resp.Body)
+		removePath := inference.ModelsPrefix + "/" + model
+		resp, err := c.doRequest(http.MethodDelete, removePath, nil)
 		if err != nil {
-			bodyStr = fmt.Sprintf("(failed to read response body: %v)", err)
-		} else {
-			bodyStr = string(body)
+			return modelRemoved, c.handleQueryError(err, removePath)
 		}
-		return "", fmt.Errorf("removing %s failed with status %s: %s", model, resp.Status, bodyStr)
+		defer resp.Body.Close()
+
+		if resp.StatusCode != http.StatusOK {
+			var bodyStr string
+			body, err := io.ReadAll(resp.Body)
+			if err != nil {
+				bodyStr = fmt.Sprintf("(failed to read response body: %v)", err)
+			} else {
+				bodyStr = string(body)
+			}
+			return modelRemoved, fmt.Errorf("removing %s failed with status %s: %s", model, resp.Status, bodyStr)
+		}
+		modelRemoved += fmt.Sprintf("Model %s removed successfully\n", model)
 	}
-
-	return fmt.Sprintf("Model %s removed successfully", model), nil
+	return modelRemoved, nil
 }
 
 func URL(path string) string {