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.CreatedEvent && msg.Payload.ParentSessionID == "" {
 93			// Check if the session is already in the list
 94			items := i.list.Items()
 95			for _, item := range items {
 96				s := item.(listItem)
 97				if s.id == msg.Payload.ID {
 98					return i, nil
 99				}
100			}
101			// insert the new session at the top of the list
102			items = append([]list.Item{listItem{
103				id:    msg.Payload.ID,
104				title: msg.Payload.Title,
105				desc:  formatTokensAndCost(msg.Payload.PromptTokens+msg.Payload.CompletionTokens, msg.Payload.Cost),
106			}}, items...)
107			return i, i.list.SetItems(items)
108		} else if msg.Type == pubsub.UpdatedEvent {
109			// update the session in the list
110			items := i.list.Items()
111			for idx, item := range items {
112				s := item.(listItem)
113				if s.id == msg.Payload.ID {
114					s.title = msg.Payload.Title
115					s.desc = formatTokensAndCost(msg.Payload.PromptTokens+msg.Payload.CompletionTokens, msg.Payload.Cost)
116					items[idx] = s
117					break
118				}
119			}
120			return i, i.list.SetItems(items)
121		}
122
123	case tea.KeyMsg:
124		switch {
125		case key.Matches(msg, sessionKeyMapValue.Select):
126			selected := i.list.SelectedItem()
127			if selected == nil {
128				return i, nil
129			}
130			return i, util.CmdHandler(SelectedSessionMsg{selected.(listItem).id})
131		}
132	}
133	if i.focused {
134		u, cmd := i.list.Update(msg)
135		i.list = u
136		return i, cmd
137	}
138	return i, nil
139}
140
141func (i *sessionsCmp) View() string {
142	return i.list.View()
143}
144
145func (i *sessionsCmp) Blur() tea.Cmd {
146	i.focused = false
147	return nil
148}
149
150func (i *sessionsCmp) Focus() tea.Cmd {
151	i.focused = true
152	return nil
153}
154
155func (i *sessionsCmp) GetSize() (int, int) {
156	return i.list.Width(), i.list.Height()
157}
158
159func (i *sessionsCmp) IsFocused() bool {
160	return i.focused
161}
162
163func (i *sessionsCmp) SetSize(width int, height int) {
164	i.list.SetSize(width, height)
165}
166
167func (i *sessionsCmp) BorderText() map[layout.BorderPosition]string {
168	totalCount := len(i.list.Items())
169	itemsPerPage := i.list.Paginator.PerPage
170	currentPage := i.list.Paginator.Page
171
172	current := min(currentPage*itemsPerPage+itemsPerPage, totalCount)
173
174	pageInfo := fmt.Sprintf(
175		"%d-%d of %d",
176		currentPage*itemsPerPage+1,
177		current,
178		totalCount,
179	)
180
181	title := "Sessions"
182	if i.focused {
183		title = lipgloss.NewStyle().Foreground(styles.Primary).Render(title)
184	}
185	return map[layout.BorderPosition]string{
186		layout.TopMiddleBorder:    title,
187		layout.BottomMiddleBorder: pageInfo,
188	}
189}
190
191func (i *sessionsCmp) BindingKeys() []key.Binding {
192	return append(layout.KeyMapToSlice(i.list.KeyMap), sessionKeyMapValue.Select)
193}
194
195func formatTokensAndCost(tokens int64, cost float64) string {
196	// Format tokens in human-readable format (e.g., 110K, 1.2M)
197	var formattedTokens string
198	switch {
199	case tokens >= 1_000_000:
200		formattedTokens = fmt.Sprintf("%.1fM", float64(tokens)/1_000_000)
201	case tokens >= 1_000:
202		formattedTokens = fmt.Sprintf("%.1fK", float64(tokens)/1_000)
203	default:
204		formattedTokens = fmt.Sprintf("%d", tokens)
205	}
206
207	// Remove .0 suffix if present
208	if strings.HasSuffix(formattedTokens, ".0K") {
209		formattedTokens = strings.Replace(formattedTokens, ".0K", "K", 1)
210	}
211	if strings.HasSuffix(formattedTokens, ".0M") {
212		formattedTokens = strings.Replace(formattedTokens, ".0M", "M", 1)
213	}
214
215	// Format cost with $ symbol and 2 decimal places
216	formattedCost := fmt.Sprintf("$%.2f", cost)
217
218	return fmt.Sprintf("Tokens: %s, Cost: %s", formattedTokens, formattedCost)
219}
220
221func NewSessionsCmp(app *app.App) SessionsCmp {
222	listDelegate := list.NewDefaultDelegate()
223	defaultItemStyle := list.NewDefaultItemStyles()
224	defaultItemStyle.SelectedTitle = defaultItemStyle.SelectedTitle.BorderForeground(styles.Secondary).Foreground(styles.Primary)
225	defaultItemStyle.SelectedDesc = defaultItemStyle.SelectedDesc.BorderForeground(styles.Secondary).Foreground(styles.Primary)
226
227	defaultStyle := list.DefaultStyles()
228	defaultStyle.FilterPrompt = defaultStyle.FilterPrompt.Foreground(styles.Secondary)
229	defaultStyle.FilterCursor = defaultStyle.FilterCursor.Foreground(styles.Flamingo)
230
231	listDelegate.Styles = defaultItemStyle
232
233	listComponent := list.New([]list.Item{}, listDelegate, 0, 0)
234	listComponent.FilterInput.PromptStyle = defaultStyle.FilterPrompt
235	listComponent.FilterInput.Cursor.Style = defaultStyle.FilterCursor
236	listComponent.SetShowTitle(false)
237	listComponent.SetShowPagination(false)
238	listComponent.SetShowHelp(false)
239	listComponent.SetShowStatusBar(false)
240	listComponent.DisableQuitKeybindings()
241
242	return &sessionsCmp{
243		app:     app,
244		list:    listComponent,
245		focused: false,
246	}
247}