1package repl
2
3import (
4 "fmt"
5 "slices"
6 "strings"
7
8 "github.com/charmbracelet/bubbles/key"
9 "github.com/charmbracelet/bubbles/viewport"
10 tea "github.com/charmbracelet/bubbletea"
11 "github.com/charmbracelet/glamour"
12 "github.com/charmbracelet/lipgloss"
13 "github.com/cloudwego/eino/schema"
14 "github.com/kujtimiihoxha/termai/internal/app"
15 "github.com/kujtimiihoxha/termai/internal/message"
16 "github.com/kujtimiihoxha/termai/internal/pubsub"
17 "github.com/kujtimiihoxha/termai/internal/session"
18 "github.com/kujtimiihoxha/termai/internal/tui/layout"
19 "github.com/kujtimiihoxha/termai/internal/tui/styles"
20)
21
22type MessagesCmp interface {
23 tea.Model
24 layout.Focusable
25 layout.Bordered
26 layout.Sizeable
27 layout.Bindings
28}
29
30type messagesCmp struct {
31 app *app.App
32 messages []message.Message
33 session session.Session
34 viewport viewport.Model
35 mdRenderer *glamour.TermRenderer
36 width int
37 height int
38 focused bool
39 cachedView string
40}
41
42func (m *messagesCmp) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
43 switch msg := msg.(type) {
44 case pubsub.Event[message.Message]:
45 if msg.Type == pubsub.CreatedEvent {
46 m.messages = append(m.messages, msg.Payload)
47 m.renderView()
48 m.viewport.GotoBottom()
49 }
50 case pubsub.Event[session.Session]:
51 if msg.Type == pubsub.UpdatedEvent {
52 if m.session.ID == msg.Payload.ID {
53 m.session = msg.Payload
54 }
55 }
56 case SelectedSessionMsg:
57 m.session, _ = m.app.Sessions.Get(msg.SessionID)
58 m.messages, _ = m.app.Messages.List(m.session.ID)
59 m.renderView()
60 m.viewport.GotoBottom()
61 }
62 if m.focused {
63 u, cmd := m.viewport.Update(msg)
64 m.viewport = u
65 return m, cmd
66 }
67 return m, nil
68}
69
70func borderColor(role schema.RoleType) lipgloss.TerminalColor {
71 switch role {
72 case schema.Assistant:
73 return styles.Mauve
74 case schema.User:
75 return styles.Rosewater
76 case schema.Tool:
77 return styles.Peach
78 }
79 return styles.Blue
80}
81
82func borderText(msgRole schema.RoleType, currentMessage int) map[layout.BorderPosition]string {
83 role := ""
84 icon := ""
85 switch msgRole {
86 case schema.Assistant:
87 role = "Assistant"
88 icon = styles.BotIcon
89 case schema.User:
90 role = "User"
91 icon = styles.UserIcon
92 }
93 return map[layout.BorderPosition]string{
94 layout.TopLeftBorder: lipgloss.NewStyle().
95 Padding(0, 1).
96 Bold(true).
97 Foreground(styles.Crust).
98 Background(borderColor(msgRole)).
99 Render(fmt.Sprintf("%s %s ", role, icon)),
100 layout.TopRightBorder: lipgloss.NewStyle().
101 Padding(0, 1).
102 Bold(true).
103 Foreground(styles.Crust).
104 Background(borderColor(msgRole)).
105 Render(fmt.Sprintf("#%d ", currentMessage)),
106 }
107}
108
109func (m *messagesCmp) renderView() {
110 stringMessages := make([]string, 0)
111 r, _ := glamour.NewTermRenderer(
112 glamour.WithStyles(styles.CatppuccinMarkdownStyle()),
113 glamour.WithWordWrap(m.width-10),
114 glamour.WithEmoji(),
115 )
116 textStyle := lipgloss.NewStyle().Width(m.width - 4)
117 currentMessage := 1
118 for _, msg := range m.messages {
119 if msg.MessageData.Role == schema.Tool {
120 continue
121 }
122 content := msg.MessageData.Content
123 if content != "" {
124 content, _ = r.Render(msg.MessageData.Content)
125 stringMessages = append(stringMessages, layout.Borderize(
126 textStyle.Render(content),
127 layout.BorderOptions{
128 InactiveBorder: lipgloss.DoubleBorder(),
129 ActiveBorder: lipgloss.DoubleBorder(),
130 ActiveColor: borderColor(msg.MessageData.Role),
131 InactiveColor: borderColor(msg.MessageData.Role),
132 EmbeddedText: borderText(msg.MessageData.Role, currentMessage),
133 },
134 ))
135 currentMessage++
136 }
137 for _, toolCall := range msg.MessageData.ToolCalls {
138 resultInx := slices.IndexFunc(m.messages, func(m message.Message) bool {
139 return m.MessageData.ToolCallID == toolCall.ID
140 })
141 content := fmt.Sprintf("**Arguments**\n```json\n%s\n```\n", toolCall.Function.Arguments)
142 if resultInx == -1 {
143 content += "Running..."
144 } else {
145 result := m.messages[resultInx].MessageData.Content
146 if result != "" {
147 lines := strings.Split(result, "\n")
148 if len(lines) > 15 {
149 result = strings.Join(lines[:15], "\n")
150 }
151 content += fmt.Sprintf("**Result**\n```\n%s\n```\n", result)
152 if len(lines) > 15 {
153 content += fmt.Sprintf("\n\n *...%d lines are truncated* ", len(lines)-15)
154 }
155 }
156 }
157 content, _ = r.Render(content)
158 stringMessages = append(stringMessages, layout.Borderize(
159 textStyle.Render(content),
160 layout.BorderOptions{
161 InactiveBorder: lipgloss.DoubleBorder(),
162 ActiveBorder: lipgloss.DoubleBorder(),
163 ActiveColor: borderColor(schema.Tool),
164 InactiveColor: borderColor(schema.Tool),
165 EmbeddedText: map[layout.BorderPosition]string{
166 layout.TopLeftBorder: lipgloss.NewStyle().
167 Padding(0, 1).
168 Bold(true).
169 Foreground(styles.Crust).
170 Background(borderColor(schema.Tool)).
171 Render(
172 fmt.Sprintf("Tool [%s] %s ", toolCall.Function.Name, styles.ToolIcon),
173 ),
174 layout.TopRightBorder: lipgloss.NewStyle().
175 Padding(0, 1).
176 Bold(true).
177 Foreground(styles.Crust).
178 Background(borderColor(schema.Tool)).
179 Render(fmt.Sprintf("#%d ", currentMessage)),
180 },
181 },
182 ))
183 currentMessage++
184 }
185 }
186 m.viewport.SetContent(lipgloss.JoinVertical(lipgloss.Top, stringMessages...))
187}
188
189func (m *messagesCmp) View() string {
190 return lipgloss.NewStyle().Padding(1).Render(m.viewport.View())
191}
192
193func (m *messagesCmp) BindingKeys() []key.Binding {
194 return layout.KeyMapToSlice(m.viewport.KeyMap)
195}
196
197func (m *messagesCmp) Blur() tea.Cmd {
198 m.focused = false
199 return nil
200}
201
202func (m *messagesCmp) BorderText() map[layout.BorderPosition]string {
203 title := m.session.Title
204 titleWidth := m.width / 2
205 if len(title) > titleWidth {
206 title = title[:titleWidth] + "..."
207 }
208 if m.focused {
209 title = lipgloss.NewStyle().Foreground(styles.Primary).Render(title)
210 }
211 return map[layout.BorderPosition]string{
212 layout.TopLeftBorder: title,
213 layout.BottomRightBorder: formatTokensAndCost(m.session.CompletionTokens+m.session.PromptTokens, m.session.Cost),
214 }
215}
216
217func (m *messagesCmp) Focus() tea.Cmd {
218 m.focused = true
219 return nil
220}
221
222func (m *messagesCmp) GetSize() (int, int) {
223 return m.width, m.height
224}
225
226func (m *messagesCmp) IsFocused() bool {
227 return m.focused
228}
229
230func (m *messagesCmp) SetSize(width int, height int) {
231 m.width = width
232 m.height = height
233 m.viewport.Width = width - 2 // padding
234 m.viewport.Height = height - 2 // padding
235}
236
237func (m *messagesCmp) Init() tea.Cmd {
238 return nil
239}
240
241func NewMessagesCmp(app *app.App) MessagesCmp {
242 return &messagesCmp{
243 app: app,
244 messages: []message.Message{},
245 viewport: viewport.New(0, 0),
246 }
247}