Fix issues with inputting and formatting multi line strings in `ollam…

…a run`

Co-authored-by: Wen Sun <[email protected]>

cmd/cmd.go

@@ -602,22 +602,20 @@ func generateInteractive(cmd *cobra.Command, model string, wordWrap bool, format
 		fmt.Fprintln(os.Stderr, "")
 	}
 
-	prompt := readline.Prompt{
+	scanner, err := readline.New(readline.Prompt{
 		Prompt:         ">>> ",
 		AltPrompt:      "... ",
 		Placeholder:    "Send a message (/? for help)",
 		AltPlaceholder: `Use """ to end multi-line input`,
-	}
-
-	scanner, err := readline.New(prompt)
+	})
 	if err != nil {
 		return err
 	}
 
 	fmt.Print(readline.StartBracketedPaste)
 	defer fmt.Printf(readline.EndBracketedPaste)
 
-	var multiLineBuffer string
+	var prompt string
 
 	for {
 		line, err := scanner.Readline()
@@ -630,27 +628,33 @@ func generateInteractive(cmd *cobra.Command, model string, wordWrap bool, format
 				fmt.Println("\nUse Ctrl-D or /bye to exit.")
 			}
 
+			scanner.Prompt.UseAlt = false
+			prompt = ""
+
 			continue
 		case err != nil:
 			return err
 		}
 
-		line = strings.TrimSpace(line)
-
 		switch {
-		case scanner.Prompt.UseAlt:
-			if strings.HasSuffix(line, `"""`) {
-				scanner.Prompt.UseAlt = false
-				multiLineBuffer += strings.TrimSuffix(line, `"""`)
-				line = multiLineBuffer
-				multiLineBuffer = ""
-			} else {
-				multiLineBuffer += line + " "
+		case strings.HasPrefix(prompt, `"""`):
+			// if the prompt so far starts with """ then we're in multiline mode
+			// and we need to keep reading until we find a line that ends with """
+			cut, found := strings.CutSuffix(line, `"""`)
+			prompt += cut + "\n"
+
+			if !found {
 				continue
 			}
-		case strings.HasPrefix(line, `"""`):
+
+			prompt = strings.TrimPrefix(prompt, `"""`)
+			scanner.Prompt.UseAlt = false
+		case strings.HasPrefix(line, `"""`) && len(prompt) == 0:
 			scanner.Prompt.UseAlt = true
-			multiLineBuffer = strings.TrimPrefix(line, `"""`) + " "
+			prompt += line + "\n"
+			continue
+		case scanner.Pasting:
+			prompt += line + "\n"
 			continue
 		case strings.HasPrefix(line, "/list"):
 			args := strings.Fields(line)
@@ -757,12 +761,17 @@ func generateInteractive(cmd *cobra.Command, model string, wordWrap bool, format
 		case strings.HasPrefix(line, "/"):
 			args := strings.Fields(line)
 			fmt.Printf("Unknown command '%s'. Type /? for help\n", args[0])
+			continue
+		default:
+			prompt += line
 		}
 
-		if len(line) > 0 && line[0] != '/' {
-			if err := generate(cmd, model, line, wordWrap, format); err != nil {
+		if len(prompt) > 0 && prompt[0] != '/' {
+			if err := generate(cmd, model, prompt, wordWrap, format); err != nil {
 				return err
 			}
+
+			prompt = ""
 		}
 	}
 }

readline/readline.go

@@ -24,6 +24,7 @@ type Instance struct {
 	Prompt   *Prompt
 	Terminal *Terminal
 	History  *History
+	Pasting  bool
 }
 
 func New(prompt Prompt) (*Instance, error) {
@@ -46,7 +47,7 @@ func New(prompt Prompt) (*Instance, error) {
 
 func (i *Instance) Readline() (string, error) {
 	prompt := i.Prompt.Prompt
-	if i.Prompt.UseAlt {
+	if i.Prompt.UseAlt || i.Pasting {
 		prompt = i.Prompt.AltPrompt
 	}
 	fmt.Print(prompt)
@@ -63,12 +64,13 @@ func (i *Instance) Readline() (string, error) {
 	var esc bool
 	var escex bool
 	var metaDel bool
-	var pasteMode PasteMode
 
 	var currentLineBuf []rune
 
 	for {
-		if buf.IsEmpty() {
+		// don't show placeholder when pasting unless we're in multiline mode
+		showPlaceholder := !i.Pasting || i.Prompt.UseAlt
+		if buf.IsEmpty() && showPlaceholder {
 			ph := i.Prompt.Placeholder
 			if i.Prompt.UseAlt {
 				ph = i.Prompt.AltPlaceholder
@@ -119,9 +121,9 @@ func (i *Instance) Readline() (string, error) {
 					code += string(r)
 				}
 				if code == CharBracketedPasteStart {
-					pasteMode = PasteModeStart
+					i.Pasting = true
 				} else if code == CharBracketedPasteEnd {
-					pasteMode = PasteModeEnd
+					i.Pasting = false
 				}
 			case KeyDel:
 				if buf.Size() > 0 {
@@ -196,12 +198,7 @@ func (i *Instance) Readline() (string, error) {
 			}
 			buf.MoveToEnd()
 			fmt.Println()
-			switch pasteMode {
-			case PasteModeStart:
-				output = `"""` + output
-			case PasteModeEnd:
-				output = output + `"""`
-			}
+
 			return output, nil
 		default:
 			if metaDel {

readline/types.go

@@ -76,11 +76,3 @@ const (
 	CharBracketedPasteStart = "00~"
 	CharBracketedPasteEnd   = "01~"
 )
-
-type PasteMode int
-
-const (
-	PastModeOff = iota
-	PasteModeStart
-	PasteModeEnd
-)