Merge pull request #18 from hanchchch/main

Trim shell prefix

suggest/api.go

@@ -3,11 +3,12 @@ package suggest
 import (
 	"context"
 	"fmt"
-	ollama "github.com/jmorganca/ollama/api"
-	"github.com/yusufcanb/tlm/shell"
 	"regexp"
 	"runtime"
 	"strings"
+
+	ollama "github.com/jmorganca/ollama/api"
+	"github.com/yusufcanb/tlm/shell"
 )
 
 const (
@@ -16,6 +17,8 @@ const (
 	Creative        = "creative"
 )
 
+var ShellPrefix = []string{"$", "❯"}
+
 func (s *Suggest) getParametersFor(preference string) map[string]interface{} {
 	switch preference {
 	case Stable:
@@ -56,13 +59,30 @@ func (s *Suggest) extractCommandsFromResponse(response string) []string {
 	var codeSnippets []string
 	for _, match := range matches {
 		if len(match) == 3 {
-			codeSnippets = append(codeSnippets, match[2])
+			codeSnippets = append(codeSnippets, s.refineCommand(match[2]))
 		}
 	}
 
 	return codeSnippets
 }
 
+func (s *Suggest) refineCommand(command string) string {
+	result := strings.Clone(command)
+
+	// Trim shell prefixes
+	for _, prefix := range ShellPrefix {
+		if strings.HasPrefix(result, prefix) {
+			result = strings.TrimPrefix(result, prefix)
+			break
+		}
+	}
+
+	// Trim leading and trailing whitespaces
+	result = strings.TrimSpace(result)
+
+	return result
+}
+
 func (s *Suggest) getCommandSuggestionFor(mode, term string, prompt string) (string, error) {
 	var responseText string
 

suggest/api_test.go

@@ -0,0 +1,32 @@
+package suggest
+
+import (
+	"testing"
+
+	ollama "github.com/jmorganca/ollama/api"
+	"github.com/yusufcanb/tlm/config"
+)
+
+func TestRefineCommand(t *testing.T) {
+	con := config.New()
+	con.LoadOrCreateConfig()
+
+	o, _ := ollama.ClientFromEnvironment()
+	s := New(o)
+
+	if s.refineCommand("ls -al") != "ls -al" {
+		t.Error("no change should be made if the command is already okay")
+	}
+
+	if s.refineCommand("$ ls -al") != "ls -al" {
+		t.Error("shell prefix should be removed")
+	}
+
+	if s.refineCommand("❯ ls -al") != "ls -al" {
+		t.Error("shell prefix should be removed")
+	}
+
+	if s.refineCommand(" ls -al") != "ls -al" {
+		t.Error("leading space should be removed")
+	}
+}