add word wrapping for lines which are longer than the terminal width (#553)

cmd/cmd.go

@@ -18,11 +18,12 @@ import (
 	"strings"
 	"time"
 
-	"github.com/pdevine/readline"
 	"github.com/dustin/go-humanize"
 	"github.com/olekukonko/tablewriter"
+	"github.com/pdevine/readline"
 	"github.com/spf13/cobra"
 	"golang.org/x/crypto/ssh"
+	"golang.org/x/term"
 
 	"github.com/jmorganca/ollama/api"
 	"github.com/jmorganca/ollama/format"
@@ -400,6 +401,29 @@ func generate(cmd *cobra.Command, model, prompt string) error {
 		generateContext = []int{}
 	}
 
+	var wrapTerm bool
+	termType := os.Getenv("TERM")
+	if termType == "xterm-256color" {
+		wrapTerm = true
+	}
+
+	termWidth, _, err := term.GetSize(int(0))
+	if err != nil {
+		wrapTerm = false
+	}
+
+	// override wrapping if the user turned it off
+	nowrap, err := cmd.Flags().GetBool("nowordwrap")
+	if err != nil {
+		return err
+	}
+	if nowrap {
+		wrapTerm = false
+	}
+
+	var currentLineLength int
+	var wordBuffer string
+
 	request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext}
 	fn := func(response api.GenerateResponse) error {
 		if !spinner.IsFinished() {
@@ -408,7 +432,31 @@ func generate(cmd *cobra.Command, model, prompt string) error {
 
 		latest = response
 
-		fmt.Print(response.Response)
+		if wrapTerm {
+			for _, ch := range response.Response {
+				if currentLineLength+1 > termWidth-5 {
+					// backtrack the length of the last word and clear to the end of the line
+					fmt.Printf("\x1b[%dD\x1b[K\n", len(wordBuffer))
+					fmt.Printf("%s%c", wordBuffer, ch)
+					currentLineLength = len(wordBuffer) + 1
+				} else {
+					fmt.Print(string(ch))
+					currentLineLength += 1
+
+					switch ch {
+					case ' ':
+						wordBuffer = ""
+					case '\n':
+						currentLineLength = 0
+					default:
+						wordBuffer += string(ch)
+					}
+				}
+			}
+		} else {
+			fmt.Print(response.Response)
+		}
+
 		return nil
 	}
 
@@ -427,7 +475,6 @@ func generate(cmd *cobra.Command, model, prompt string) error {
 		}
 		return err
 	}
-
 	if prompt != "" {
 		fmt.Println()
 		fmt.Println()
@@ -470,13 +517,10 @@ func generateInteractive(cmd *cobra.Command, model string) error {
 		readline.PcItem("/set",
 			readline.PcItem("history"),
 			readline.PcItem("nohistory"),
+			readline.PcItem("wordwrap"),
+			readline.PcItem("nowordwrap"),
 			readline.PcItem("verbose"),
 			readline.PcItem("quiet"),
-			readline.PcItem("mode",
-				readline.PcItem("vim"),
-				readline.PcItem("emacs"),
-				readline.PcItem("default"),
-			),
 		),
 		readline.PcItem("/show",
 			readline.PcItem("license"),
@@ -535,6 +579,7 @@ func generateInteractive(cmd *cobra.Command, model string) error {
 				line = multiLineBuffer
 				multiLineBuffer = ""
 				scanner.SetPrompt(">>> ")
+				continue
 			} else {
 				multiLineBuffer += line + " "
 				continue
@@ -549,45 +594,42 @@ func generateInteractive(cmd *cobra.Command, model string) error {
 			if err := ListHandler(cmd, args[1:]); err != nil {
 				return err
 			}
-
-			continue
 		case strings.HasPrefix(line, "/set"):
 			args := strings.Fields(line)
 			if len(args) > 1 {
 				switch args[1] {
 				case "history":
 					scanner.HistoryEnable()
-					continue
 				case "nohistory":
 					scanner.HistoryDisable()
-					continue
+				case "wordwrap":
+					cmd.Flags().Set("nowordwrap", "false")
+					fmt.Println("Set 'wordwrap' mode.")
+				case "nowordwrap":
+					cmd.Flags().Set("nowordwrap", "true")
+					fmt.Println("Set 'nowordwrap' mode.")
 				case "verbose":
 					cmd.Flags().Set("verbose", "true")
-					continue
+					fmt.Println("Set 'verbose' mode.")
 				case "quiet":
 					cmd.Flags().Set("verbose", "false")
-					continue
+					fmt.Println("Set 'quiet' mode.")
 				case "mode":
 					if len(args) > 2 {
 						switch args[2] {
 						case "vim":
 							scanner.SetVimMode(true)
-							continue
 						case "emacs", "default":
 							scanner.SetVimMode(false)
-							continue
 						default:
 							usage()
-							continue
 						}
 					} else {
 						usage()
-						continue
 					}
 				}
 			} else {
 				usage()
-				continue
 			}
 		case strings.HasPrefix(line, "/show"):
 			args := strings.Fields(line)
@@ -595,7 +637,6 @@ func generateInteractive(cmd *cobra.Command, model string) error {
 				resp, err := server.GetModelInfo(model)
 				if err != nil {
 					fmt.Println("error: couldn't get model")
-					continue
 				}
 
 				switch args[1] {
@@ -612,17 +653,16 @@ func generateInteractive(cmd *cobra.Command, model string) error {
 				default:
 					fmt.Println("error: unknown command")
 				}
-
-				continue
 			} else {
 				usage()
-				continue
 			}
 		case line == "/help", line == "/?":
 			usage()
-			continue
 		case line == "/exit", line == "/bye":
 			return nil
+		case strings.HasPrefix(line, "/"):
+			args := strings.Fields(line)
+			fmt.Printf("Unknown command '%s'. Type /? for help\n", args[0])
 		}
 
 		if len(line) > 0 && line[0] != '/' {
@@ -828,6 +868,7 @@ func NewCLI() *cobra.Command {
 
 	runCmd.Flags().Bool("verbose", false, "Show timings for response")
 	runCmd.Flags().Bool("insecure", false, "Use an insecure registry")
+	runCmd.Flags().Bool("nowordwrap", false, "Don't wrap words to the next line automatically")
 
 	serveCmd := &cobra.Command{
 		Use:     "serve",

go.sum

@@ -120,7 +120,6 @@ golang.org/x/arch v0.3.0/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8=
 golang.org/x/crypto v0.0.0-20210711020723-a769d52b0f97/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/crypto v0.10.0 h1:LKqV2xt9+kDzSTfOhx4FrkEBcMrAgHSYgzywV9zcGmM=
 golang.org/x/crypto v0.10.0/go.mod h1:o4eNf7Ede1fv+hwOwZsTHl9EsPFO6q6ZvYR8vYfY45I=
-golang.org/x/exp v0.0.0-20230321023759-10a507213a29 h1:ooxPy7fPvB4kwsA2h+iBNHkAbp/4JxTSwCmvdjEYmug=
 golang.org/x/exp v0.0.0-20230817173708-d852ddb80c63 h1:m64FZMko/V45gv0bNmrNYoDEq8U5YUhetc9cBWKS1TQ=
 golang.org/x/exp v0.0.0-20230817173708-d852ddb80c63/go.mod h1:0v4NqG35kSWCMzLaMeX+IQrlSnVE/bqGSyC2cz/9Le8=
 golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=