From 11eecdde8624ed27e97cbbeb507ebaad2f162429 Mon Sep 17 00:00:00 2001
From: ParthSareen <parth.sareen@ollama.com>
Date: Thu, 12 Dec 2024 15:47:12 -0800
Subject: [PATCH] cmd: enable use of structured outputs

---
 cmd/cmd.go         | 16 +++++++++-------
 cmd/interactive.go | 13 +++++++++----
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/cmd/cmd.go b/cmd/cmd.go
index 2f77640c2..fc53fda55 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -1038,14 +1038,15 @@ func chat(cmd *cobra.Command, opts runOptions) (*api.Message, error) {
 		return nil
 	}
 
-	if opts.Format == "json" {
-		opts.Format = `"` + opts.Format + `"`
+	var format json.RawMessage
+	if opts.Format != "" {
+		format = json.RawMessage(opts.Format)
 	}
 
 	req := &api.ChatRequest{
 		Model:    opts.Model,
 		Messages: opts.Messages,
-		Format:   json.RawMessage(opts.Format),
+		Format:   format,
 		Options:  opts.Options,
 	}
 
@@ -1127,8 +1128,9 @@ func generate(cmd *cobra.Command, opts runOptions) error {
 		}
 	}
 
-	if opts.Format == "json" {
-		opts.Format = `"` + opts.Format + `"`
+	var format json.RawMessage
+	if opts.Format != "" {
+		format = json.RawMessage(opts.Format)
 	}
 
 	request := api.GenerateRequest{
@@ -1136,7 +1138,7 @@ func generate(cmd *cobra.Command, opts runOptions) error {
 		Prompt:    opts.Prompt,
 		Context:   generateContext,
 		Images:    opts.Images,
-		Format:    json.RawMessage(opts.Format),
+		Format:    format,
 		System:    opts.System,
 		Options:   opts.Options,
 		KeepAlive: opts.KeepAlive,
@@ -1353,7 +1355,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)")
+	runCmd.Flags().String("format", "", `Response format ("json" or a JSON Schema)`)
 
 	stopCmd := &cobra.Command{
 		Use:     "stop MODEL",
diff --git a/cmd/interactive.go b/cmd/interactive.go
index 9035b4c52..85a00e491 100644
--- a/cmd/interactive.go
+++ b/cmd/interactive.go
@@ -261,11 +261,16 @@ func generateInteractive(cmd *cobra.Command, opts runOptions) error {
 					}
 					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'")
+					if len(args) < 3 {
+						fmt.Println("Invalid or missing format. For 'json' mode use '/set format json or provide a JSON schema'")
+					} else if len(args) == 3 && (args[2] == "json" || args[2] == `"json"`) {
+						opts.Format = `"json"`
+						fmt.Println("Set format to 'json' mode.")
+					} else if len(args) > 3 && strings.HasPrefix(args[2], "{") {
+						opts.Format = strings.Join(args[2:], " ")
+						fmt.Printf("Set format to schema: \n'%s'.\n", opts.Format)
 					} else {
-						opts.Format = args[2]
-						fmt.Printf("Set format to '%s' mode.\n", args[2])
+						fmt.Println("Invalid or missing format. For 'json' mode use '/set format json or provide a JSON schema'")
 					}
 				case "noformat":
 					opts.Format = ""