use readline

This commit is contained in:
Michael Yang
2023-07-18 14:01:19 -07:00
parent 9c5572d51f
commit 84200dcde6
3 changed files with 103 additions and 10 deletions

View File

@@ -5,6 +5,7 @@ import (
"context"
"errors"
"fmt"
"io"
"log"
"net"
"net/http"
@@ -13,11 +14,11 @@ import (
"strings"
"time"
"github.com/chzyer/readline"
"github.com/dustin/go-humanize"
"github.com/olekukonko/tablewriter"
"github.com/schollz/progressbar/v3"
"github.com/spf13/cobra"
"golang.org/x/term"
"github.com/jmorganca/ollama/api"
"github.com/jmorganca/ollama/format"
@@ -111,7 +112,9 @@ func list(cmd *cobra.Command, args []string) error {
var data [][]string
for _, m := range models.Models {
data = append(data, []string{m.Name, humanize.Bytes(uint64(m.Size)), format.HumanTime(m.ModifiedAt, "Never")})
if len(args) == 0 || strings.HasPrefix(m.Name, args[0]) {
data = append(data, []string{m.Name, humanize.Bytes(uint64(m.Size)), format.HumanTime(m.ModifiedAt, "Never")})
}
}
table := tablewriter.NewWriter(os.Stdout)
@@ -169,7 +172,7 @@ func RunGenerate(cmd *cobra.Command, args []string) error {
return generate(cmd, args[0], strings.Join(args[1:], " "))
}
if term.IsTerminal(int(os.Stdin.Fd())) {
if readline.IsTerminal(int(os.Stdin.Fd())) {
return generateInteractive(cmd, args[0])
}
@@ -227,17 +230,99 @@ func generate(cmd *cobra.Command, model, prompt string) error {
}
func generateInteractive(cmd *cobra.Command, model string) error {
fmt.Print(">>> ")
scanner := bufio.NewScanner(os.Stdin)
for scanner.Scan() {
if err := generate(cmd, model, scanner.Text()); err != nil {
home, err := os.UserHomeDir()
if err != nil {
return err
}
completer := readline.NewPrefixCompleter(
readline.PcItem("/help"),
readline.PcItem("/list"),
readline.PcItem("/set",
readline.PcItem("history"),
readline.PcItem("nohistory"),
readline.PcItem("mode",
readline.PcItem("vim"),
readline.PcItem("emacs"),
readline.PcItem("default"),
),
),
readline.PcItem("/exit"),
readline.PcItem("/bye"),
)
usage := func() {
fmt.Fprintln(os.Stderr, "commands:")
fmt.Fprintln(os.Stderr, completer.Tree(" "))
}
config := readline.Config{
Prompt: ">>> ",
HistoryFile: filepath.Join(home, ".ollama", "history"),
AutoComplete: completer,
}
scanner, err := readline.NewEx(&config)
if err != nil {
return err
}
defer scanner.Close()
for {
line, err := scanner.Readline()
switch {
case errors.Is(err, io.EOF):
return nil
case errors.Is(err, readline.ErrInterrupt):
continue
case err != nil:
return err
}
fmt.Print(">>> ")
}
line = strings.TrimSpace(line)
return nil
switch {
case strings.HasPrefix(line, "/list"):
args := strings.Fields(line)
if err := list(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 "mode":
if len(args) > 2 {
switch args[2] {
case "vim":
scanner.SetVimMode(true)
continue
case "emacs", "default":
scanner.SetVimMode(false)
continue
}
}
}
}
case line == "/help", line == "/?":
usage()
continue
case line == "/exit", line == "/bye":
return nil
}
if err := generate(cmd, model, line); err != nil {
return err
}
}
}
func generateBatch(cmd *cobra.Command, model string) error {