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