91 lines
1.4 KiB
Go
91 lines
1.4 KiB
Go
|
package internal
|
||
|
|
|
||
|
|
import (
|
||
|
|
"fmt"
|
||
|
|
"os"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"charm.land/bubbles/v2/spinner"
|
||
|
|
"charm.land/lipgloss/v2"
|
||
|
|
)
|
||
|
|
|
||
|
|
type SpinnerState int
|
||
|
|
|
||
|
|
const (
|
||
|
|
StateThinking SpinnerState = iota
|
||
|
|
StateAnswering
|
||
|
|
StateDone
|
||
|
|
)
|
||
|
|
|
||
|
|
type Spinner struct {
|
||
|
|
text string
|
||
|
|
state SpinnerState
|
||
|
|
spinner spinner.Model
|
||
|
|
stopCh chan struct{}
|
||
|
|
doneCh chan struct{}
|
||
|
|
}
|
||
|
|
|
||
|
|
func NewSpinner(text string) *Spinner {
|
||
|
|
s := spinner.New(
|
||
|
|
spinner.WithSpinner(spinner.MiniDot),
|
||
|
|
spinner.WithStyle(lipgloss.NewStyle().Foreground(lipgloss.Color("243"))),
|
||
|
|
)
|
||
|
|
return &Spinner{
|
||
|
|
text: text,
|
||
|
|
state: StateThinking,
|
||
|
|
spinner: s,
|
||
|
|
stopCh: make(chan struct{}),
|
||
|
|
doneCh: make(chan struct{}),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Spinner) Start() {
|
||
|
|
go s.run()
|
||
|
|
s.tick()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Spinner) Update(text string) {
|
||
|
|
s.text = text
|
||
|
|
s.state = StateAnswering
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Spinner) Stop() {
|
||
|
|
close(s.stopCh)
|
||
|
|
<-s.doneCh
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Spinner) tick() {
|
||
|
|
msg := s.spinner.Tick()
|
||
|
|
if msg, ok := msg.(spinner.TickMsg); ok {
|
||
|
|
s.spinner, _ = s.spinner.Update(msg)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Spinner) run() {
|
||
|
|
defer close(s.doneCh)
|
||
|
|
|
||
|
|
ticker := time.NewTicker(time.Second / 12)
|
||
|
|
defer ticker.Stop()
|
||
|
|
|
||
|
|
for {
|
||
|
|
select {
|
||
|
|
case <-ticker.C:
|
||
|
|
s.tick()
|
||
|
|
s.render()
|
||
|
|
case <-s.stopCh:
|
||
|
|
s.clear()
|
||
|
|
return
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Spinner) render() {
|
||
|
|
fmt.Printf("\r%s %s", s.spinner.View(), s.text)
|
||
|
|
os.Stdout.Sync()
|
||
|
|
}
|
||
|
|
|
||
|
|
func (s *Spinner) clear() {
|
||
|
|
fmt.Printf("\r%s 思考完成.\n", s.spinner.View())
|
||
|
|
os.Stdout.Sync()
|
||
|
|
}
|