0decf78 · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

refactor: add assistant info item (#1881)

0decf78Kujtim Hoxha9 months ago

internal/ui/AGENTS.md+2

compare with main

unifiedsplit

@@ -4,6 +4,8 @@

4- Never use commands to send messages when you can directly mutate children or state.

5- Keep things simple; do not overcomplicate.

6- Create files if needed to separate logic; do not nest models.

7+- Always do IO in commands

8+- Never change the model state inside of a command use messages and than update the state in the main loop

8## Architecture

10## Architecture

internal/ui/chat/messages.go+84−1

@@ -4,13 +4,19 @@

4package chat

6import (

7+ "fmt"

7 "image"

8 "image"

8 "strings"

9 "strings"

10+ "time"

10 tea "charm.land/bubbletea/v2"

12 tea "charm.land/bubbletea/v2"

13+ "charm.land/lipgloss/v2"

14+ "github.com/charmbracelet/catwalk/pkg/catwalk"

15+ "github.com/charmbracelet/crush/internal/config"

11 "github.com/charmbracelet/crush/internal/message"

16 "github.com/charmbracelet/crush/internal/message"

12 "github.com/charmbracelet/crush/internal/ui/anim"

17 "github.com/charmbracelet/crush/internal/ui/anim"

13 "github.com/charmbracelet/crush/internal/ui/attachments"

18 "github.com/charmbracelet/crush/internal/ui/attachments"

19+ "github.com/charmbracelet/crush/internal/ui/common"

14 "github.com/charmbracelet/crush/internal/ui/list"

20 "github.com/charmbracelet/crush/internal/ui/list"

15 "github.com/charmbracelet/crush/internal/ui/styles"

21 "github.com/charmbracelet/crush/internal/ui/styles"

16)

22)

@@ -42,9 +48,19 @@ type Expandable interface {

42// UI and be part of a [list.List] identifiable by a unique ID.

48// UI and be part of a [list.List] identifiable by a unique ID.

43type MessageItem interface {

49type MessageItem interface {

44 list.Item

50 list.Item

51+ Identifiable

52+}

53+

54+// HighlightableMessageItem is a message item that supports highlighting.

55+type HighlightableMessageItem interface {

56+ MessageItem

45 list.Highlightable

57 list.Highlightable

58+}

59+

60+// FocusableMessageItem is a message item that supports focus.

61+type FocusableMessageItem interface {

62+ MessageItem

46 list.Focusable

63 list.Focusable

47− Identifiable

48}

64}

50// SendMsg represents a message to send a chat message.

66// SendMsg represents a message to send a chat message.

@@ -146,6 +162,73 @@ func (f *focusableMessageItem) SetFocused(focused bool) {

146 f.focused = focused

162 f.focused = focused

147}

163}

148

164

165+// AssistantInfoID returns a stable ID for assistant info items.

166+func AssistantInfoID(messageID string) string {

167+ return fmt.Sprintf("%s:assistant-info", messageID)

168+}

169+

170+// AssistantInfoItem renders model info and response time after assistant completes.

171+type AssistantInfoItem struct {

172+ *cachedMessageItem

173+

174+ id string

175+ message *message.Message

176+ sty *styles.Styles

177+ lastUserMessageTime time.Time

178+}

179+

180+// NewAssistantInfoItem creates a new AssistantInfoItem.

181+func NewAssistantInfoItem(sty *styles.Styles, message *message.Message, lastUserMessageTime time.Time) MessageItem {

182+ return &AssistantInfoItem{

183+ cachedMessageItem: &cachedMessageItem{},

184+ id: AssistantInfoID(message.ID),

185+ message: message,

186+ sty: sty,

187+ lastUserMessageTime: lastUserMessageTime,

188+ }

189+}

190+

191+// ID implements MessageItem.

192+func (a *AssistantInfoItem) ID() string {

193+ return a.id

194+}

195+

196+// Render implements MessageItem.

