tui: add tool rendering

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

4
go.mod
View File

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

View File

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

View File

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

View File

@ -10,6 +10,7 @@ package tui
import ( import (
"context" "context"
"encoding/json"
"fmt" "fmt"
"strings" "strings"
"time" "time"
@ -24,6 +25,7 @@ import (
"github.com/charmbracelet/lipgloss" "github.com/charmbracelet/lipgloss"
"github.com/muesli/reflow/ansi" "github.com/muesli/reflow/ansi"
"github.com/muesli/reflow/wordwrap" "github.com/muesli/reflow/wordwrap"
"gopkg.in/yaml.v2"
) )
type focusState int type focusState int
@ -70,7 +72,8 @@ type model struct {
focus focusState focus focusState
wrap bool // whether message content is wrapped to viewport width wrap bool // whether message content is wrapped to viewport width
status string // a general status message 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 messageOffsets []int
selectedMessage int selectedMessage int
@ -109,14 +112,18 @@ type (
// styles // styles
var ( var (
userStyle = lipgloss.NewStyle().Faint(true).Bold(true).Foreground(lipgloss.Color("10")) headingStyle = lipgloss.NewStyle().
assistantStyle = lipgloss.NewStyle().Faint(true).Bold(true).Foreground(lipgloss.Color("12")) MarginTop(1).
messageStyle = lipgloss.NewStyle().PaddingLeft(2).PaddingRight(2) 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(). headerStyle = lipgloss.NewStyle().
Background(lipgloss.Color("0")) Background(lipgloss.Color("0"))
conversationStyle = lipgloss.NewStyle().
MarginTop(1).
MarginBottom(1)
footerStyle = lipgloss.NewStyle(). footerStyle = lipgloss.NewStyle().
BorderTop(true). BorderTop(true).
BorderStyle(lipgloss.NormalBorder()) BorderStyle(lipgloss.NormalBorder())
@ -171,6 +178,11 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.persistence = !m.persistence m.persistence = !m.persistence
case "ctrl+w": case "ctrl+w":
m.wrap = !m.wrap m.wrap = !m.wrap
m.rebuildMessageCache()
m.updateContent()
case "ctrl+t":
m.showToolResults = !m.showToolResults
m.rebuildMessageCache()
m.updateContent() m.updateContent()
case "q": case "q":
if m.focus != focusInput { if m.focus != focusInput {
@ -193,6 +205,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
m.height = msg.Height m.height = msg.Height
m.content.Width = msg.Width m.content.Width = msg.Width
m.input.SetWidth(msg.Width - 1) m.input.SetWidth(msg.Width - 1)
m.rebuildMessageCache()
m.updateContent() m.updateContent()
case msgConversationLoaded: case msgConversationLoaded:
m.conversation = (*models.Conversation)(msg) m.conversation = (*models.Conversation)(msg)
@ -521,7 +534,7 @@ func (m *model) handleMessagesKey(msg tea.KeyMsg) tea.Cmd {
return nil return nil
} }
m.messages = m.messages[:m.selectedMessage+1] m.messages = m.messages[:m.selectedMessage+1]
m.highlightCache = m.highlightCache[:m.selectedMessage+1] m.messageCache = m.messageCache[:m.selectedMessage+1]
m.updateContent() m.updateContent()
m.content.GotoBottom() m.content.GotoBottom()
return m.promptLLM() return m.promptLLM()
@ -533,9 +546,13 @@ func (m *model) handleInputKey(msg tea.KeyMsg) tea.Cmd {
switch msg.String() { switch msg.String() {
case "esc": case "esc":
m.focus = focusMessages m.focus = focusMessages
if len(m.messages) > 0 {
if m.selectedMessage < 0 || m.selectedMessage >= len(m.messages) { if m.selectedMessage < 0 || m.selectedMessage >= len(m.messages) {
m.selectedMessage = len(m.messages) - 1 m.selectedMessage = len(m.messages) - 1
} }
offset := m.messageOffsets[m.selectedMessage]
scrollIntoView(&m.content, offset, 0.1)
}
m.updateContent() m.updateContent()
m.input.Blur() m.input.Blur()
case "ctrl+s": case "ctrl+s":
@ -734,37 +751,150 @@ func (m *model) persistConversation() tea.Cmd {
return nil 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) { func (m *model) setMessages(messages []models.Message) {
m.messages = messages m.messages = messages
m.highlightCache = make([]string, len(messages)) m.rebuildMessageCache()
for i, msg := range m.messages {
highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
m.highlightCache[i] = highlighted
}
} }
func (m *model) setMessage(i int, msg models.Message) { func (m *model) setMessage(i int, msg models.Message) {
if i >= len(m.messages) { if i >= len(m.messages) {
panic("i out of range") panic("i out of range")
} }
highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
m.messages[i] = msg m.messages[i] = msg
m.highlightCache[i] = highlighted m.messageCache[i] = m.renderMessage(&msg)
} }
func (m *model) addMessage(msg models.Message) { func (m *model) addMessage(msg models.Message) {
highlighted, _ := m.ctx.Chroma.HighlightS(msg.Content)
m.messages = append(m.messages, msg) 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) { func (m *model) setMessageContents(i int, content string) {
if i >= len(m.messages) { if i >= len(m.messages) {
panic("i out of range") panic("i out of range")
} }
highlighted, _ := m.ctx.Chroma.HighlightS(content)
m.messages[i].Content = 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() { func (m *model) updateContent() {
@ -779,80 +909,35 @@ func (m *model) updateContent() {
// render the conversation into a string // render the conversation into a string
func (m *model) conversationView() string { func (m *model) conversationView() string {
sb := strings.Builder{} sb := strings.Builder{}
msgCnt := len(m.messages)
m.messageOffsets = make([]int, len(m.messages)) m.messageOffsets = make([]int, len(m.messages))
lineCnt := conversationStyle.GetMarginTop() lineCnt := 1
for i, message := range m.messages { for i, message := range m.messages {
m.messageOffsets[i] = lineCnt m.messageOffsets[i] = lineCnt
icon := "⚙️"
friendly := message.Role.FriendlyRole()
style := lipgloss.NewStyle().Bold(true).Faint(true)
switch message.Role { switch message.Role {
case models.MessageRoleUser: case models.MessageRoleToolCall:
icon = "" if !m.showToolResults && message.Content == "" {
style = userStyle continue
case models.MessageRoleAssistant: }
icon = "" case models.MessageRoleToolResult:
style = assistantStyle if !m.showToolResults {
case models.MessageRoleToolCall, models.MessageRoleToolResult: continue
icon = "🔧" }
} }
// write message heading with space for content heading := m.renderMessageHeading(i, &message)
user := style.Render(icon + friendly) sb.WriteString(heading)
sb.WriteString("\n")
lineCnt += lipgloss.Height(heading)
var prefix string cached := m.messageCache[i]
var suffix string sb.WriteString(cached)
sb.WriteString("\n")
faint := lipgloss.NewStyle().Faint(true) lineCnt += lipgloss.Height(cached)
if m.focus == focusMessages {
if i == m.selectedMessage {
prefix = "> "
}
suffix += faint.Render(fmt.Sprintf(" (%d/%d)", i+1, msgCnt))
} }
if message.ID == 0 { return sb.String()
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())
} }
func Launch(ctx *lmcli.Context, convShortname string) error { func Launch(ctx *lmcli.Context, convShortname string) error {

View File

@ -58,3 +58,29 @@ func (s *ChromaHighlighter) HighlightS(text string) (string, error) {
s.formatter.Format(&sb, s.style, it) s.formatter.Format(&sb, s.style, it)
return sb.String(), nil 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
}