baalajimaestro/ollama
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

add json mode to cli (#1095)

Browse source
This commit is contained in:
Jeffrey Morgan 2023-11-13 21:54:02 -05:00 committed by GitHub
parent cb745965ce
commit c1844bbee2
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 25 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

31
cmd/cmd.go
View file

@ -365,7 +365,12 @@ func RunGenerate(cmd *cobra.Command, args []string) error {
wordWrap = false
}
return generate(cmd, args[0], strings.Join(args[1:], " "), wordWrap)
format, err := cmd.Flags().GetString("format")
if err != nil {
return err
}
return generate(cmd, args[0], strings.Join(args[1:], " "), wordWrap, format)
}
if readline.IsTerminal(int(os.Stdin.Fd())) {
@ -377,7 +382,7 @@ func RunGenerate(cmd *cobra.Command, args []string) error {
type generateContextKey string
func generate(cmd *cobra.Command, model, prompt string, wordWrap bool) error {
func generate(cmd *cobra.Command, model, prompt string, wordWrap bool, format string) error {
client, err := api.ClientFromEnvironment()
if err != nil {
return err
@ -414,7 +419,7 @@ func generate(cmd *cobra.Command, model, prompt string, wordWrap bool) error {
var currentLineLength int
var wordBuffer string
request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext}
request := api.GenerateRequest{Model: model, Prompt: prompt, Context: generateContext, Format: format}
fn := func(response api.GenerateResponse) error {
if !spinner.IsFinished() {
spinner.Finish()
@ -487,7 +492,7 @@ func generate(cmd *cobra.Command, model, prompt string, wordWrap bool) error {
func generateInteractive(cmd *cobra.Command, model string) error {
// load the model
if err := generate(cmd, model, "", false); err != nil {
if err := generate(cmd, model, "", false, ""); err != nil {
return err
}
@ -508,6 +513,8 @@ func generateInteractive(cmd *cobra.Command, model string) error {
fmt.Fprintln(os.Stderr, " /set nohistory Disable history")
fmt.Fprintln(os.Stderr, " /set wordwrap Enable wordwrap")
fmt.Fprintln(os.Stderr, " /set nowordwrap Disable wordwrap")
fmt.Fprintln(os.Stderr, " /set format json Enable JSON mode")
fmt.Fprintln(os.Stderr, " /set noformat Disable formatting")
fmt.Fprintln(os.Stderr, " /set verbose Show LLM stats")
fmt.Fprintln(os.Stderr, " /set quiet Disable LLM stats")
fmt.Fprintln(os.Stderr, "")
@ -535,6 +542,7 @@ func generateInteractive(cmd *cobra.Command, model string) error {
return err
}
var format string
var wordWrap bool
termType := os.Getenv("TERM")
if termType == "xterm-256color" {
@ -613,6 +621,16 @@ func generateInteractive(cmd *cobra.Command, model string) error {
case "quiet":
cmd.Flags().Set("verbose", "false")
fmt.Println("Set 'quiet' mode.")
case "format":
if len(args) < 3 || args[2] != "json" {
fmt.Println("Invalid or missing format. For 'json' mode use '/set format json'")
} else {
format = args[2]
fmt.Printf("Set format to '%s' mode.\n", args[2])
}
case "noformat":
format = ""
fmt.Println("Disabled format.")
default:
fmt.Printf("Unknown command '/set %s'. Type /? for help\n", args[1])
}
@ -686,7 +704,7 @@ func generateInteractive(cmd *cobra.Command, model string) error {
}
if len(line) > 0 && line[0] != '/' {
if err := generate(cmd, model, line, wordWrap); err != nil {
if err := generate(cmd, model, line, wordWrap, format); err != nil {
return err
}
}
@ -698,7 +716,7 @@ func generateBatch(cmd *cobra.Command, model string) error {
for scanner.Scan() {
prompt := scanner.Text()
fmt.Printf(">>> %s\n", prompt)
if err := generate(cmd, model, prompt, false); err != nil {
if err := generate(cmd, model, prompt, false, ""); err != nil {
return err
}
}
@ -883,6 +901,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")
runCmd.Flags().String("format", "", "Response format (e.g. json)")
serveCmd := &cobra.Command{
Use: "serve",