197+func (a *AssistantInfoItem) Render(width int) string {

198+ innerWidth := max(0, width-messageLeftPaddingTotal)

199+ content, _, ok := a.getCachedRender(innerWidth)

200+ if !ok {

201+ content = a.renderContent(innerWidth)

202+ height := lipgloss.Height(content)

203+ a.setCachedRender(content, innerWidth, height)

204+ }

205+

206+ return a.sty.Chat.Message.SectionHeader.Render(content)

207+}

208+

209+func (a *AssistantInfoItem) renderContent(width int) string {

210+ finishData := a.message.FinishPart()

211+ if finishData == nil {

212+ return ""

213+ }

214+ finishTime := time.Unix(finishData.Time, 0)

215+ duration := finishTime.Sub(a.lastUserMessageTime)

216+ infoMsg := a.sty.Chat.Message.AssistantInfoDuration.Render(duration.String())

217+ icon := a.sty.Chat.Message.AssistantInfoIcon.Render(styles.ModelIcon)

218+ model := config.Get().GetModel(a.message.Provider, a.message.Model)

219+ if model == nil {

220+ model = &catwalk.Model{Name: "Unknown Model"}

221+ }

222+ modelFormatted := a.sty.Chat.Message.AssistantInfoModel.Render(model.Name)

223+ providerName := a.message.Provider

224+ if providerConfig, ok := config.Get().Providers.Get(a.message.Provider); ok {

225+ providerName = providerConfig.Name

226+ }

227+ provider := a.sty.Chat.Message.AssistantInfoProvider.Render(fmt.Sprintf("via %s", providerName))

228+ assistant := fmt.Sprintf("%s %s %s %s", icon, modelFormatted, provider, infoMsg)

229+ return common.Section(a.sty, assistant, width)

230+}

231+

149// cappedMessageWidth returns the maximum width for message content for readability.

232// cappedMessageWidth returns the maximum width for message content for readability.

