mlow/lmcli
Private
Public Access
1
0
Fork 0
Code Activity

tui: add tool rendering

Browse Source 
cleaned up message rendering and changed cache semantics

other smaller tweaks
This commit is contained in:
Matt Low
2024-03-26 07:59:39 +00:00
parent 6310021dca
commit a669313a0b
5 changed files with 207 additions and 96 deletions
Download Patch File Download Diff File Expand all files Collapse all files

255
pkg/tui/tui.go
View File

@@ -10,6 +10,7 @@ package tui
import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
@@ -24,6 +25,7 @@ import (
"github.com/charmbracelet/lipgloss"
"github.com/muesli/reflow/ansi"
"github.com/muesli/reflow/wordwrap"
"gopkg.in/yaml.v2"
)
type focusState int
@@ -70,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
@@ -109,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)
headerStyle = lipgloss.NewStyle().
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())
@@ -171,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 {
@@ -193,6 +205,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.height = msg.Height
m.content.Width = msg.Width
m.input.SetWidth(msg.Width - 1)
m.rebuildMessageCache()
m.updateContent()
case msgConversationLoaded:
m.conversation = (*models.Conversation)(msg)
@@ -521,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()
@@ -533,8 +546,12 @@ func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd {
switch msg.String() {
case "esc":
m.focus = focusMessages
if m.selectedMessage < 0 || m.selectedMessage >= len(m.messages) {
m.selectedMessage = len(m.messages) - 1
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()
@@ -734,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() {
@@ -779,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 = "🔧"
}
// write message heading with space for content
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 = "> "
case models.MessageRoleToolCall:
if !m.showToolResults && message.Content == "" {
continue
}
suffix += faint.Render(fmt.Sprintf(" (%d/%d)", i+1, msgCnt))
}
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]
case models.MessageRoleToolResult:
if !m.showToolResults {
continue
}
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
}
heading := m.renderMessageHeading(i, &message)
sb.WriteString(heading)
sb.WriteString("\n")
lineCnt += lipgloss.Height(heading)
cached := m.messageCache[i]
sb.WriteString(cached)
sb.WriteString("\n")
lineCnt += lipgloss.Height(cached)
}
return conversationStyle.Render(sb.String())
return sb.String()
}
func Launch(ctx *lmcli.Context, convShortname string) error {