From 17726f032dda4e606c923667bde6db263ef08e28 Mon Sep 17 00:00:00 2001 From: Sam McLeod Date: Fri, 31 May 2024 23:47:25 +1000 Subject: [PATCH] minor updates --- app_model.go | 53 +++++++++++++++++++++++++------- helpers.go | 22 -------------- item_delegate.go | 49 ++++++++++++------------------ keymap.go | 2 ++ main.go | 1 + operations.go | 70 +++++++++++++++++++++--------------------- styles.go | 79 +++++++++++++++++++++++++++++++++++++++++++++++- 7 files changed, 177 insertions(+), 99 deletions(-) diff --git a/app_model.go b/app_model.go index 77793b3..bcbd72b 100644 --- a/app_model.go +++ b/app_model.go @@ -4,6 +4,7 @@ import ( "fmt" "gollama/logging" "sort" + "strings" "github.com/charmbracelet/bubbles/key" "github.com/charmbracelet/bubbles/list" @@ -18,7 +19,7 @@ func (m *AppModel) Init() tea.Cmd { func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { switch msg := msg.(type) { case tea.KeyMsg: - logging.DebugLogger.Printf("AppModel received key: %s\n", msg.String()) // Add this line + logging.DebugLogger.Printf("AppModel received key: %s\n", msg.String()) switch { case key.Matches(msg, m.keys.Space): if item, ok := m.list.SelectedItem().(Model); ok { @@ -94,23 +95,43 @@ func (m *AppModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } case key.Matches(msg, m.keys.LinkModel): if item, ok := m.list.SelectedItem().(Model); ok { - err := linkModel(item.Name, m.lmStudioModelsDir, m.noCleanup) + message, err := linkModel(item.Name, m.lmStudioModelsDir, m.noCleanup) if err != nil { - lipgloss.NewStyle().Foreground(lipgloss.Color("red")).Render(fmt.Sprintf("Error linking model: %v", err)) + m.message = fmt.Sprintf("Error linking model: %v", err) + } else if message != "" { + break + } else { + m.message = fmt.Sprintf("Model %s linked successfully", item.Name) } } - m.refreshList() + return m.clearScreen(), tea.ClearScreen case key.Matches(msg, m.keys.LinkAllModels): - // make sure we start with a clean terminal - + var messages []string for _, model := range m.models { - err := linkModel(model.Name, m.lmStudioModelsDir, m.noCleanup) + message, err := linkModel(model.Name, m.lmStudioModelsDir, m.noCleanup) + // if the message is empty, don't add it to the list if err != nil { - lipgloss.NewStyle().Foreground(lipgloss.Color("red")).Render(fmt.Sprintf("Error linking model: %v", err)) + messages = append(messages, fmt.Sprintf("Error linking model %s: %v", model.Name, err)) + } else if message != "" { + continue + } else { + messages = append(messages, message) } } - m.refreshList() - m.View( /* force refresh */ ) + // remove any empty messages or duplicates + for i := 0; i < len(messages); i++ { + for j := i + 1; j < len(messages); j++ { + if messages[i] == messages[j] { + messages = append(messages[:j], messages[j+1:]...) + j-- + } + } + } + messages = append(messages, "Linking complete") + m.message = strings.Join(messages, "\n") + return m.clearScreen(), tea.ClearScreen + case key.Matches(msg, m.keys.ClearScreen): + return m.clearScreen(), tea.ClearScreen } case tea.WindowSizeMsg: m.width = msg.Width @@ -143,7 +164,11 @@ func (m *AppModel) View() string { idWidth, "ID", ), ) - return header + "\n" + m.list.View() + message := "" + if m.message != "" { + message = lipgloss.NewStyle().Foreground(lipgloss.Color("green")).Render(m.message) + "\n" + } + return message + header + "\n" + m.list.View() } func (m *AppModel) refreshList() { @@ -153,3 +178,9 @@ func (m *AppModel) refreshList() { } m.list.SetItems(items) } + +func (m *AppModel) clearScreen() tea.Model { + m.list.ResetFilter() + m.refreshList() + return m +} diff --git a/helpers.go b/helpers.go index e4021ff..e10e7eb 100644 --- a/helpers.go +++ b/helpers.go @@ -2,7 +2,6 @@ package main import ( "gollama/logging" - "sort" "github.com/charmbracelet/lipgloss" "github.com/ollama/ollama/api" @@ -111,24 +110,3 @@ func wrapText(text string, width int) string { wrapped += text return wrapped } - -func sortModels(models *[]Model, order string) { - switch order { - case "name": - sort.Slice(*models, func(i, j int) bool { - return (*models)[i].Name < (*models)[j].Name - }) - case "size": - sort.Slice(*models, func(i, j int) bool { - return (*models)[i].Size > (*models)[j].Size - }) - case "modified": - sort.Slice(*models, func(i, j int) bool { - return (*models)[i].Modified.After((*models)[j].Modified) - }) - case "family": - sort.Slice(*models, func(i, j int) bool { - return (*models)[i].Family < (*models)[j].Family - }) - } -} diff --git a/item_delegate.go b/item_delegate.go index 638355e..7dda002 100644 --- a/item_delegate.go +++ b/item_delegate.go @@ -4,7 +4,6 @@ import ( "fmt" "io" - "gollama/colors" "gollama/logging" "github.com/charmbracelet/bubbles/list" @@ -49,41 +48,31 @@ func (d itemDelegate) Render(w io.Writer, m list.Model, index int, item list.Ite if !ok { return } - var nameStyle, idStyle, sizeStyle, quantStyle, modifiedStyle, familyStyle lipgloss.Style - familyColor, exists := colors.FamilyColors[model.Family] - if !exists { - // pick a random colour between 10 and 200 - familyColor = lipgloss.Color(fmt.Sprintf("%d", 10+index%190)) - } - // set the quantization level colour - quantColor := colors.QuantColor(model.QuantizationLevel) - sizeColor := colors.SizeColor(model.Size) + nameStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("254")) + idStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("254")).Faint(true) + sizeStyle := lipgloss.NewStyle().Foreground(sizeColor(model.Size)).Faint(true) + familyStyle := lipgloss.NewStyle().Foreground(familyColor(model.Family, index)) + quantStyle := lipgloss.NewStyle().Foreground(quantColor(model.QuantizationLevel)) + modifiedStyle := lipgloss.NewStyle().Foreground(lipgloss.Color("254")).Faint(true) if index == m.Index() { - // Highlight the selected item - nameStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("205")).Bold(true).BorderLeft(true).BorderStyle(lipgloss.OuterHalfBlockBorder()).BorderLeftBackground(lipgloss.Color("200")) - sizeStyle = lipgloss.NewStyle().Foreground(sizeColor).Bold(true).BorderLeft(true).BorderLeftBackground(lipgloss.Color("200")) - quantStyle = lipgloss.NewStyle().Foreground(quantColor).Bold(true).BorderLeft(true).BorderLeftBackground(lipgloss.Color("200")) - familyStyle = lipgloss.NewStyle().Foreground(familyColor).Bold(true).BorderLeft(true).BorderLeftBackground(lipgloss.Color("200")) - modifiedStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("115")).BorderLeft(true) - idStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("225")).BorderLeft(true) - } else { - nameStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("254")) - idStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("254")).Faint(true) - sizeStyle = lipgloss.NewStyle().Foreground(sizeColor).Faint(true) - familyStyle = lipgloss.NewStyle().Foreground(familyColor) - quantStyle = lipgloss.NewStyle().Foreground(quantColor) - modifiedStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("254")).Faint(true) + nameStyle = nameStyle.Bold(true).BorderLeft(true).BorderStyle(lipgloss.OuterHalfBlockBorder()).BorderLeftBackground(lipgloss.Color("200")) + sizeStyle = sizeStyle.Bold(true).BorderLeft(true).BorderLeftBackground(lipgloss.Color("200")) + quantStyle = quantStyle.Bold(true).BorderLeft(true).BorderLeftBackground(lipgloss.Color("200")) + familyStyle = familyStyle.Bold(true).BorderLeft(true).BorderLeftBackground(lipgloss.Color("200")) + modifiedStyle = modifiedStyle.Foreground(lipgloss.Color("115")).BorderLeft(true) + idStyle = idStyle.Foreground(lipgloss.Color("225")).BorderLeft(true) } if model.Selected { - nameStyle = nameStyle.Background(lipgloss.Color("236")).Bold(true).Italic(true) - idStyle = idStyle.Background(lipgloss.Color("236")).Bold(true).Italic(true) - sizeStyle = sizeStyle.Background(lipgloss.Color("236")).Bold(true).Italic(true) - familyStyle = familyStyle.Background(lipgloss.Color("236")).Bold(true).Italic(true) - quantStyle = quantStyle.Background(lipgloss.Color("236")).Bold(true).Italic(true) - modifiedStyle = modifiedStyle.Background(lipgloss.Color("236")).Bold(true).Italic(true) + selectedStyle := lipgloss.NewStyle().Background(lipgloss.Color("236")).Bold(true).Italic(true) + nameStyle = nameStyle.Inherit(selectedStyle) + idStyle = idStyle.Inherit(selectedStyle) + sizeStyle = sizeStyle.Inherit(selectedStyle) + familyStyle = familyStyle.Inherit(selectedStyle) + quantStyle = quantStyle.Inherit(selectedStyle) + modifiedStyle = modifiedStyle.Inherit(selectedStyle) } nameWidth, sizeWidth, quantWidth, modifiedWidth, idWidth, familyWidth := calculateColumnWidths(m.Width()) diff --git a/keymap.go b/keymap.go index 8a6d1d1..455844b 100644 --- a/keymap.go +++ b/keymap.go @@ -15,6 +15,7 @@ type KeyMap struct { ConfirmNo key.Binding LinkModel key.Binding LinkAllModels key.Binding + ClearScreen key.Binding SortOrder string } @@ -32,6 +33,7 @@ func NewKeyMap() *KeyMap { ConfirmNo: key.NewBinding(key.WithKeys("n"), key.WithHelp("n", "cancel deletion")), LinkModel: key.NewBinding(key.WithKeys("l"), key.WithHelp("l", "link model to LM Studio")), LinkAllModels: key.NewBinding(key.WithKeys("L"), key.WithHelp("L", "link all models to LM Studio")), + ClearScreen: key.NewBinding(key.WithKeys("c"), key.WithHelp("c", "clear screen")), } } diff --git a/main.go b/main.go index 05bc3c1..35a11d8 100644 --- a/main.go +++ b/main.go @@ -31,6 +31,7 @@ type AppModel struct { lmStudioModelsDir string noCleanup bool cfg *config.Config + message string } func main() { diff --git a/operations.go b/operations.go index 1ef27e5..2b7c241 100644 --- a/operations.go +++ b/operations.go @@ -10,7 +10,6 @@ import ( "gollama/logging" - "github.com/charmbracelet/lipgloss" "github.com/ollama/ollama/api" "golang.org/x/term" ) @@ -69,11 +68,10 @@ func deleteModel(client *api.Client, name string) error { return nil } -func linkModel(modelName, lmStudioModelsDir string, noCleanup bool) error { +func linkModel(modelName, lmStudioModelsDir string, noCleanup bool) (string, error) { modelPath, err := getModelPath(modelName) if err != nil { - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("205")).Render(fmt.Sprintf("Error getting model path for %s: %v", modelName, err))) - return err + return "", fmt.Errorf("error getting model path for %s: %v", modelName, err) } parts := strings.Split(modelName, ":") @@ -85,29 +83,26 @@ func linkModel(modelName, lmStudioModelsDir string, noCleanup bool) error { lmStudioModelName := strings.ReplaceAll(strings.ReplaceAll(modelName, ":", "-"), "_", "-") lmStudioModelDir := filepath.Join(lmStudioModelsDir, author, lmStudioModelName+"-GGUF") - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("cyan")).Render(fmt.Sprintf("Model: %s", modelName))) - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("green")).Render(fmt.Sprintf("Path: %s", modelPath))) - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("green")).Render(fmt.Sprintf("LM Studio model directory: %s", lmStudioModelDir))) - // Check if the model path is a valid file fileInfo, err := os.Stat(modelPath) if err != nil || fileInfo.IsDir() { - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("red")).Render(fmt.Sprintf("Invalid model path for %s: %s", modelName, modelPath))) - return fmt.Errorf("invalid model path for %s: %s", modelName, modelPath) + return "", fmt.Errorf("invalid model path for %s: %s", modelName, modelPath) } // Check if the symlink already exists and is valid lmStudioModelPath := filepath.Join(lmStudioModelDir, filepath.Base(lmStudioModelName)+".gguf") if _, err := os.Lstat(lmStudioModelPath); err == nil { if isValidSymlink(lmStudioModelPath, modelPath) { - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("grey")).Render(fmt.Sprintf("Found model %s, already linked, skipping...", modelName))) - return nil + message := "Model %s is already symlinked to %s" + logging.InfoLogger.Printf(message+"\n", modelName, lmStudioModelPath) + return "", nil } // Remove the invalid symlink err = os.Remove(lmStudioModelPath) if err != nil { - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("red")).Render(fmt.Sprintf("Failed to remove invalid symlink %s: %v", lmStudioModelPath, err))) - return err + message := "failed to remove invalid symlink %s: %v" + logging.ErrorLogger.Printf(message+"\n", lmStudioModelPath, err) + return "", fmt.Errorf(message, lmStudioModelPath, err) } } @@ -130,39 +125,41 @@ func linkModel(modelName, lmStudioModelsDir string, noCleanup bool) error { return nil }) if err != nil { - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("red")).Render(fmt.Sprintf("Error checking for duplicated symlinks: %v", err))) - return err + message := "error walking LM Studio models directory: %v" + logging.ErrorLogger.Printf(message+"\n", err) + return "", fmt.Errorf(message, err) } if existingSymlinkPath != "" { // Remove the duplicated model directory err = os.RemoveAll(lmStudioModelDir) if err != nil { - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("red")).Render(fmt.Sprintf("Failed to remove duplicated model directory %s: %v", lmStudioModelDir, err))) - return err + message := "failed to remove duplicated model directory %s: %v" + logging.ErrorLogger.Printf(message+"\n", lmStudioModelDir, err) + return "", fmt.Errorf(message, lmStudioModelDir, err) } - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("yellow")).Render(fmt.Sprintf("Removed duplicated model directory %s", lmStudioModelDir))) - return nil + return fmt.Sprintf("Removed duplicated model directory %s", lmStudioModelDir), nil } // Create the symlink err = os.MkdirAll(lmStudioModelDir, os.ModePerm) if err != nil { - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("red")).Render(fmt.Sprintf("Failed to create directory %s: %v", lmStudioModelDir, err))) - return err + message := "failed to create directory %s: %v" + logging.ErrorLogger.Printf(message+"\n", lmStudioModelDir, err) + return "", fmt.Errorf(message, lmStudioModelDir, err) } err = os.Symlink(modelPath, lmStudioModelPath) if err != nil { - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("red")).Render(fmt.Sprintf("Failed to symlink %s: %v", modelName, err))) - return err + message := "failed to symlink %s: %v" + logging.ErrorLogger.Printf(message+"\n", modelName, err) + return "", fmt.Errorf(message, modelName, err) } - fmt.Println(lipgloss.NewStyle().Foreground(lipgloss.Color("green")).Render(fmt.Sprintf("Symlinked %s to %s", modelName, lmStudioModelPath))) - if !noCleanup { cleanBrokenSymlinks(lmStudioModelsDir) } - - return nil + message := "Symlinked %s to %s" + logging.InfoLogger.Printf(message+"\n", modelName, lmStudioModelPath) + return "", nil } func getModelPath(modelName string) (string, error) { @@ -177,7 +174,9 @@ func getModelPath(modelName string) (string, error) { return strings.TrimSpace(line[5:]), nil } } - return "", fmt.Errorf("model path not found for %s", modelName) + message := "failed to get model path for %s: no 'FROM' line in output" + logging.ErrorLogger.Printf(message+"\n", modelName) + return "", fmt.Errorf(message, modelName) } func cleanBrokenSymlinks(lmStudioModelsDir string) { @@ -191,7 +190,7 @@ func cleanBrokenSymlinks(lmStudioModelsDir string) { return err } if len(files) == 0 { - fmt.Printf("%sRemoving empty directory: %s%s\n", lipgloss.Color("yellow"), path, lipgloss.Color("reset")) + logging.InfoLogger.Printf("Removing empty directory: %s\n", path) err = os.Remove(path) if err != nil { return err @@ -203,7 +202,7 @@ func cleanBrokenSymlinks(lmStudioModelsDir string) { return err } if !isValidSymlink(path, linkPath) { - fmt.Printf("%sRemoving invalid symlink: %s%s\n", lipgloss.Color("yellow"), path, lipgloss.Color("reset")) + logging.InfoLogger.Printf("Removing invalid symlink: %s\n", path) err = os.Remove(path) if err != nil { return err @@ -213,7 +212,8 @@ func cleanBrokenSymlinks(lmStudioModelsDir string) { return nil }) if err != nil { - fmt.Printf("%sError walking LM Studio models directory: %v%s\n", lipgloss.Color("red"), err, lipgloss.Color("reset")) + logging.ErrorLogger.Printf("Error walking LM Studio models directory: %v\n", err) + return } } @@ -251,7 +251,7 @@ func cleanupSymlinkedModels(lmStudioModelsDir string) { return err } if len(files) == 0 { - fmt.Printf("%sRemoving empty directory: %s%s\n", lipgloss.Color("yellow"), path, lipgloss.Color("reset")) + logging.InfoLogger.Printf("Removing empty directory: %s\n", path) err = os.Remove(path) if err != nil { return err @@ -259,7 +259,7 @@ func cleanupSymlinkedModels(lmStudioModelsDir string) { hasEmptyDir = true } } else if info.Mode()&os.ModeSymlink != 0 { - fmt.Printf("%sRemoving symlinked model: %s%s\n", lipgloss.Color("yellow"), path, lipgloss.Color("reset")) + logging.InfoLogger.Printf("Removing symlinked model: %s\n", path) err = os.Remove(path) if err != nil { return err @@ -268,7 +268,7 @@ func cleanupSymlinkedModels(lmStudioModelsDir string) { return nil }) if err != nil { - fmt.Printf("%sError walking LM Studio models directory: %v%s\n", lipgloss.Color("red"), err, lipgloss.Color("reset")) + logging.ErrorLogger.Printf("Error walking LM Studio models directory: %v\n", err) return } if !hasEmptyDir { diff --git a/styles.go b/styles.go index cec6999..17314ad 100644 --- a/styles.go +++ b/styles.go @@ -1,6 +1,11 @@ package main -import "github.com/charmbracelet/lipgloss" +import ( + "fmt" + "math" + + "github.com/charmbracelet/lipgloss" +) var ( headerStyle = lipgloss.NewStyle(). @@ -19,4 +24,76 @@ var ( Foreground(lipgloss.Color("#000000")). Background(lipgloss.Color("#FF00FF")). Bold(true) + + errorStyle = lipgloss.NewStyle().Foreground(lipgloss.Color("red")) + + // Define brighter colors for different model families + familyColors = map[string]lipgloss.Color{ + "alpaca": lipgloss.Color("#7B6BA5"), + "bert": lipgloss.Color("#ED5234"), + "command-r": lipgloss.Color("#C575B7"), + "gemma": lipgloss.Color("#FFC060"), + "guanaco": lipgloss.Color("#F8DAD9"), + "llama": lipgloss.Color("#FF7F6F"), + "nomic-bert": lipgloss.Color("#E66086"), + "phi": lipgloss.Color("#A2B8E0"), + "phi2": lipgloss.Color("#A9B1F9"), + "phi3": lipgloss.Color("#A2B8E1"), + "qwen": lipgloss.Color("#FFA07A"), + "qwen2": lipgloss.Color("#A91511"), + "starcoder": lipgloss.Color("#55C8BC"), + "starcoder2": lipgloss.Color("#10C787"), + "vicuna": lipgloss.Color("#98C05B"), + "granite": lipgloss.Color("#A9A9A9"), + } + + // Define color gradients + orangeGradient = []string{ + "#FFD700", "#FF0000", "#FF1269", "#FF2285", "#FF409D", "#FF60B5", + "#FF80CD", "#FF90D9", "#FFA0E5", "#FFB0F1", "#FFBFF1", "#FFFEF1", + } + + purpleGradient = []string{ + "#FF0000", "#FFFF00", "#00FF00", "#00FF7F", "#00FFBF", "#00BFFF", + "#007FFF", "#003FFF", "#0000FF", + } ) + +func quantColor(quant string) lipgloss.Color { + quantMap := map[string]int{ + "IQ1_XXS": 0, "IQ1_XS": 0, "IQ1_S": 0, "IQ1_NL": 0, + "Q2_K": 1, "Q2_K_S": 1, "Q2_K_M": 1, "Q2_K_L": 1, + "Q3_0": 2, "IQ2_XXS": 2, "Q3_K_S": 2, + "IQ2_XS": 3, "IQ2_S": 3, "IQ2_NL": 3, + "Q3_K_M": 4, "Q3_K_L": 4, + "Q4_0": 5, "IQ3_XXS": 5, "IQ3_XS": 5, "IQ3_NL": 5, "IQ3_S": 6, + "Q4_K_S": 6, "Q4_1": 6, "IQ4_XXS": 6, "Q4_K_M": 6, + "IQ4_XS": 7, "IQ4_S": 7, "IQ4_NL": 7, "Q4_K_L": 7, + "Q5_K_S": 8, "Q5_K_M": 8, "Q5_1": 8, "Q5_2": 8, "Q5_K_L": 8, + "Q6_0": 9, "Q6_1": 9, "Q6_K": 9, + "Q8": 11, "Q8_0": 11, "Q8_1": 11, "Q8_K": 11, + "FP16": 12, + } + + index, exists := quantMap[quant] + if !exists { + index = 0 // Default to lightest if unknown quant + } + return lipgloss.Color(orangeGradient[index]) +} + +func sizeColor(size float64) lipgloss.Color { + index := int(math.Log10(size+1) * 2.5) + if index >= len(purpleGradient) { + index = len(purpleGradient) - 1 + } + return lipgloss.Color(purpleGradient[index]) +} + +func familyColor(family string, index int) lipgloss.Color { + color, exists := familyColors[family] + if !exists { + color = lipgloss.Color(fmt.Sprintf("%d", 10+index%190)) // Random color + } + return color +}