tui: add tool rendering

cleaned up message rendering and changed cache semantics

other smaller tweaks

go.mod

@@ -8,9 +8,11 @@ require (
 	github.com/charmbracelet/bubbletea v0.25.0
 	github.com/charmbracelet/lipgloss v0.10.0
 	github.com/go-yaml/yaml v2.1.0+incompatible
+	github.com/muesli/reflow v0.3.0
 	github.com/sashabaranov/go-openai v1.17.7
 	github.com/spf13/cobra v1.8.0
 	github.com/sqids/sqids-go v0.4.1
+	gopkg.in/yaml.v2 v2.2.2
 	gorm.io/driver/sqlite v1.5.4
 	gorm.io/gorm v1.25.5
 )
@@ -31,7 +33,6 @@ require (
 	github.com/mattn/go-sqlite3 v1.14.18 // indirect
 	github.com/muesli/ansi v0.0.0-20211018074035-2e021307bc4b // indirect
 	github.com/muesli/cancelreader v0.2.2 // indirect
-	github.com/muesli/reflow v0.3.0 // indirect
 	github.com/muesli/termenv v0.15.2 // indirect
 	github.com/rivo/uniseg v0.4.7 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
@@ -40,5 +41,4 @@ require (
 	golang.org/x/term v0.6.0 // indirect
 	golang.org/x/text v0.3.8 // indirect
 	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
-	gopkg.in/yaml.v2 v2.2.2 // indirect
 )

pkg/lmcli/model/conversation.go

@@ -50,8 +50,8 @@ func (m *MessageRole) IsAssistant() bool {
 }
 
 // FriendlyRole returns a human friendly signifier for the message's role.
-func (m *MessageRole) FriendlyRole() string {
-	switch *m {
+func (m MessageRole) FriendlyRole() string {
+	switch m {
 	case MessageRoleUser:
 		return "You"
 	case MessageRoleSystem:
@@ -63,6 +63,6 @@ func (m *MessageRole) FriendlyRole() string {
 	case MessageRoleToolResult:
 		return "Tool Result"
 	default:
-		return string(*m)
+		return string(m)
 	}
 }

pkg/lmcli/model/tool.go

@@ -22,9 +22,9 @@ type ToolParameter struct {
 }
 
 type ToolCall struct {
-	ID         string                 `json:"id"`
-	Name       string                 `json:"name"`
-	Parameters map[string]interface{} `json:"parameters"`
+	ID         string                 `json:"id" yaml:"-"`
+	Name       string                 `json:"name" yaml:"tool"`
+	Parameters map[string]interface{} `json:"parameters" yaml:"parameters"`
 }
 
 type ToolCalls []ToolCall
@@ -51,9 +51,9 @@ func (tc ToolCalls) Value() (driver.Value, error) {
 }
 
 type ToolResult struct {
-	ToolCallID string `json:"toolCallID"`
-	ToolName   string `json:"toolName,omitempty"`
-	Result     string `json:"result,omitempty"`
+	ToolCallID string `json:"toolCallID" yaml:"-"`
+	ToolName   string `json:"toolName,omitempty" yaml:"tool"`
+	Result     string `json:"result,omitempty" yaml:"result"`
 }
 
 type ToolResults []ToolResult

pkg/tui/tui.go

@@ -10,6 +10,7 @@ package tui
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 	"strings"
 	"time"
@@ -22,7 +23,9 @@ import (
 	"github.com/charmbracelet/bubbles/viewport"
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
+	"github.com/muesli/reflow/ansi"
 	"github.com/muesli/reflow/wordwrap"
+	"gopkg.in/yaml.v2"
 )
 
 type focusState int
@@ -39,6 +42,14 @@ const (
 	selectedMessage
 )
 
+type uiCache struct {
+	header  string
+	content string
+	error   string
+	input   string
+	footer  string
+}
+
 type model struct {
 	width  int
 	height int
@@ -61,7 +72,8 @@ type model struct {
 	focus           focusState
 	wrap            bool     // whether message content is wrapped to viewport width
 	status          string   // a general status message
-	highlightCache  []string // a cache of syntax highlighted message content
+	showToolResults bool     // whether tool calls and results are shown
+	messageCache    []string // cache of syntax highlighted and wrapped message content
 	messageOffsets  []int
 	selectedMessage int
 
@@ -69,6 +81,8 @@ type model struct {
 	content viewport.Model
 	input   textarea.Model
 	spinner spinner.Model
+
+	cache *uiCache
 }
 
 type message struct {
@@ -98,14 +112,18 @@ type (
 
 // styles
 var (
-	userStyle      = lipgloss.NewStyle().Faint(true).Bold(true).Foreground(lipgloss.Color("10"))
-	assistantStyle = lipgloss.NewStyle().Faint(true).Bold(true).Foreground(lipgloss.Color("12"))
-	messageStyle   = lipgloss.NewStyle().PaddingLeft(2).PaddingRight(2)
+	headingStyle = lipgloss.NewStyle().
+			MarginTop(1).
+			MarginBottom(1).
+			PaddingLeft(1).
+			Bold(true)
+	userStyle      = lipgloss.NewStyle().Faint(true).Foreground(lipgloss.Color("10"))
+	assistantStyle = lipgloss.NewStyle().Faint(true).Foreground(lipgloss.Color("12"))
+	messageStyle   = lipgloss.NewStyle().
+			PaddingLeft(2).
+			PaddingRight(2)
 	headerStyle = lipgloss.NewStyle().
 			Background(lipgloss.Color("0"))
-	conversationStyle = lipgloss.NewStyle().
-				MarginTop(1).
-				MarginBottom(1)
 	footerStyle = lipgloss.NewStyle().
 			BorderTop(true).
 			BorderStyle(lipgloss.NormalBorder())
@@ -160,6 +178,11 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 			m.persistence = !m.persistence
 		case "ctrl+w":
 			m.wrap = !m.wrap
+			m.rebuildMessageCache()
+			m.updateContent()
+		case "ctrl+t":
+			m.showToolResults = !m.showToolResults
+			m.rebuildMessageCache()
 			m.updateContent()
 		case "q":
 			if m.focus != focusInput {
@@ -181,8 +204,8 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.width = msg.Width
 		m.height = msg.Height
 		m.content.Width = msg.Width
-		m.content.Height = msg.Height - m.getFixedComponentHeight()
 		m.input.SetWidth(msg.Width - 1)
+		m.rebuildMessageCache()
 		m.updateContent()
 	case msgConversationLoaded:
 		m.conversation = (*models.Conversation)(msg)
@@ -284,9 +307,26 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		}
 	}
 
+	if m.width > 0 {
+		m.cache.header = m.headerView()
+		m.cache.footer = m.footerView()
+		m.cache.error = m.errorView()
+		m.cache.input = m.inputView()
+		fixedHeight := height(m.cache.header) + height(m.cache.error) + height(m.cache.input) + height(m.cache.footer)
+		m.content.Height = m.height - fixedHeight
+		m.cache.content = m.contentView()
+	}
+
 	return m, tea.Batch(cmds...)
 }
 
+func height(str string) int {
+	if str == "" {
+		return 0
+	}
+	return strings.Count(str, "\n") + 1
+}
+
 func (m model) View() string {
 	if m.width == 0 {
 		// this is the case upon initial startup, but it's also a safe bet that
@@ -296,14 +336,13 @@ func (m model) View() string {
 	}
 
 	sections := make([]string, 0, 6)
-	sections = append(sections, m.headerView())
-	sections = append(sections, m.contentView())
-	error := m.errorView()
-	if error != "" {
-		sections = append(sections, error)
+	sections = append(sections, m.cache.header)
+	sections = append(sections, m.cache.content)
+	if m.cache.error != "" {
+		sections = append(sections, m.cache.error)
 	}
-	sections = append(sections, m.inputView())
-	sections = append(sections, m.footerView())
+	sections = append(sections, m.cache.input)
+	sections = append(sections, m.cache.footer)
 
 	return lipgloss.JoinVertical(
 		lipgloss.Left,
@@ -311,20 +350,6 @@ func (m model) View() string {
 	)
 }
 
-// returns the total height of "fixed" components, which are those which don't
-// change height dependent on window size.
-func (m *model) getFixedComponentHeight() int {
-	h := 0
-	h += m.input.Height()
-	h += lipgloss.Height(m.headerView())
-	h += lipgloss.Height(m.footerView())
-	errorView := m.errorView()
-	if errorView != "" {
-		h += lipgloss.Height(errorView)
-	}
-	return h
-}
-
 func (m *model) headerView() string {
 	titleStyle := lipgloss.NewStyle().
 		PaddingLeft(1).
@@ -400,9 +425,15 @@ func (m *model) footerView() string {
 	footer := left + padding + right
 	if remaining < 0 {
 		ellipses := "... "
-		// this doesn't work very well, due to trying to trim a string with
-		// ansii chars already in it
-		footer = footer[:(len(footer)+remaining)-len(ellipses)-3] + ellipses
+		for {
+			truncWidth := ansi.PrintableRuneWidth(footer) + len(ellipses)
+			if truncWidth <= m.width {
+				break
+			}
+			// truncate the minimum amount, not accounting for printed width
+			footer = footer[:len(footer)-(truncWidth-m.width)]
+		}
+		footer += ellipses
 	}
 	return footerStyle.Width(m.width).Render(footer)
 }
@@ -420,6 +451,8 @@ func initialModel(ctx *lmcli.Context, convShortname string) model {
 
 		wrap:            true,
 		selectedMessage: -1,
+
+		cache: &uiCache{},
 	}
 
 	m.content = viewport.New(0, 0)
@@ -501,7 +534,7 @@ func (m *model) handleMessagesKey(msg tea.KeyMsg) tea.Cmd {
 			return nil
 		}
 		m.messages = m.messages[:m.selectedMessage+1]
-		m.highlightCache = m.highlightCache[:m.selectedMessage+1]
+		m.messageCache = m.messageCache[:m.selectedMessage+1]
 		m.updateContent()
 		m.content.GotoBottom()
 		return m.promptLLM()
@@ -513,9 +546,13 @@ func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd {
 	switch msg.String() {
 	case "esc":
 		m.focus = focusMessages
+		if len(m.messages) > 0 {
 			if m.selectedMessage < 0 || m.selectedMessage >= len(m.messages) {
 				m.selectedMessage = len(m.messages) - 1
 			}
+			offset := m.messageOffsets[m.selectedMessage]
+			scrollIntoView(&m.content, offset, 0.1)
+		}
 		m.updateContent()
 		m.input.Blur()
 	case "ctrl+s":
@@ -714,37 +751,150 @@ func (m *model) persistConversation() tea.Cmd {
 	return nil
 }
 
+func (m *model) renderMessageHeading(i int, message *models.Message) string {
+	icon := ""
+	friendly := message.Role.FriendlyRole()
+	style := lipgloss.NewStyle().Faint(true).Bold(true)
+
+	switch message.Role {
+	case models.MessageRoleSystem:
+		icon = "⚙️"
+	case models.MessageRoleUser:
+		style = userStyle
+	case models.MessageRoleAssistant:
+		style = assistantStyle
+	case models.MessageRoleToolCall:
+		style = assistantStyle
+		friendly = models.MessageRoleAssistant.FriendlyRole()
+	case models.MessageRoleToolResult:
+		icon = "🔧"
+	}
+
+	user := style.Render(icon + friendly)
+
+	var prefix string
+	var suffix string
+
+	faint := lipgloss.NewStyle().Faint(true)
+	if m.focus == focusMessages {
+		if i == m.selectedMessage {
+			prefix = "> "
+		}
+	}
+
+	if message.ID == 0 {
+		suffix += faint.Render(" (not saved)")
+	}
+
+	return headingStyle.Render(prefix + user + suffix)
+}
+
+func (m *model) renderMessage(msg *models.Message) string {
+	sb := &strings.Builder{}
+	sb.Grow(len(msg.Content) * 2)
+	if msg.Content != "" {
+		err := m.ctx.Chroma.Highlight(sb, msg.Content)
+		if err != nil {
+			sb.Reset()
+			sb.WriteString(msg.Content)
+		}
+	}
+
+	var toolString string
+	switch msg.Role {
+	case models.MessageRoleToolCall:
+		bytes, err := yaml.Marshal(msg.ToolCalls)
+		if err != nil {
+			toolString = "Could not serialize ToolCalls"
+		} else {
+			toolString = "tool_calls:\n" + string(bytes)
+		}
+	case models.MessageRoleToolResult:
+		if !m.showToolResults {
+			break
+		}
+
+		type renderedResult struct {
+			ToolName string `yaml:"tool"`
+			Result   any
+		}
+
+		var toolResults []renderedResult
+		for _, result := range msg.ToolResults {
+			var jsonResult interface{}
+			err := json.Unmarshal([]byte(result.Result), &jsonResult)
+			if err != nil {
+				// If parsing as JSON fails, treat Result as a plain string
+				toolResults = append(toolResults, renderedResult{
+					ToolName: result.ToolName,
+					Result:   result.Result,
+				})
+			} else {
+				// If parsing as JSON succeeds, marshal the parsed JSON into YAML
+				toolResults = append(toolResults, renderedResult{
+					ToolName: result.ToolName,
+					Result:   &jsonResult,
+				})
+			}
+		}
+
+		bytes, err := yaml.Marshal(toolResults)
+		if err != nil {
+			toolString = "Could not serialize ToolResults"
+		} else {
+			toolString = "tool_results:\n" + string(bytes)
+		}
+	}
+
+	if toolString != "" {
+		toolString = strings.TrimRight(toolString, "\n")
+		if msg.Content != "" {
+			sb.WriteString("\n\n")
+		}
+		_ = m.ctx.Chroma.HighlightLang(sb, toolString, "yaml")
+	}
+
+	content := strings.TrimRight(sb.String(), "\n")
+
+	if m.wrap {
+		wrapWidth := m.content.Width - messageStyle.GetHorizontalPadding() - 2
+		content = wordwrap.String(content, wrapWidth)
+	}
+
+	return messageStyle.Width(0).Render(content)
+}
+
 func (m *model) setMessages(messages []models.Message) {
 	m.messages = messages
-	m.highlightCache = make([]string, len(messages))
-	for i, msg := range m.messages {
-		highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
-		m.highlightCache[i] = highlighted
-	}
+	m.rebuildMessageCache()
 }
 
 func (m *model) setMessage(i int, msg models.Message) {
 	if i >= len(m.messages) {
 		panic("i out of range")
 	}
-	highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
 	m.messages[i] = msg
-	m.highlightCache[i] = highlighted
+	m.messageCache[i] = m.renderMessage(&msg)
 }
 
 func (m *model) addMessage(msg models.Message) {
-	highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
 	m.messages = append(m.messages, msg)
-	m.highlightCache = append(m.highlightCache, highlighted)
+	m.messageCache = append(m.messageCache, m.renderMessage(&msg))
 }
 
 func (m *model) setMessageContents(i int, content string) {
 	if i >= len(m.messages) {
 		panic("i out of range")
 	}
-	highlighted, _ := m.ctx.Chroma.HighlightS(content)
 	m.messages[i].Content = content
-	m.highlightCache[i] = highlighted
+	m.messageCache[i] = m.renderMessage(&m.messages[i])
+}
+
+func (m *model) rebuildMessageCache() {
+	m.messageCache = make([]string, len(m.messages))
+	for i, msg := range m.messages {
+		m.messageCache[i] = m.renderMessage(&msg)
+	}
 }
 
 func (m *model) updateContent() {
@@ -759,80 +909,35 @@ func (m *model) updateContent() {
 // render the conversation into a string
 func (m *model) conversationView() string {
 	sb := strings.Builder{}
-	msgCnt := len(m.messages)
 
 	m.messageOffsets = make([]int, len(m.messages))
-	lineCnt := conversationStyle.GetMarginTop()
+	lineCnt := 1
 	for i, message := range m.messages {
 		m.messageOffsets[i] = lineCnt
 
-		icon := "⚙️"
-		friendly := message.Role.FriendlyRole()
-		style := lipgloss.NewStyle().Bold(true).Faint(true)
-
 		switch message.Role {
-		case models.MessageRoleUser:
-			icon = ""
-			style = userStyle
-		case models.MessageRoleAssistant:
-			icon = ""
-			style = assistantStyle
-		case models.MessageRoleToolCall, models.MessageRoleToolResult:
-			icon = "🔧"
+		case models.MessageRoleToolCall:
+			if !m.showToolResults && message.Content == "" {
+				continue
+			}
+		case models.MessageRoleToolResult:
+			if !m.showToolResults {
+				continue
+			}
 		}
 
-		// write message heading with space for content
-		user := style.Render(icon + friendly)
+		heading := m.renderMessageHeading(i, &message)
+		sb.WriteString(heading)
+		sb.WriteString("\n")
+		lineCnt += lipgloss.Height(heading)
 
-		var prefix string
-		var suffix string
-
-		faint := lipgloss.NewStyle().Faint(true)
-		if m.focus == focusMessages {
-			if i == m.selectedMessage {
-				prefix = "> "
-			}
-			suffix += faint.Render(fmt.Sprintf(" (%d/%d)", i+1, msgCnt))
+		cached := m.messageCache[i]
+		sb.WriteString(cached)
+		sb.WriteString("\n")
+		lineCnt += lipgloss.Height(cached)
 	}
 
-		if message.ID == 0 {
-			suffix += faint.Render(" (not saved)")
-		}
-
-		header := lipgloss.NewStyle().PaddingLeft(1).Render(prefix + user + suffix)
-		sb.WriteString(header)
-		lineCnt += lipgloss.Height(header)
-
-		// TODO: special rendering for tool calls/results?
-		if message.Content != "" {
-			sb.WriteString("\n\n")
-			lineCnt += 1
-
-			// write message contents
-			var highlighted string
-			if m.highlightCache[i] == "" {
-				highlighted = message.Content
-			} else {
-				highlighted = m.highlightCache[i]
-			}
-			var contents string
-			if m.wrap {
-				wrapWidth := m.content.Width - messageStyle.GetHorizontalPadding() - 2
-				wrapped := wordwrap.String(highlighted, wrapWidth)
-				contents = wrapped
-			} else {
-				contents = highlighted
-			}
-			sb.WriteString(messageStyle.Width(0).Render(contents))
-			lineCnt += lipgloss.Height(contents)
-		}
-
-		if i < msgCnt-1 {
-			sb.WriteString("\n\n")
-			lineCnt += 1
-		}
-	}
-	return conversationStyle.Render(sb.String())
+	return sb.String()
 }
 
 func Launch(ctx *lmcli.Context, convShortname string) error {

pkg/util/tty/highlight.go

@@ -58,3 +58,29 @@ func (s *ChromaHighlighter) HighlightS(text string) (string, error) {
 	s.formatter.Format(&sb, s.style, it)
 	return sb.String(), nil
 }
+
+func (s *ChromaHighlighter) HighlightLang(w io.Writer, text string, lang string) (error) {
+	l := lexers.Get(lang)
+	if l == nil {
+		l = lexers.Fallback
+	}
+	l = chroma.Coalesce(l)
+	old := s.lexer
+	s.lexer = l
+	err := s.Highlight(w, text)
+	s.lexer = old
+	return err
+}
+
+func (s *ChromaHighlighter) HighlightLangS(text string, lang string) (string, error) {
+	l := lexers.Get(lang)
+	if l == nil {
+		l = lexers.Fallback
+	}
+	l = chroma.Coalesce(l)
+	old := s.lexer
+	s.lexer = l
+	highlighted, err := s.HighlightS(text)
+	s.lexer = old
+	return highlighted, err
+}