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

stop all spinners on progress stop

Browse source
This commit is contained in:
Michael Yang 2023-11-15 16:59:49 -08:00
parent 4d677ee389
commit 976068369b
2 changed files with 24 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
cmd/cmd.go
View file

@ -53,7 +53,7 @@ func CreateHandler(cmd *cobra.Command, args []string) error {
bars := make(map[string]*progress.Bar) bars := make(map[string]*progress.Bar)
status := "transferring context" status := fmt.Sprintf("creating %s", args[0])
spinner := progress.NewSpinner(status) spinner := progress.NewSpinner(status)
p.Add(status, spinner) p.Add(status, spinner)
@ -72,6 +72,12 @@ func CreateHandler(cmd *cobra.Command, args []string) error {
return err return err
} }
spinner.Stop()
status = "transferring context"
spinner = progress.NewSpinner(status)
p.Add(status, spinner)
for _, c := range commands { for _, c := range commands {
switch c.Name { switch c.Name {
case "model", "adapter": case "model", "adapter":
@ -133,10 +139,6 @@ func CreateHandler(cmd *cobra.Command, args []string) error {
return err return err
} }
if spinner != nil {
spinner.Stop()
}
return nil return nil
} }
@ -178,15 +180,13 @@ func PushHandler(cmd *cobra.Command, args []string) error {
bars := make(map[string]*progress.Bar) bars := make(map[string]*progress.Bar)
var status string status := fmt.Sprintf("pushing %s", args[0])
var spinner *progress.Spinner spinner := progress.NewSpinner(status)
p.Add(status, spinner)
fn := func(resp api.ProgressResponse) error { fn := func(resp api.ProgressResponse) error {
if resp.Digest != "" { if resp.Digest != "" {
if spinner != nil {
spinner.Stop() spinner.Stop()
spinner = nil
}
bar, ok := bars[resp.Digest] bar, ok := bars[resp.Digest]
if !ok { if !ok {
@ -197,10 +197,7 @@ func PushHandler(cmd *cobra.Command, args []string) error {
bar.Set(resp.Completed) bar.Set(resp.Completed)
} else if status != resp.Status { } else if status != resp.Status {
if spinner != nil {
spinner.Stop() spinner.Stop()
spinner = nil
}
status = resp.Status status = resp.Status
spinner = progress.NewSpinner(status) spinner = progress.NewSpinner(status)
@ -215,10 +212,7 @@ func PushHandler(cmd *cobra.Command, args []string) error {
return err return err
} }
if spinner != nil {
spinner.Stop() spinner.Stop()
}
return nil return nil
} }
@ -379,15 +373,13 @@ func PullHandler(cmd *cobra.Command, args []string) error {
bars := make(map[string]*progress.Bar) bars := make(map[string]*progress.Bar)
var status string status := fmt.Sprintf("pulling %s", args[0])
var spinner *progress.Spinner spinner := progress.NewSpinner(status)
p.Add(status, spinner)
fn := func(resp api.ProgressResponse) error { fn := func(resp api.ProgressResponse) error {
if resp.Digest != "" { if resp.Digest != "" {
if spinner != nil {
spinner.Stop() spinner.Stop()
spinner = nil
}
bar, ok := bars[resp.Digest] bar, ok := bars[resp.Digest]
if !ok { if !ok {
@ -398,10 +390,7 @@ func PullHandler(cmd *cobra.Command, args []string) error {
bar.Set(resp.Completed) bar.Set(resp.Completed)
} else if status != resp.Status { } else if status != resp.Status {
if spinner != nil {
spinner.Stop() spinner.Stop()
spinner = nil
}
status = resp.Status status = resp.Status
spinner = progress.NewSpinner(status) spinner = progress.NewSpinner(status)
@ -416,10 +405,6 @@ func PullHandler(cmd *cobra.Command, args []string) error {
return err return err
} }
if spinner != nil {
spinner.Stop()
}
return nil return nil
} }
@ -476,7 +461,6 @@ func generate(cmd *cobra.Command, model, prompt string, wordWrap bool, format st
defer p.Stop() defer p.Stop()
spinner := progress.NewSpinner("") spinner := progress.NewSpinner("")
defer spinner.Stop()
p.Add("", spinner) p.Add("", spinner)
var latest api.GenerateResponse var latest api.GenerateResponse

6
progress/progress.go
View file

@ -31,6 +31,12 @@ func NewProgress(w io.Writer) *Progress {
} }
func (p *Progress) Stop() bool { func (p *Progress) Stop() bool {
for _, state := range p.states {
if spinner, ok := state.(*Spinner); ok {
spinner.Stop()
}
}
if p.ticker != nil { if p.ticker != nil {
p.ticker.Stop() p.ticker.Stop()
p.ticker = nil p.ticker = nil