1package chat
2
3import (
4 "fmt"
5 "image"
6 "strings"
7 "time"
8
9 tea "charm.land/bubbletea/v2"
10 "charm.land/lipgloss/v2"
11 "github.com/charmbracelet/catwalk/pkg/catwalk"
12 "github.com/charmbracelet/crush/internal/config"
13 "github.com/charmbracelet/crush/internal/message"
14 "github.com/charmbracelet/crush/internal/ui/anim"
15 "github.com/charmbracelet/crush/internal/ui/attachments"
16 "github.com/charmbracelet/crush/internal/ui/common"
17 "github.com/charmbracelet/crush/internal/ui/list"
18 "github.com/charmbracelet/crush/internal/ui/styles"
19)
20
21// this is the total width that is taken up by the border + padding
22// we also cap the width so text is readable to the maxTextWidth(120)
23const messageLeftPaddingTotal = 2
24
25// maxTextWidth is the maximum width text messages can be
26const maxTextWidth = 120
27
28// Identifiable is an interface for items that can provide a unique identifier.
29type Identifiable interface {
30 ID() string
31}
32
33// Animatable is an interface for items that support animation.
34type Animatable interface {
35 StartAnimation() tea.Cmd
36 Animate(msg anim.StepMsg) tea.Cmd
37}
38
39// Expandable is an interface for items that can be expanded or collapsed.
40type Expandable interface {
41 ToggleExpanded()
42}
43
44// KeyEventHandler is an interface for items that can handle key events.
45type KeyEventHandler interface {
46 HandleKeyEvent(key tea.KeyMsg) (bool, tea.Cmd)
47}
48
49// MessageItem represents a [message.Message] item that can be displayed in the
50// UI and be part of a [list.List] identifiable by a unique ID.
51type MessageItem interface {
52 list.Item
53 list.RawRenderable
54 Identifiable
55}
56
57// HighlightableMessageItem is a message item that supports highlighting.
58type HighlightableMessageItem interface {
59 MessageItem
60 list.Highlightable
61}
62
63// FocusableMessageItem is a message item that supports focus.
64type FocusableMessageItem interface {
65 MessageItem
66 list.Focusable
67}
68
69// SendMsg represents a message to send a chat message.
70type SendMsg struct {
71 Text string
72 Attachments []message.Attachment
73}
74
75type highlightableMessageItem struct {
76 startLine int
77 startCol int
78 endLine int
79 endCol int
80 highlighter list.Highlighter
81}
82
83var _ list.Highlightable = (*highlightableMessageItem)(nil)
84
85// isHighlighted returns true if the item has a highlight range set.
86func (h *highlightableMessageItem) isHighlighted() bool {
87 return h.startLine != -1 || h.endLine != -1
88}
89
90// renderHighlighted highlights the content if necessary.
91func (h *highlightableMessageItem) renderHighlighted(content string, width, height int) string {
92 if !h.isHighlighted() {
93 return content
94 }
95 area := image.Rect(0, 0, width, height)
96 return list.Highlight(content, area, h.startLine, h.startCol, h.endLine, h.endCol, h.highlighter)
97}
98
99// SetHighlight implements list.Highlightable.
100func (h *highlightableMessageItem) SetHighlight(startLine int, startCol int, endLine int, endCol int) {
101 // Adjust columns for the style's left inset (border + padding) since we
102 // highlight the content only.
103 offset := messageLeftPaddingTotal
104 h.startLine = startLine
105 h.startCol = max(0, startCol-offset)
106 h.endLine = endLine
107 if endCol >= 0 {
108 h.endCol = max(0, endCol-offset)
109 } else {
110 h.endCol = endCol
111 }
112}
113
114// Highlight implements list.Highlightable.
115func (h *highlightableMessageItem) Highlight() (startLine int, startCol int, endLine int, endCol int) {
116 return h.startLine, h.startCol, h.endLine, h.endCol
117}
118
119func defaultHighlighter(sty *styles.Styles) *highlightableMessageItem {
120 return &highlightableMessageItem{
121 startLine: -1,
122 startCol: -1,
123 endLine: -1,
124 endCol: -1,
125 highlighter: list.ToHighlighter(sty.TextSelection),
126 }
127}
128
129// cachedMessageItem caches rendered message content to avoid re-rendering.
130//
131// This should be used by any message that can store a cached version of its render. e.x user,assistant... and so on
132//
133// THOUGHT(kujtim): we should consider if its efficient to store the render for different widths
134// the issue with that could be memory usage
135type cachedMessageItem struct {
136 // rendered is the cached rendered string
137 rendered string
138 // width and height are the dimensions of the cached render
139 width int
140 height int
141}
142
143// getCachedRender returns the cached render if it exists for the given width.
144func (c *cachedMessageItem) getCachedRender(width int) (string, int, bool) {
145 if c.width == width && c.rendered != "" {
146 return c.rendered, c.height, true
147 }
148 return "", 0, false
149}
150
151// setCachedRender sets the cached render.
152func (c *cachedMessageItem) setCachedRender(rendered string, width, height int) {
153 c.rendered = rendered
154 c.width = width
155 c.height = height
156}
157
158// clearCache clears the cached render.
159func (c *cachedMessageItem) clearCache() {
160 c.rendered = ""
161 c.width = 0
162 c.height = 0
163}
164
165// focusableMessageItem is a base struct for message items that can be focused.
166type focusableMessageItem struct {
167 focused bool
168}
169
170// SetFocused implements MessageItem.
171func (f *focusableMessageItem) SetFocused(focused bool) {
172 f.focused = focused
173}
174
175// AssistantInfoID returns a stable ID for assistant info items.
176func AssistantInfoID(messageID string) string {
177 return fmt.Sprintf("%s:assistant-info", messageID)
178}
179
180// AssistantInfoItem renders model info and response time after assistant completes.
181type AssistantInfoItem struct {
182 *cachedMessageItem
183
184 id string
185 message *message.Message
186 sty *styles.Styles
187 lastUserMessageTime time.Time
188}
189
190// NewAssistantInfoItem creates a new AssistantInfoItem.
191func NewAssistantInfoItem(sty *styles.Styles, message *message.Message, lastUserMessageTime time.Time) MessageItem {
192 return &AssistantInfoItem{
193 cachedMessageItem: &cachedMessageItem{},
194 id: AssistantInfoID(message.ID),
195 message: message,
196 sty: sty,
197 lastUserMessageTime: lastUserMessageTime,
198 }
199}
200
201// ID implements MessageItem.
202func (a *AssistantInfoItem) ID() string {
203 return a.id
204}
205
206// RawRender implements MessageItem.
207func (a *AssistantInfoItem) RawRender(width int) string {
208 innerWidth := max(0, width-messageLeftPaddingTotal)
209 content, _, ok := a.getCachedRender(innerWidth)
210 if !ok {
211 content = a.renderContent(innerWidth)
212 height := lipgloss.Height(content)
213 a.setCachedRender(content, innerWidth, height)
214 }
215 return content
216}
217
218// Render implements MessageItem.
219func (a *AssistantInfoItem) Render(width int) string {
220 return a.sty.Chat.Message.SectionHeader.Render(a.RawRender(width))
221}
222
223func (a *AssistantInfoItem) renderContent(width int) string {
224 finishData := a.message.FinishPart()
225 if finishData == nil {
226 return ""
227 }
228 finishTime := time.Unix(finishData.Time, 0)
229 duration := finishTime.Sub(a.lastUserMessageTime)
230 infoMsg := a.sty.Chat.Message.AssistantInfoDuration.Render(duration.String())
231 icon := a.sty.Chat.Message.AssistantInfoIcon.Render(styles.ModelIcon)
232 model := config.Get().GetModel(a.message.Provider, a.message.Model)
233 if model == nil {
234 model = &catwalk.Model{Name: "Unknown Model"}
235 }
236 modelFormatted := a.sty.Chat.Message.AssistantInfoModel.Render(model.Name)
237 providerName := a.message.Provider
238 if providerConfig, ok := config.Get().Providers.Get(a.message.Provider); ok {
239 providerName = providerConfig.Name
240 }
241 provider := a.sty.Chat.Message.AssistantInfoProvider.Render(fmt.Sprintf("via %s", providerName))
242 assistant := fmt.Sprintf("%s %s %s %s", icon, modelFormatted, provider, infoMsg)
243 return common.Section(a.sty, assistant, width)
244}
245
246// cappedMessageWidth returns the maximum width for message content for readability.
247func cappedMessageWidth(availableWidth int) int {
248 return min(availableWidth-messageLeftPaddingTotal, maxTextWidth)
249}
250
251// ExtractMessageItems extracts [MessageItem]s from a [message.Message]. It
252// returns all parts of the message as [MessageItem]s.
253//
254// For assistant messages with tool calls, pass a toolResults map to link results.
255// Use BuildToolResultMap to create this map from all messages in a session.
256func ExtractMessageItems(sty *styles.Styles, msg *message.Message, toolResults map[string]message.ToolResult) []MessageItem {
257 switch msg.Role {
258 case message.User:
259 r := attachments.NewRenderer(
260 sty.Attachments.Normal,
261 sty.Attachments.Deleting,
262 sty.Attachments.Image,
263 sty.Attachments.Text,
264 )
265 return []MessageItem{NewUserMessageItem(sty, msg, r)}
266 case message.Assistant:
267 var items []MessageItem
268 if ShouldRenderAssistantMessage(msg) {
269 items = append(items, NewAssistantMessageItem(sty, msg))
270 }
271 for _, tc := range msg.ToolCalls() {
272 var result *message.ToolResult
273 if tr, ok := toolResults[tc.ID]; ok {
274 result = &tr
275 }
276 items = append(items, NewToolMessageItem(
277 sty,
278 msg.ID,
279 tc,
280 result,
281 msg.FinishReason() == message.FinishReasonCanceled,
282 ))
283 }
284 return items
285 }
286 return []MessageItem{}
287}
288
289// ShouldRenderAssistantMessage determines if an assistant message should be rendered
290//
291// In some cases the assistant message only has tools so we do not want to render an
292// empty message.
293func ShouldRenderAssistantMessage(msg *message.Message) bool {
294 content := strings.TrimSpace(msg.Content().Text)
295 thinking := strings.TrimSpace(msg.ReasoningContent().Thinking)
296 isError := msg.FinishReason() == message.FinishReasonError
297 isCancelled := msg.FinishReason() == message.FinishReasonCanceled
298 hasToolCalls := len(msg.ToolCalls()) > 0
299 return !hasToolCalls || content != "" || thinking != "" || msg.IsThinking() || isError || isCancelled
300}
301
302// BuildToolResultMap creates a map of tool call IDs to their results from a list of messages.
303// Tool result messages (role == message.Tool) contain the results that should be linked
304// to tool calls in assistant messages.
305func BuildToolResultMap(messages []*message.Message) map[string]message.ToolResult {
306 resultMap := make(map[string]message.ToolResult)
307 for _, msg := range messages {
308 if msg.Role == message.Tool {
309 for _, result := range msg.ToolResults() {
310 if result.ToolCallID != "" {
311 resultMap[result.ToolCallID] = result
312 }
313 }
314 }
315 }
316 return resultMap
317}