chat.go

  1package page
  2
  3import (
  4	"context"
  5
  6	"github.com/charmbracelet/bubbles/key"
  7	tea "github.com/charmbracelet/bubbletea"
  8	"github.com/kujtimiihoxha/termai/internal/app"
  9	"github.com/kujtimiihoxha/termai/internal/llm/agent"
 10	"github.com/kujtimiihoxha/termai/internal/session"
 11	"github.com/kujtimiihoxha/termai/internal/tui/components/chat"
 12	"github.com/kujtimiihoxha/termai/internal/tui/layout"
 13	"github.com/kujtimiihoxha/termai/internal/tui/util"
 14)
 15
 16var ChatPage PageID = "chat"
 17
 18type chatPage struct {
 19	app     *app.App
 20	layout  layout.SplitPaneLayout
 21	session session.Session
 22}
 23
 24type ChatKeyMap struct {
 25	NewSession key.Binding
 26}
 27
 28var keyMap = ChatKeyMap{
 29	NewSession: key.NewBinding(
 30		key.WithKeys("ctrl+n"),
 31		key.WithHelp("ctrl+n", "new session"),
 32	),
 33}
 34
 35func (p *chatPage) Init() tea.Cmd {
 36	// TODO: remove
 37	cmds := []tea.Cmd{
 38		p.layout.Init(),
 39	}
 40
 41	sessions, _ := p.app.Sessions.List(context.Background())
 42	if len(sessions) > 0 {
 43		p.session = sessions[0]
 44		cmd := p.setSidebar()
 45		cmds = append(cmds, util.CmdHandler(chat.SessionSelectedMsg(p.session)), cmd)
 46	}
 47	return tea.Batch(
 48		cmds...,
 49	)
 50}
 51
 52func (p *chatPage) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 53	switch msg := msg.(type) {
 54	case tea.WindowSizeMsg:
 55		p.layout.SetSize(msg.Width, msg.Height)
 56	case chat.SendMsg:
 57		cmd := p.sendMessage(msg.Text)
 58		if cmd != nil {
 59			return p, cmd
 60		}
 61	case tea.KeyMsg:
 62		switch {
 63		case key.Matches(msg, keyMap.NewSession):
 64			p.session = session.Session{}
 65			p.clearSidebar()
 66			return p, util.CmdHandler(chat.SessionClearedMsg{})
 67		}
 68	}
 69	u, cmd := p.layout.Update(msg)
 70	p.layout = u.(layout.SplitPaneLayout)
 71	if cmd != nil {
 72		return p, cmd
 73	}
 74	return p, nil
 75}
 76
 77func (p *chatPage) setSidebar() tea.Cmd {
 78	sidebarContainer := layout.NewContainer(
 79		chat.NewSidebarCmp(p.session),
 80		layout.WithPadding(1, 1, 1, 1),
 81	)
 82	p.layout.SetRightPanel(sidebarContainer)
 83	width, height := p.layout.GetSize()
 84	p.layout.SetSize(width, height)
 85	return sidebarContainer.Init()
 86}
 87
 88func (p *chatPage) clearSidebar() {
 89	p.layout.SetRightPanel(nil)
 90	width, height := p.layout.GetSize()
 91	p.layout.SetSize(width, height)
 92}
 93
 94func (p *chatPage) sendMessage(text string) tea.Cmd {
 95	var cmds []tea.Cmd
 96	if p.session.ID == "" {
 97		session, err := p.app.Sessions.Create(context.Background(), "New Session")
 98		if err != nil {
 99			return util.ReportError(err)
100		}
101
102		p.session = session
103		cmd := p.setSidebar()
104		if cmd != nil {
105			cmds = append(cmds, cmd)
106		}
107		cmds = append(cmds, util.CmdHandler(chat.SessionSelectedMsg(session)))
108	}
109	// TODO: move this to a service
110	a, err := agent.NewCoderAgent(p.app)
111	if err != nil {
112		return util.ReportError(err)
113	}
114	go func() {
115		a.Generate(context.Background(), p.session.ID, text)
116	}()
117
118	return tea.Batch(cmds...)
119}
120
121func (p *chatPage) View() string {
122	return p.layout.View()
123}
124
125func NewChatPage(app *app.App) tea.Model {
126	messagesContainer := layout.NewContainer(
127		chat.NewMessagesCmp(app),
128		layout.WithPadding(1, 1, 0, 1),
129	)
130
131	editorContainer := layout.NewContainer(
132		chat.NewEditorCmp(),
133		layout.WithBorder(true, false, false, false),
134	)
135	return &chatPage{
136		app: app,
137		layout: layout.NewSplitPane(
138			layout.WithLeftPanel(messagesContainer),
139			layout.WithBottomPanel(editorContainer),
140		),
141	}
142}