grokkit/cmd/chat.go

package cmd

import (
	"os"
	"strings"

	"github.com/charmbracelet/bubbles/spinner"
	"github.com/charmbracelet/bubbles/textinput"
	"github.com/charmbracelet/bubbles/viewport"
	tea "github.com/charmbracelet/bubbletea"
	"github.com/charmbracelet/lipgloss"
	"github.com/fatih/color"
	"github.com/spf13/cobra"
	"github.com/spf13/viper"
	"gmgauthier.com/grokkit/internal/grok"
)

type (
	Message struct {
		Role    string
		Content string
	}
	model struct {
		messages  []Message
		history   []map[string]string
		input     textinput.Model
		spinner   spinner.Model
		viewport  viewport.Model
		streamCh  <-chan string
		streaming bool
		client    *grok.Client
		gmodel    string
	}
	streamTickMsg struct{}
)

var (
	userStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("172")).
			Padding(0, 1)
	grokStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("45")).
			Padding(0, 1)
	headerStyle = lipgloss.NewStyle().
			Foreground(lipgloss.Color("99")).
			Bold(true)
)

func initialModel() *model {
	systemContent := "You are Grok, a helpful and maximally truthful AI built by xAI, not based on any other companies and their models."
	system := Message{Role: "system", Content: systemContent}
	h := []map[string]string{{"role": "system", "content": systemContent}}
	ti := textinput.New()
	ti.Placeholder = "Type message (Enter send, q/Ctrl+C quit)"
	ti.Focus()
	sp := spinner.New()
	sp.Spinner = spinner.Dot
	sp.Style = lipgloss.NewStyle().Foreground(lipgloss.Color("63"))
	vp := viewport.New(0, 0)
	vp.MouseWheelEnabled = true
	vp.GotoBottom()
	return &model{
		messages: []Message{system},
		history:  h,
		input:    ti,
		spinner:  sp,
		viewport: vp,
		client:   grok.NewClient(),
		gmodel:   viper.GetString("model"),
	}
}

func (m model) renderMessages() string {
	var lines []string
	for _, msg := range m.messages {
		content := strings.ReplaceAll(msg.Content, "\n", "\n  ")
		switch msg.Role {
		case "user":
			lines = append(lines, userStyle.Render("You: "+content))
		case "assistant":
			lines = append(lines, grokStyle.Render("Grok: "+content))
		case "system":
			lines = append(lines, content)
		}
		lines = append(lines, "")
	}
	return strings.Join(lines, "\n")
}

func streamTickMsgCmd() tea.Cmd {
	return func() tea.Msg {
		return streamTickMsg{}
	}
}

func (m *model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
	var cmds []tea.Cmd
	switch msg := msg.(type) {
	case tea.WindowSizeMsg:
		m.viewport.Width = msg.Width
		if msg.Height > 3 {
			m.viewport.Height = msg.Height - 3
		}
		m.viewport.SetContent(m.renderMessages())
		return m, nil
	case tea.KeyMsg:
		if m.streaming {
			if msg.String() == "q" || msg.String() == "ctrl+c" || msg.String() == "esc" {
				return m, tea.Quit
			}
			return m, nil
		}
		switch msg.String() {
		case "q", "ctrl+c", "esc":
			return m, tea.Quit
		case "enter":
			text := m.input.Value()
			if text == "" {
				return m, nil
			}
			userMsg := Message{Role: "user", Content: text}
			m.messages = append(m.messages, userMsg)
			m.history = append(m.history, map[string]string{"role": "user", "content": text})
			assMsg := Message{Role: "assistant", Content: ""}
			m.messages = append(m.messages, assMsg)
			m.input.Reset()
			m.streaming = true
			m.streamCh = m.client.StreamChan(m.history, m.gmodel)
			m.viewport.SetContent(m.renderMessages())
			cmds = append(cmds, streamTickMsgCmd())
			return m, tea.Batch(cmds...)
		default:
			var cmd tea.Cmd
			m.input, cmd = m.input.Update(msg)
			return m, cmd
		}
	case streamTickMsg:
		if !m.streaming || m.streamCh == nil {
			return m, nil
		}
		select {
		case chunk, ok := <-m.streamCh:
			if !ok {
				m.streaming = false
				m.streamCh = nil
				m.viewport.SetContent(m.renderMessages())
				return m, nil
			}
			m.messages[len(m.messages)-1].Content += chunk
			m.viewport.SetContent(m.renderMessages())
			m.viewport.GotoBottom()
			return m, streamTickMsgCmd()
		default:
			return m, streamTickMsgCmd()
		}
	}
	return m, tea.Batch(cmds...)
}

func (m model) View() string {
	header := headerStyle.Render("=== Grokkit TUI Chat ===  q/Ctrl+C to quit")
	vpView := m.viewport.View()
	spinnerLine := ""
	if m.streaming {
		spinnerLine = m.spinner.View() + " Grok typing..."
	}
	inputLine := m.input.View()
	return lipgloss.JoinVertical(lipgloss.Left, header, vpView, spinnerLine, inputLine)
}

func (m model) Init() tea.Cmd {
	return nil
}

var chatCmd = &cobra.Command{
	Use:   "chat",
	Short: "Interactive TUI chat with Grok",
	Run: func(cmd *cobra.Command, args []string) {
		p := tea.NewProgram(initialModel(), tea.WithAltScreen(), tea.WithMouseCellMotion())
		if _, err := p.Run(); err != nil {
			color.Red("Chat failed: %v", err)
			os.Exit(1)
		}
	},
}