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