mlow/lmcli
1
0
Fork 0
Code Issues 3 Pull Requests Packages Projects Releases Wiki 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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
go.mod
View File

@ -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
)

6
pkg/lmcli/model/conversation.go
View File

@ -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)
}
}

12
pkg/lmcli/model/tool.go
View File

@ -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

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 {

26
pkg/util/tty/highlight.go
View File

@ -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
}