From 76b85bc0e97e7499e100ecb2001c9573aeaba68b Mon Sep 17 00:00:00 2001
From: Michael Yang <mxyng@pm.me>
Date: Mon, 14 Aug 2023 11:15:53 -0700
Subject: [PATCH] set non-zero error code on error

---
 cmd/cmd.go | 26 +++++++++++++++++++++++---
 main.go    |  3 ++-
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/cmd/cmd.go b/cmd/cmd.go
index 53c039ba..ecdd938d 100644
--- a/cmd/cmd.go
+++ b/cmd/cmd.go
@@ -75,6 +75,7 @@ func CreateHandler(cmd *cobra.Command, args []string) error {
 			spinner = NewSpinner(resp.Status)
 			go spinner.Spin(100 * time.Millisecond)
 		}
+
 		return nil
 	}
 
@@ -84,6 +85,9 @@ func CreateHandler(cmd *cobra.Command, args []string) error {
 
 	if spinner != nil {
 		spinner.Stop()
+		if spinner.description != "success" {
+			return errors.New("unexpected end to create model")
+		}
 	}
 
 	return nil
@@ -149,6 +153,11 @@ func PushHandler(cmd *cobra.Command, args []string) error {
 	if err := client.Push(context.Background(), &request, fn); err != nil {
 		return err
 	}
+
+	if bar != nil && !bar.IsFinished() {
+		return errors.New("unexpected end to push model")
+	}
+
 	return nil
 }
 
@@ -235,12 +244,18 @@ func pull(model string, insecure bool) error {
 			currentDigest = ""
 			fmt.Println(resp.Status)
 		}
+
 		return nil
 	}
 
 	if err := client.Pull(context.Background(), &request, fn); err != nil {
 		return err
 	}
+
+	if bar != nil && !bar.IsFinished() {
+		return errors.New("unexpected end to pull model")
+	}
+
 	return nil
 }
 
@@ -304,6 +319,10 @@ func generate(cmd *cobra.Command, model, prompt string) error {
 		fmt.Println()
 		fmt.Println()
 
+		if !latest.Done {
+			return errors.New("unexpected end of response")
+		}
+
 		verbose, err := cmd.Flags().GetBool("verbose")
 		if err != nil {
 			return err
@@ -664,9 +683,10 @@ func NewCLI() *cobra.Command {
 	log.SetFlags(log.LstdFlags | log.Lshortfile)
 
 	rootCmd := &cobra.Command{
-		Use:          "ollama",
-		Short:        "Large language model runner",
-		SilenceUsage: true,
+		Use:           "ollama",
+		Short:         "Large language model runner",
+		SilenceUsage:  true,
+		SilenceErrors: true,
 		CompletionOptions: cobra.CompletionOptions{
 			DisableDefaultCmd: true,
 		},
diff --git a/main.go b/main.go
index e56c9b4f..ecdcb71d 100644
--- a/main.go
+++ b/main.go
@@ -4,8 +4,9 @@ import (
 	"context"
 
 	"github.com/jmorganca/ollama/cmd"
+	"github.com/spf13/cobra"
 )
 
 func main() {
-	cmd.NewCLI().ExecuteContext(context.Background())
+	cobra.CheckErr(cmd.NewCLI().ExecuteContext(context.Background()))
 }