Merge pull request #34 from RoseSecurity/create-list-models-command

feat: add kuzco list models command

Author: RoseSecurity

cmd/list.go

@@ -0,0 +1,17 @@
+package cmd
+
+import (
+	"github.com/spf13/cobra"
+)
+
+// listCmd commands lists models
+var listCmd = &cobra.Command{
+	Use:                "list",
+	Short:              "Lists available Ollama models",
+	Long:               `Queries and lists available Ollama models for Kuzco use`,
+	FParseErrWhitelist: struct{ UnknownFlags bool }{UnknownFlags: false},
+}
+
+func init() {
+	rootCmd.AddCommand(listCmd)
+}

cmd/list_models.go

@@ -0,0 +1,56 @@
+package cmd
+
+import (
+	"fmt"
+	"strings"
+
+	"github.com/charmbracelet/glamour"
+	"github.com/spf13/cobra"
+
+	u "github.com/RoseSecurity/kuzco/pkg/utils"
+)
+
+// ListModelsCmd lists Ollama models
+var listModelsCmd = &cobra.Command{
+	Use:                "models",
+	Short:              "Lists available Ollama models",
+	Long:               `This command lists Ollama models`,
+	Example:            "list models",
+	FParseErrWhitelist: struct{ UnknownFlags bool }{UnknownFlags: false},
+	Run: func(cmd *cobra.Command, args []string) {
+		models, err := u.ListModels(addr)
+		if err != nil {
+			u.LogErrorAndExit(err)
+		}
+
+		// Prepare a markdown list of models
+		var modelList strings.Builder
+		modelList.WriteString("### Available Ollama Models\n\n")
+
+		if len(models) == 0 {
+			modelList.WriteString("No models available.\n")
+		} else {
+			for _, model := range models {
+				modelList.WriteString(fmt.Sprintf("- %s\n", model))
+			}
+		}
+
+		renderer, err := glamour.NewTermRenderer(glamour.WithAutoStyle())
+		if err != nil {
+			u.LogErrorAndExit(fmt.Errorf("error creating markdown renderer: %v", err))
+		}
+
+		rendered, err := renderer.Render(modelList.String())
+		if err != nil {
+			u.LogErrorAndExit(fmt.Errorf("error rendering markdown: %v", err))
+		}
+
+		fmt.Print(rendered)
+	},
+}
+
+func init() {
+	listModelsCmd.Flags().StringVarP(&addr, "address", "a", "http://localhost:11434", "IP Address and port to use for the LLM model (ex: http://localhost:11434)")
+	listCmd.AddCommand(listModelsCmd)
+	listModelsCmd.DisableFlagParsing = false
+}

pkg/utils/list_models.go

@@ -0,0 +1,44 @@
+package utils
+
+import (
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/url"
+)
+
+type Model struct {
+	Name string `json:"name"`
+}
+
+// ListModels lists available models from Ollama
+func ListModels(addr string) ([]string, error) {
+	if _, err := url.Parse(addr); err != nil {
+		return nil, fmt.Errorf("invalid address provided: %v", err)
+	}
+
+	resp, err := http.Get(fmt.Sprintf("%s/api/tags", addr))
+	if err != nil {
+		return nil, fmt.Errorf("error fetching models from Ollama: %v", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("failed to retrieve models: status code %d", resp.StatusCode)
+	}
+
+	var result struct {
+		Models []Model `json:"models"`
+	}
+
+	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
+		return nil, fmt.Errorf("error decoding response: %v", err)
+	}
+
+	var modelNames []string
+	for _, model := range result.Models {
+		modelNames = append(modelNames, model.Name)
+	}
+
+	return modelNames, nil
+}