Repository: ollama/ollama
Pull request: add word wrapping for lines which are longer than the terminal width

This change makes it so the REPL will properly wrap a line on a word boundary. The way it works is that it walks through each character of each token returned by the server, and then keeps a buffer of the last word. If the maximum boundary length is exceeded, it will backtrack using ANSI escape codes to the length of the current word buffer, erase to the end of the line, give a line feed, and then add the word fragment to the new line. This requires that the terminal allow ANSI graphics which should be OK for any modern terminal. If you run this headless, it will default to not wrapping any lines.

Right now I've set the width to 5 characters less than the terminal width, but we can potentially make this a setting in the future.

This fixes issue #150 


Changed files:
filename: cmd/cmd.go
patch: @@ -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,53 +594,49 @@ 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)
 			if len(args) > 1 {
 				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",

filename: go.sum
patch: @@ -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=