Skip to content

Commit

Permalink
minor updates
Browse files Browse the repository at this point in the history
  • Loading branch information
@sammcj
sammcj committed May 31, 2024
1 parent ed7d567 commit 17726f0
Show file tree
Hide file tree
Showing 7 changed files with 177 additions and 99 deletions.
53 changes: 42 additions & 11 deletions app_model.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ import (
"fmt"
"gollama/logging"
"sort"
"strings"

"github.com/charmbracelet/bubbles/key"
"github.com/charmbracelet/bubbles/list"
Expand All @@ -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 {
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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() {
Expand All @@ -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
}
22 changes: 0 additions & 22 deletions helpers.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@ package main

import (
"gollama/logging"
"sort"

"github.com/charmbracelet/lipgloss"
"github.com/ollama/ollama/api"
Expand Down Expand Up @@ -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
})
}
}
49 changes: 19 additions & 30 deletions item_delegate.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@ import (
"fmt"
"io"

"gollama/colors"
"gollama/logging"

"github.com/charmbracelet/bubbles/list"
Expand Down Expand Up @@ -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())
Expand Down
2 changes: 2 additions & 0 deletions keymap.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ type KeyMap struct {
ConfirmNo key.Binding
LinkModel key.Binding
LinkAllModels key.Binding
ClearScreen key.Binding
SortOrder string
}

Expand All @@ -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")),
}
}

Expand Down
1 change: 1 addition & 0 deletions main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@ type AppModel struct {
lmStudioModelsDir string
noCleanup bool
cfg *config.Config
message string
}

func main() {
Expand Down
70 changes: 35 additions & 35 deletions operations.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ import (

"gollama/logging"

"github.com/charmbracelet/lipgloss"
"github.com/ollama/ollama/api"
"golang.org/x/term"
)
Expand Down Expand Up @@ -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, ":")
Expand All @@ -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)
}
}

Expand All @@ -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) {
Expand All @@ -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) {
Expand All @@ -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
Expand All @@ -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
Expand All @@ -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
}
}

Expand Down Expand Up @@ -251,15 +251,15 @@ 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
}
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
Expand All @@ -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 {
Expand Down
Loading

0 comments on commit 17726f0

Please sign in to comment.