diff --git a/suggest/api.go b/suggest/api.go index 67e710d..191b6a1 100644 --- a/suggest/api.go +++ b/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 diff --git a/suggest/api_test.go b/suggest/api_test.go new file mode 100644 index 0000000..c3b3e85 --- /dev/null +++ b/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") + } +}