Skip to content

Commit

Permalink
add word wrapping for lines which are longer than the terminal width (#…
Browse files Browse the repository at this point in the history
  • Loading branch information
pdevine committed Sep 22, 2023
1 parent e1a0846 commit c928ceb
Show file tree
Hide file tree
Showing 2 changed files with 65 additions and 25 deletions.
89 changes: 65 additions & 24 deletions cmd/cmd.go
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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() {
Expand All @@ -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
}

Expand All @@ -427,7 +475,6 @@ func generate(cmd *cobra.Command, model, prompt string) error {
}
return err
}

if prompt != "" {
fmt.Println()
fmt.Println()
Expand Down Expand Up @@ -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"),
Expand Down Expand Up @@ -535,6 +579,7 @@ func generateInteractive(cmd *cobra.Command, model string) error {
line = multiLineBuffer
multiLineBuffer = ""
scanner.SetPrompt(">>> ")
continue
} else {
multiLineBuffer += line + " "
continue
Expand All @@ -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] {
Expand All @@ -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] != '/' {
Expand Down Expand Up @@ -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",
Expand Down
1 change: 0 additions & 1 deletion go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down

0 comments on commit c928ceb

Please sign in to comment.