150func cappedMessageWidth(availableWidth int) int {

233func cappedMessageWidth(availableWidth int) int {

151 return min(availableWidth-messageLeftPaddingTotal, maxTextWidth)

234 return min(availableWidth-messageLeftPaddingTotal, maxTextWidth)

internal/ui/common/elements.go+26−6

@@ -26,15 +26,35 @@ type ModelContextInfo struct {

26 Cost float64

27}

29−// ModelInfo renders model information including name, reasoning settings, and

29+// ModelInfo renders model information including name, provider, reasoning

30−// optional context usage/cost.

30+// settings, and optional context usage/cost.

31−func ModelInfo(t *styles.Styles, modelName string, reasoningInfo string, context *ModelContextInfo, width int) string {

31+func ModelInfo(t *styles.Styles, modelName, providerName, reasoningInfo string, context *ModelContextInfo, width int) string {

32 modelIcon := t.Subtle.Render(styles.ModelIcon)

33 modelName = t.Base.Render(modelName)

34− modelInfo := fmt.Sprintf("%s %s", modelIcon, modelName)

35+ // Build first line with model name and optionally provider on the same line

36+ var firstLine string

37+ if providerName != "" {

38+ providerInfo := t.Muted.Render(fmt.Sprintf("via %s", providerName))

39+ modelWithProvider := fmt.Sprintf("%s %s %s", modelIcon, modelName, providerInfo)

40+

41+ // Check if it fits on one line

42+ if lipgloss.Width(modelWithProvider) <= width {

43+ firstLine = modelWithProvider

44+ } else {

45+ // If it doesn't fit, put provider on next line

46+ firstLine = fmt.Sprintf("%s %s", modelIcon, modelName)

47+ }

48+ } else {

49+ firstLine = fmt.Sprintf("%s %s", modelIcon, modelName)

50+ }

51+

36− parts := []string{

52+ parts := []string{firstLine}

37− modelInfo,

53+

54+ // If provider didn't fit on first line, add it as second line

55+ if providerName != "" && !strings.Contains(firstLine, "via") {

56+ providerInfo := fmt.Sprintf("via %s", providerName)

57+ parts = append(parts, t.Muted.PaddingLeft(2).Render(providerInfo))

38 }

58 }

40 if reasoningInfo != "" {

60 if reasoningInfo != "" {

internal/ui/model/chat.go+91−6

@@ -238,39 +238,121 @@ func (m *Chat) SelectedItemInView() bool {

238 return m.list.SelectedItemInView()

239}

240

241+func (m *Chat) isSelectable(index int) bool {

242+ item := m.list.ItemAt(index)

243+ if item == nil {

244+ return false

245+ }

246+ _, ok := item.(list.Focusable)

247+ return ok

248+}

249+

241// SetSelected sets the selected message index in the chat list.

250// SetSelected sets the selected message index in the chat list.

242func (m *Chat) SetSelected(index int) {

251func (m *Chat) SetSelected(index int) {

243 m.list.SetSelected(index)

252 m.list.SetSelected(index)

253+ if index < 0 || index >= m.list.Len() {

254+ return

255+ }

256+ for {

257+ if m.isSelectable(m.list.Selected()) {

258+ return

259+ }

260+ if m.list.SelectNext() {

261+ continue

262+ }

263+ // If we're at the end and the last item isn't selectable, walk backwards

264+ // to find the nearest selectable item.

265+ for {

266+ if !m.list.SelectPrev() {

267+ return

268+ }

269+ if m.isSelectable(m.list.Selected()) {

270+ return

271+ }

272+ }

273+ }

244}

274}

245

275

246// SelectPrev selects the previous message in the chat list.

276// SelectPrev selects the previous message in the chat list.

247func (m *Chat) SelectPrev() {

277func (m *Chat) SelectPrev() {

278+ for {

248− m.list.SelectPrev()

279+ if !m.list.SelectPrev() {

280+ return

281+ }

282+ if m.isSelectable(m.list.Selected()) {

283+ return

284+ }

285+ }

249}

286}

250

287

251// SelectNext selects the next message in the chat list.

288// SelectNext selects the next message in the chat list.

252func (m *Chat) SelectNext() {

289func (m *Chat) SelectNext() {

290+ for {

253− m.list.SelectNext()

291+ if !m.list.SelectNext() {

292+ return

293+ }

294+ if m.isSelectable(m.list.Selected()) {

295+ return

296+ }

297+ }

254}

298}

255

299

256// SelectFirst selects the first message in the chat list.

300// SelectFirst selects the first message in the chat list.

257func (m *Chat) SelectFirst() {

301func (m *Chat) SelectFirst() {

258− m.list.SelectFirst()

302+ if !m.list.SelectFirst() {

303+ return

304+ }

305+ if m.isSelectable(m.list.Selected()) {

306+ return

307+ }

308+ for {

309+ if !m.list.SelectNext() {

310+ return

311+ }

312+ if m.isSelectable(m.list.Selected()) {

313+ return

314+ }

315+ }

259}

316}

260

317

261// SelectLast selects the last message in the chat list.

318// SelectLast selects the last message in the chat list.

262func (m *Chat) SelectLast() {

319func (m *Chat) SelectLast() {

263− m.list.SelectLast()

320+ if !m.list.SelectLast() {

321+ return

322+ }

323+ if m.isSelectable(m.list.Selected()) {

324+ return

325+ }

326+ for {

327+ if !m.list.SelectPrev() {

328+ return

329+ }

330+ if m.isSelectable(m.list.Selected()) {

331+ return

332+ }

333+ }

264}

334}

265

335

266// SelectFirstInView selects the first message currently in view.

336// SelectFirstInView selects the first message currently in view.

267func (m *Chat) SelectFirstInView() {

337func (m *Chat) SelectFirstInView() {

268− m.list.SelectFirstInView()

338+ startIdx, endIdx := m.list.VisibleItemIndices()

339+ for i := startIdx; i <= endIdx; i++ {

340+ if m.isSelectable(i) {

341+ m.list.SetSelected(i)

342+ return

343+ }

344+ }

269}

345}

270

346

271// SelectLastInView selects the last message currently in view.

347// SelectLastInView selects the last message currently in view.

272func (m *Chat) SelectLastInView() {

348func (m *Chat) SelectLastInView() {

273− m.list.SelectLastInView()

349+ startIdx, endIdx := m.list.VisibleItemIndices()

350+ for i := endIdx; i >= startIdx; i-- {

351+ if m.isSelectable(i) {

352+ m.list.SetSelected(i)

353+ return

354+ }

355+ }

274}

356}

275

357

276// ClearMessages removes all messages from the chat list.

358// ClearMessages removes all messages from the chat list.

@@ -335,6 +417,9 @@ func (m *Chat) HandleMouseDown(x, y int) bool {

335 if itemIdx < 0 {

417 if itemIdx < 0 {

336 return false

418 return false

337 }

419 }

420+ if !m.isSelectable(itemIdx) {

421+ return false

422+ }

338

423

339 m.mouseDown = true

424 m.mouseDown = true

340 m.mouseDownItem = itemIdx

425 m.mouseDownItem = itemIdx

internal/ui/model/sidebar.go+21−12

@@ -18,23 +18,32 @@ import (

18func (m *UI) modelInfo(width int) string {

19 model := m.selectedLargeModel()

20 reasoningInfo := ""

21+ providerName := ""

22+

21− if model != nil && model.CatwalkCfg.CanReason {

23+ if model != nil {

24+ // Get provider name first

22 providerConfig, ok := m.com.Config().Providers.Get(model.ModelCfg.Provider)

25 providerConfig, ok := m.com.Config().Providers.Get(model.ModelCfg.Provider)

23 if ok {

26 if ok {

27+ providerName = providerConfig.Name

28+

29+ // Only check reasoning if model can reason

30+ if model.CatwalkCfg.CanReason {

24− switch providerConfig.Type {

31+ switch providerConfig.Type {

25− case catwalk.TypeAnthropic:

32+ case catwalk.TypeAnthropic:

26− if model.ModelCfg.Think {

33+ if model.ModelCfg.Think {

27− reasoningInfo = "Thinking On"

34+ reasoningInfo = "Thinking On"

28− } else {

35+ } else {

29− reasoningInfo = "Thinking Off"

36+ reasoningInfo = "Thinking Off"

37+ }

38+ default:

39+ formatter := cases.Title(language.English, cases.NoLower)

40+ reasoningEffort := cmp.Or(model.ModelCfg.ReasoningEffort, model.CatwalkCfg.DefaultReasoningEffort)

41+ reasoningInfo = formatter.String(fmt.Sprintf("Reasoning %s", reasoningEffort))

30 }

42 }

31− default:

32− formatter := cases.Title(language.English, cases.NoLower)

33− reasoningEffort := cmp.Or(model.ModelCfg.ReasoningEffort, model.CatwalkCfg.DefaultReasoningEffort)

34− reasoningInfo = formatter.String(fmt.Sprintf("Reasoning %s", reasoningEffort))

35 }

43 }

36 }

44 }

37 }

45 }

46+

38 var modelContext *common.ModelContextInfo

47 var modelContext *common.ModelContextInfo

39 if m.session != nil {

48 if m.session != nil {

40 modelContext = &common.ModelContextInfo{

49 modelContext = &common.ModelContextInfo{

@@ -43,7 +52,7 @@ func (m *UI) modelInfo(width int) string {

43 ModelContext: model.CatwalkCfg.ContextWindow,

52 ModelContext: model.CatwalkCfg.ContextWindow,

44 }

53 }

45 }

54 }

46− return common.ModelInfo(m.com.Styles, model.CatwalkCfg.Name, reasoningInfo, modelContext, width)

55+ return common.ModelInfo(m.com.Styles, model.CatwalkCfg.Name, providerName, reasoningInfo, modelContext, width)

47}

56}

49// getDynamicHeightLimits will give us the num of items to show in each section based on the hight

58// getDynamicHeightLimits will give us the num of items to show in each section based on the hight

internal/ui/model/ui.go+55−3

@@ -110,6 +110,8 @@ type UI struct {

110 session *session.Session

111 sessionFiles []SessionFile

112

113+ lastUserMessageTime int64

114+

113 // The width and height of the terminal in cells.

115 // The width and height of the terminal in cells.

114 width int

116 width int

115 height int

117 height int

@@ -596,11 +598,26 @@ func (m *UI) setSessionMessages(msgs []message.Message) tea.Cmd {

596 msgPtrs[i] = &msgs[i]

598 msgPtrs[i] = &msgs[i]

597 }

599 }

598 toolResultMap := chat.BuildToolResultMap(msgPtrs)

600 toolResultMap := chat.BuildToolResultMap(msgPtrs)

601+ if len(msgPtrs) > 0 {

602+ m.lastUserMessageTime = msgPtrs[0].CreatedAt

603+ }

599

604

600 // Add messages to chat with linked tool results

605 // Add messages to chat with linked tool results

601 items := make([]chat.MessageItem, 0, len(msgs)*2)

606 items := make([]chat.MessageItem, 0, len(msgs)*2)

602 for _, msg := range msgPtrs {

607 for _, msg := range msgPtrs {

608+ switch msg.Role {

609+ case message.User:

610+ m.lastUserMessageTime = msg.CreatedAt

611+ items = append(items, chat.ExtractMessageItems(m.com.Styles, msg, toolResultMap)...)

612+ case message.Assistant:

613+ items = append(items, chat.ExtractMessageItems(m.com.Styles, msg, toolResultMap)...)

614+ if msg.FinishPart() != nil && msg.FinishPart().Reason == message.FinishReasonEndTurn {

615+ infoItem := chat.NewAssistantInfoItem(m.com.Styles, msg, time.Unix(m.lastUserMessageTime, 0))

616+ items = append(items, infoItem)

617+ }

618+ default:

603− items = append(items, chat.ExtractMessageItems(m.com.Styles, msg, toolResultMap)...)

619+ items = append(items, chat.ExtractMessageItems(m.com.Styles, msg, toolResultMap)...)

620+ }

604 }

621 }

605

622

606 // Load nested tool calls for agent/agentic_fetch tools.

623 // Load nested tool calls for agent/agentic_fetch tools.

@@ -692,7 +709,21 @@ func (m *UI) appendSessionMessage(msg message.Message) tea.Cmd {

692 return nil

709 return nil

693 }

710 }

694 switch msg.Role {

711 switch msg.Role {

695− case message.User, message.Assistant:

712+ case message.User:

713+ m.lastUserMessageTime = msg.CreatedAt

714+ items := chat.ExtractMessageItems(m.com.Styles, &msg, nil)

715+ for _, item := range items {

716+ if animatable, ok := item.(chat.Animatable); ok {

717+ if cmd := animatable.StartAnimation(); cmd != nil {

718+ cmds = append(cmds, cmd)

719+ }

720+ }

721+ }

722+ m.chat.AppendMessages(items...)

723+ if cmd := m.chat.ScrollToBottomAndAnimate(); cmd != nil {

724+ cmds = append(cmds, cmd)

725+ }

726+ case message.Assistant:

696 items := chat.ExtractMessageItems(m.com.Styles, &msg, nil)

727 items := chat.ExtractMessageItems(m.com.Styles, &msg, nil)

697 for _, item := range items {

728 for _, item := range items {

698 if animatable, ok := item.(chat.Animatable); ok {

729 if animatable, ok := item.(chat.Animatable); ok {

@@ -705,6 +736,13 @@ func (m *UI) appendSessionMessage(msg message.Message) tea.Cmd {

705 if cmd := m.chat.ScrollToBottomAndAnimate(); cmd != nil {

736 if cmd := m.chat.ScrollToBottomAndAnimate(); cmd != nil {

706 cmds = append(cmds, cmd)

737 cmds = append(cmds, cmd)

707 }

738 }

739+ if msg.FinishPart() != nil && msg.FinishPart().Reason == message.FinishReasonEndTurn {

740+ infoItem := chat.NewAssistantInfoItem(m.com.Styles, &msg, time.Unix(m.lastUserMessageTime, 0))

741+ m.chat.AppendMessages(infoItem)

742+ if cmd := m.chat.ScrollToBottomAndAnimate(); cmd != nil {

743+ cmds = append(cmds, cmd)

744+ }

745+ }

708 case message.Tool:

746 case message.Tool:

709 for _, tr := range msg.ToolResults() {

747 for _, tr := range msg.ToolResults() {

710 toolItem := m.chat.MessageItem(tr.ToolCallID)

748 toolItem := m.chat.MessageItem(tr.ToolCallID)

@@ -733,9 +771,23 @@ func (m *UI) updateSessionMessage(msg message.Message) tea.Cmd {

733 }

771 }

734 }

772 }

735

773

774+ shouldRenderAssistant := chat.ShouldRenderAssistantMessage(&msg)

736 // if the message of the assistant does not have any response just tool calls we need to remove it

775 // if the message of the assistant does not have any response just tool calls we need to remove it

737− if !chat.ShouldRenderAssistantMessage(&msg) && len(msg.ToolCalls()) > 0 && existingItem != nil {

776+ if !shouldRenderAssistant && len(msg.ToolCalls()) > 0 && existingItem != nil {

738 m.chat.RemoveMessage(msg.ID)

777 m.chat.RemoveMessage(msg.ID)

778+ if infoItem := m.chat.MessageItem(chat.AssistantInfoID(msg.ID)); infoItem != nil {

779+ m.chat.RemoveMessage(chat.AssistantInfoID(msg.ID))

780+ }

781+ }

782+

783+ if shouldRenderAssistant && msg.FinishPart() != nil && msg.FinishPart().Reason == message.FinishReasonEndTurn {

784+ if infoItem := m.chat.MessageItem(chat.AssistantInfoID(msg.ID)); infoItem == nil {

785+ newInfoItem := chat.NewAssistantInfoItem(m.com.Styles, &msg, time.Unix(m.lastUserMessageTime, 0))

786+ m.chat.AppendMessages(newInfoItem)

787+ if cmd := m.chat.ScrollToBottomAndAnimate(); cmd != nil {

788+ cmds = append(cmds, cmd)

789+ }

790+ }

739 }

791 }

740

792

741 var items []chat.MessageItem

793 var items []chat.MessageItem

internal/ui/styles/styles.go+8

@@ -238,6 +238,10 @@ type Styles struct {

238 ThinkingTruncationHint lipgloss.Style // "… (N lines hidden)" hint

239 ThinkingFooterTitle lipgloss.Style // "Thought for" text

240 ThinkingFooterDuration lipgloss.Style // Duration value

241+ AssistantInfoIcon lipgloss.Style

242+ AssistantInfoModel lipgloss.Style

243+ AssistantInfoProvider lipgloss.Style

244+ AssistantInfoDuration lipgloss.Style

241 }

245 }

242 }

246 }

243

247

@@ -1193,6 +1197,10 @@ func DefaultStyles() Styles {

1193 // No padding or border for compact tool calls within messages

1197 // No padding or border for compact tool calls within messages

1194 s.Chat.Message.ToolCallCompact = s.Muted

1198 s.Chat.Message.ToolCallCompact = s.Muted

1195 s.Chat.Message.SectionHeader = s.Base.PaddingLeft(2)

1199 s.Chat.Message.SectionHeader = s.Base.PaddingLeft(2)

1200+ s.Chat.Message.AssistantInfoIcon = s.Subtle

1201+ s.Chat.Message.AssistantInfoModel = s.Muted

1202+ s.Chat.Message.AssistantInfoProvider = s.Subtle

1203+ s.Chat.Message.AssistantInfoDuration = s.Subtle

1196

1204

1197 // Thinking section styles

1205 // Thinking section styles

1198 s.Chat.Message.ThinkingBox = s.Subtle.Background(bgBaseLighter)

1206 s.Chat.Message.ThinkingBox = s.Subtle.Background(bgBaseLighter)