sessions.go

  1package repl
  2
  3import (
  4	"fmt"
  5	"strings"
  6
  7	"github.com/charmbracelet/bubbles/key"
  8	"github.com/charmbracelet/bubbles/list"
  9	tea "github.com/charmbracelet/bubbletea"
 10	"github.com/charmbracelet/lipgloss"
 11	"github.com/kujtimiihoxha/termai/internal/app"
 12	"github.com/kujtimiihoxha/termai/internal/pubsub"
 13	"github.com/kujtimiihoxha/termai/internal/session"
 14	"github.com/kujtimiihoxha/termai/internal/tui/layout"
 15	"github.com/kujtimiihoxha/termai/internal/tui/styles"
 16	"github.com/kujtimiihoxha/termai/internal/tui/util"
 17)
 18
 19type SessionsCmp interface {
 20	tea.Model
 21	layout.Sizeable
 22	layout.Focusable
 23	layout.Bordered
 24	layout.Bindings
 25}
 26type sessionsCmp struct {
 27	app     *app.App
 28	list    list.Model
 29	focused bool
 30}
 31
 32type listItem struct {
 33	id, title, desc string
 34}
 35
 36func (i listItem) Title() string       { return i.title }
 37func (i listItem) Description() string { return i.desc }
 38func (i listItem) FilterValue() string { return i.title }
 39
 40type InsertSessionsMsg struct {
 41	sessions []session.Session
 42}
 43
 44type SelectedSessionMsg struct {
 45	SessionID string
 46}
 47
 48type sessionsKeyMap struct {
 49	Select key.Binding
 50}
 51
 52var sessionKeyMapValue = sessionsKeyMap{
 53	Select: key.NewBinding(
 54		key.WithKeys("enter", " "),
 55		key.WithHelp("enter/space", "select session"),
 56	),
 57}
 58
 59func (i *sessionsCmp) Init() tea.Cmd {
 60	existing, err := i.app.Sessions.List()
 61	if err != nil {
 62		return util.ReportError(err)
 63	}
 64	if len(existing) == 0 || existing[0].MessageCount > 0 {
 65		newSession, err := i.app.Sessions.Create(
 66			"New Session",
 67		)
 68		if err != nil {
 69			return util.ReportError(err)
 70		}
 71		existing = append([]session.Session{newSession}, existing...)
 72	}
 73	return tea.Batch(
 74		util.CmdHandler(InsertSessionsMsg{existing}),
 75		util.CmdHandler(SelectedSessionMsg{existing[0].ID}),
 76	)
 77}
 78
 79func (i *sessionsCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 80	switch msg := msg.(type) {
 81	case InsertSessionsMsg:
 82		items := make([]list.Item, len(msg.sessions))
 83		for i, s := range msg.sessions {
 84			items[i] = listItem{
 85				id:    s.ID,
 86				title: s.Title,
 87				desc:  formatTokensAndCost(s.PromptTokens+s.CompletionTokens, s.Cost),
 88			}
 89		}
 90		return i, i.list.SetItems(items)
 91	case pubsub.Event[session.Session]:
 92		if msg.Type == pubsub.UpdatedEvent {
 93			// update the session in the list
 94			items := i.list.Items()
 95			for idx, item := range items {
 96				s := item.(listItem)
 97				if s.id == msg.Payload.ID {
 98					s.title = msg.Payload.Title
 99					s.desc = formatTokensAndCost(msg.Payload.PromptTokens+msg.Payload.CompletionTokens, msg.Payload.Cost)
100					items[idx] = s
101					break
102				}
103			}
104			return i, i.list.SetItems(items)
105		}
106
107	case tea.KeyMsg:
108		switch {
109		case key.Matches(msg, sessionKeyMapValue.Select):
110			selected := i.list.SelectedItem()
111			if selected == nil {
112				return i, nil
113			}
114			return i, util.CmdHandler(SelectedSessionMsg{selected.(listItem).id})
115		}
116	}
117	if i.focused {
118		u, cmd := i.list.Update(msg)
119		i.list = u
120		return i, cmd
121	}
122	return i, nil
123}
124
125func (i *sessionsCmp) View() string {
126	return i.list.View()
127}
128
129func (i *sessionsCmp) Blur() tea.Cmd {
130	i.focused = false
131	return nil
132}
133
134func (i *sessionsCmp) Focus() tea.Cmd {
135	i.focused = true
136	return nil
137}
138
139func (i *sessionsCmp) GetSize() (int, int) {
140	return i.list.Width(), i.list.Height()
141}
142
143func (i *sessionsCmp) IsFocused() bool {
144	return i.focused
145}
146
147func (i *sessionsCmp) SetSize(width int, height int) {
148	i.list.SetSize(width, height)
149}
150
151func (i *sessionsCmp) BorderText() map[layout.BorderPosition]string {
152	totalCount := len(i.list.Items())
153	itemsPerPage := i.list.Paginator.PerPage
154	currentPage := i.list.Paginator.Page
155
156	current := min(currentPage*itemsPerPage+itemsPerPage, totalCount)
157
158	pageInfo := fmt.Sprintf(
159		"%d-%d of %d",
160		currentPage*itemsPerPage+1,
161		current,
162		totalCount,
163	)
164
165	title := "Sessions"
166	if i.focused {
167		title = lipgloss.NewStyle().Foreground(styles.Primary).Render(title)
168	}
169	return map[layout.BorderPosition]string{
170		layout.TopMiddleBorder:    title,
171		layout.BottomMiddleBorder: pageInfo,
172	}
173}
174
175func (i *sessionsCmp) BindingKeys() []key.Binding {
176	return append(layout.KeyMapToSlice(i.list.KeyMap), sessionKeyMapValue.Select)
177}
178
179func formatTokensAndCost(tokens int64, cost float64) string {
180	// Format tokens in human-readable format (e.g., 110K, 1.2M)
181	var formattedTokens string
182	switch {
183	case tokens >= 1_000_000:
184		formattedTokens = fmt.Sprintf("%.1fM", float64(tokens)/1_000_000)
185	case tokens >= 1_000:
186		formattedTokens = fmt.Sprintf("%.1fK", float64(tokens)/1_000)
187	default:
188		formattedTokens = fmt.Sprintf("%d", tokens)
189	}
190
191	// Remove .0 suffix if present
192	if strings.HasSuffix(formattedTokens, ".0K") {
193		formattedTokens = strings.Replace(formattedTokens, ".0K", "K", 1)
194	}
195	if strings.HasSuffix(formattedTokens, ".0M") {
196		formattedTokens = strings.Replace(formattedTokens, ".0M", "M", 1)
197	}
198
199	// Format cost with $ symbol and 2 decimal places
200	formattedCost := fmt.Sprintf("$%.2f", cost)
201
202	return fmt.Sprintf("Tokens: %s, Cost: %s", formattedTokens, formattedCost)
203}
204
205func NewSessionsCmp(app *app.App) SessionsCmp {
206	listDelegate := list.NewDefaultDelegate()
207	defaultItemStyle := list.NewDefaultItemStyles()
208	defaultItemStyle.SelectedTitle = defaultItemStyle.SelectedTitle.BorderForeground(styles.Secondary).Foreground(styles.Primary)
209	defaultItemStyle.SelectedDesc = defaultItemStyle.SelectedDesc.BorderForeground(styles.Secondary).Foreground(styles.Primary)
210
211	defaultStyle := list.DefaultStyles()
212	defaultStyle.FilterPrompt = defaultStyle.FilterPrompt.Foreground(styles.Secondary)
213	defaultStyle.FilterCursor = defaultStyle.FilterCursor.Foreground(styles.Flamingo)
214
215	listDelegate.Styles = defaultItemStyle
216
217	listComponent := list.New([]list.Item{}, listDelegate, 0, 0)
218	listComponent.FilterInput.PromptStyle = defaultStyle.FilterPrompt
219	listComponent.FilterInput.Cursor.Style = defaultStyle.FilterCursor
220	listComponent.SetShowTitle(false)
221	listComponent.SetShowPagination(false)
222	listComponent.SetShowHelp(false)
223	listComponent.SetShowStatusBar(false)
224	listComponent.DisableQuitKeybindings()
225
226	return &sessionsCmp{
227		app:     app,
228		list:    listComponent,
229		focused: false,
230	}
231}