feat: ctrl+k to pull a model while keeping any changes (#178)

* ctrl+k to pull a model while keeping any changes

* ctrl+k to pull a model while keeping any changes

* ctrl+k to pull a model while keeping any changes

* fix: #175

Author: sammcj

README.md

@@ -116,6 +116,7 @@ echo "alias g=gollama" >> ~/.zshrc
 - `c`: Copy model
 - `U`: Unload all models
 - `p`: Pull an existing model
+- `ctrl+k`: Pull model & preserve user configuration
 - `ctrl+p`: Pull (get) new model
 - `P`: Push model
 - `n`: Sort by name

app_model.go

@@ -260,6 +260,8 @@ func (m *AppModel) handleKeyMsg(msg tea.KeyMsg) (tea.Model, tea.Cmd) {
 		return m.handlePushModelKey()
 	case key.Matches(msg, m.keys.PullModel):
 		return m.handlePullModelKey()
+	case key.Matches(msg, m.keys.PullKeepConfig):
+		return m.handlePullKeepConfigKey()
 	case key.Matches(msg, m.keys.RenameModel):
 		return m.handleRenameModelKey()
 	case key.Matches(msg, m.keys.PullNewModel):
@@ -555,6 +557,9 @@ func (m *AppModel) handleUpdateModelKey() (tea.Model, tea.Cmd) {
 			m.message = fmt.Sprintf("Error updating model: %v", err)
 		} else {
 			m.message = message
+			// Automatically return to main view after editing
+			m.view = MainView
+			m.editing = false
 		}
 		m.clearScreen()
 		m.refreshList()
@@ -669,6 +674,18 @@ func (m *AppModel) handlePullModelKey() (tea.Model, tea.Cmd) {
 	return m, nil
 }
 
+// handlePullKeepConfigKey handles the shift+p key to pull a model while preserving user config
+func (m *AppModel) handlePullKeepConfigKey() (tea.Model, tea.Cmd) {
+	logging.DebugLogger.Println("PullKeepConfig key matched")
+	if item, ok := m.list.SelectedItem().(Model); ok {
+		m.message = styles.InfoStyle().Render(fmt.Sprintf("Pulling model & preserving config: %s\n", item.Name))
+		m.pulling = true
+		m.pullProgress = 0
+		return m, m.startPullModelPreserveConfig(item.Name)
+	}
+	return m, nil
+}
+
 func (m *AppModel) handlePullNewModelKey() (tea.Model, tea.Cmd) {
 	m.pullInput = textinput.New()
 	m.pullInput.Placeholder = "Enter model name (e.g. llama3:8b-instruct)"

keymap.go

@@ -25,6 +25,7 @@ type KeyMap struct {
 	CopyModel        key.Binding
 	PushModel        key.Binding
 	PullModel        key.Binding
+	PullKeepConfig   key.Binding
 	Top              key.Binding
 	AltScreen        key.Binding
 	EditModel        key.Binding
@@ -57,6 +58,7 @@ func NewKeyMap() *KeyMap {
 		LinkModel:        key.NewBinding(key.WithKeys("l"), key.WithHelp("l", "link (L=all)")),
 		PushModel:        key.NewBinding(key.WithKeys("P"), key.WithHelp("P", "push")),
 		PullModel:        key.NewBinding(key.WithKeys("p"), key.WithHelp("p", "pull")),
+		PullKeepConfig:   key.NewBinding(key.WithKeys("ctrl+k"), key.WithHelp("ctrl+k", "pull & keep config")),
 		PullNewModel:     key.NewBinding(key.WithKeys("ctrl+p"), key.WithHelp("ctrl+p", "pull new model")),
 		Quit:             key.NewBinding(key.WithKeys("q")),
 		RunModel:         key.NewBinding(key.WithKeys("enter"), key.WithHelp("enter", "run")),

operations.go

@@ -235,6 +235,13 @@ func editModelfile(client *api.Client, modelName string) (string, error) {
 	// Write the fetched content to a temporary file
 	tempDir := os.TempDir()
 	newModelfilePath := filepath.Join(tempDir, fmt.Sprintf("%s_modelfile.txt", modelName))
+
+	// Ensure parent directories exist
+	parentDir := filepath.Dir(newModelfilePath)
+	if err := os.MkdirAll(parentDir, 0755); err != nil {
+		return "", fmt.Errorf("error creating directory for modelfile: %v", err)
+	}
+
 	err = os.WriteFile(newModelfilePath, []byte(modelfileContent), 0644)
 	if err != nil {
 		return "", fmt.Errorf("error writing modelfile to temp file: %v", err)
@@ -327,7 +334,7 @@ func editModelfile(client *api.Client, modelName string) (string, error) {
 	// log to the console if we're not in a tea app
 	fmt.Printf("Model %s updated successfully\n", modelName)
 
-	return fmt.Sprintf("Model %s updated successfully, Press 'q' to return to the models list", modelName), nil
+	return fmt.Sprintf("Model %s updated successfully", modelName), nil
 }
 
 func isLocalhost(url string) bool {
@@ -524,6 +531,78 @@ func (m *AppModel) startPullModel(modelName string) tea.Cmd {
 	}
 }
 
+// startPullModelPreserveConfig pulls a model while preserving user-modified configuration
+func (m *AppModel) startPullModelPreserveConfig(modelName string) tea.Cmd {
+	return func() tea.Msg {
+		ctx := context.Background()
+
+		// Step 1: Extract current parameters and template before pulling
+		logging.InfoLogger.Printf("Extracting parameters for model %s before pulling", modelName)
+		currentParams, currentTemplate, systemPrompt, err := getModelParamsWithSystem(modelName, m.client)
+		if err != nil {
+			logging.ErrorLogger.Printf("Error extracting parameters for model %s: %v", modelName, err)
+			return pullErrorMsg{fmt.Errorf("failed to extract parameters: %v", err)}
+		}
+
+		// Step 2: Pull the updated model
+		logging.InfoLogger.Printf("Pulling updated model: %s", modelName)
+		req := &api.PullRequest{Name: modelName}
+		err = m.client.Pull(ctx, req, func(resp api.ProgressResponse) error {
+			m.pullProgress = float64(resp.Completed) / float64(resp.Total)
+			return nil
+		})
+		if err != nil {
+			return pullErrorMsg{err}
+		}
+
+		// Step 3: Apply the saved configuration back to the updated model
+		logging.InfoLogger.Printf("Restoring configuration for model: %s", modelName)
+
+		// Create request with base fields
+		createReq := &api.CreateRequest{
+			Model: modelName, // The model to update
+			From:  modelName, // Use the same model name as base (it's now been updated)
+		}
+
+		// Add template if it exists
+		if currentTemplate != "" {
+			createReq.Template = currentTemplate
+		}
+
+		// Add system prompt if it exists
+		if systemPrompt != "" {
+			createReq.System = systemPrompt
+		}
+
+		// Add parameters if any were found
+		if len(currentParams) > 0 {
+			// Convert map[string]string to map[string]any
+			parameters := make(map[string]any)
+			for k, v := range currentParams {
+				// Try to convert numeric values
+				if floatVal, err := strconv.ParseFloat(v, 64); err == nil {
+					parameters[k] = floatVal
+				} else if intVal, err := strconv.Atoi(v); err == nil {
+					parameters[k] = intVal
+				} else {
+					parameters[k] = v
+				}
+			}
+			createReq.Parameters = parameters
+		}
+
+		// Apply the configuration
+		err = m.client.Create(ctx, createReq, func(resp api.ProgressResponse) error {
+			return nil
+		})
+		if err != nil {
+			return pullErrorMsg{fmt.Errorf("failed to restore configuration: %v", err)}
+		}
+
+		return pullSuccessMsg{modelName}
+	}
+}
+
 type editorFinishedMsg struct{ err error }
 
 func cleanupSymlinkedModels(lmStudioModelsDir string) {
@@ -796,6 +875,80 @@ func extractTemplateAndSystem(content string) (template string, system string) {
 	return template, system
 }
 
+// getModelParamsWithSystem extracts parameters, template, and system prompt from a model's modelfile
+func getModelParamsWithSystem(modelName string, client *api.Client) (map[string]string, string, string, error) {
+	logging.InfoLogger.Printf("Getting parameters and system prompt for model: %s\n", modelName)
+	ctx := context.Background()
+	req := &api.ShowRequest{Name: modelName}
+	resp, err := client.Show(ctx, req)
+	if err != nil {
+		logging.ErrorLogger.Printf("Error getting modelfile for %s: %v\n", modelName, err)
+		return nil, "", "", err
+	}
+	output := []byte(resp.Modelfile)
+	lines := strings.Split(strings.TrimSpace(string(output)), "\n")
+	params := make(map[string]string)
+	var template string
+	var system string
+
+	inTemplate := false
+	inMultilineTemplate := false
+	var templateLines []string
+
+	for _, line := range lines {
+		trimmed := strings.TrimSpace(line)
+
+		// Handle TEMPLATE directive
+		if strings.HasPrefix(trimmed, "TEMPLATE") {
+			if strings.Contains(trimmed, `"""`) {
+				// Multi-line template
+				templateContent := strings.TrimPrefix(trimmed, "TEMPLATE ")
+				templateContent = strings.TrimSpace(templateContent)
+				if strings.HasPrefix(templateContent, `"""`) {
+					templateContent = strings.TrimPrefix(templateContent, `"""`)
+				}
+				inTemplate = true
+				inMultilineTemplate = true
+				if templateContent != "" {
+					templateLines = append(templateLines, templateContent)
+				}
+			} else {
+				// Single-line template
+				template = strings.TrimPrefix(trimmed, "TEMPLATE ")
+				template = strings.Trim(template, `"`)
+			}
+		} else if inTemplate {
+			if inMultilineTemplate && strings.HasSuffix(trimmed, `"""`) {
+				line = strings.TrimSuffix(line, `"""`)
+				if line != "" {
+					templateLines = append(templateLines, line)
+				}
+				inTemplate = false
+				inMultilineTemplate = false
+			} else {
+				templateLines = append(templateLines, line)
+			}
+		} else if strings.HasPrefix(trimmed, "SYSTEM") {
+			system = strings.TrimPrefix(trimmed, "SYSTEM ")
+			// Remove surrounding quotes if present
+			system = strings.Trim(system, `"`)
+		} else if strings.HasPrefix(trimmed, "PARAMETER") {
+			parts := strings.SplitN(trimmed, " ", 3)
+			if len(parts) >= 3 {
+				key := parts[1]
+				value := strings.TrimSpace(parts[2])
+				params[key] = value
+			}
+		}
+	}
+
+	if len(templateLines) > 0 {
+		template = strings.Join(templateLines, "\n")
+	}
+
+	return params, template, system, nil
+}
+
 // extractParameters extracts parameter values from modelfile content
 func extractParameters(content string) map[string]any {
 	parameters := make(map[string]any)