From 8bf4b40d94e1020d536f36846d76e0d4c129de2f Mon Sep 17 00:00:00 2001 From: Arpit Mohapatra Date: Fri, 16 Aug 2024 11:54:48 +0530 Subject: [PATCH] added markdown rendering for model output --- internal/chat/interface.go | 180 ++++++++++++++++++++++++++----------- 1 file changed, 126 insertions(+), 54 deletions(-) diff --git a/internal/chat/interface.go b/internal/chat/interface.go index dd7a9e0..e9965d2 100644 --- a/internal/chat/interface.go +++ b/internal/chat/interface.go @@ -8,6 +8,7 @@ import ( "github.com/charmbracelet/bubbles/textinput" "github.com/charmbracelet/bubbles/viewport" tea "github.com/charmbracelet/bubbletea" + "github.com/charmbracelet/glamour" "github.com/charmbracelet/lipgloss" "github.com/marpit19/charmlama/internal/ollama" ) @@ -21,10 +22,18 @@ var ( Foreground(lipgloss.Color("#00FFFF")). // Cyan Bold(true) - inputStyle = lipgloss.NewStyle(). - BorderStyle(lipgloss.NormalBorder()). - BorderForeground(lipgloss.Color("#FF1493")). // Deep Pink - Foreground(lipgloss.Color("#FFFFFF")) // White text + activeInputStyle = lipgloss.NewStyle(). + BorderStyle(lipgloss.NormalBorder()). + BorderForeground(lipgloss.Color("#FF1493")). // Deep Pink + Foreground(lipgloss.Color("#FFFFFF")) // White text + + disabledInputStyle = lipgloss.NewStyle(). + BorderStyle(lipgloss.NormalBorder()). + BorderForeground(lipgloss.Color("#696969")). // Dim Gray + Foreground(lipgloss.Color("#A9A9A9")) // Dark Gray text + + statusStyle = lipgloss.NewStyle(). + Foreground(lipgloss.Color("#98FB98")) ) type ChatInterface struct { @@ -39,6 +48,7 @@ type ChatInterface struct { spinner spinner.Model width int height int + renderer *glamour.TermRenderer } func NewChatInterface(model string, manager *ollama.Manager) *ChatInterface { @@ -47,13 +57,18 @@ func NewChatInterface(model string, manager *ollama.Manager) *ChatInterface { input.Focus() vp := viewport.New(80, 20) - vp.KeyMap.PageDown.SetEnabled(false) - vp.KeyMap.PageUp.SetEnabled(false) + // vp.KeyMap.PageDown.SetEnabled(false) + // vp.KeyMap.PageUp.SetEnabled(false) s := spinner.New() s.Spinner = spinner.Dot s.Style = lipgloss.NewStyle().Foreground(lipgloss.Color("#FFA500")) + renderer, _ := glamour.NewTermRenderer( + glamour.WithAutoStyle(), + glamour.WithWordWrap(80), + ) + return &ChatInterface{ model: model, manager: manager, @@ -61,6 +76,7 @@ func NewChatInterface(model string, manager *ollama.Manager) *ChatInterface { viewport: vp, input: input, spinner: s, + renderer: renderer, } } @@ -73,6 +89,17 @@ func (c *ChatInterface) Update(msg tea.Msg) (tea.Model, tea.Cmd) { switch msg := msg.(type) { case tea.KeyMsg: + if c.waiting { + // Ignore most key presses while waiting + switch msg.String() { + case "ctrl+c": + c.quitting = true + return c, tea.Quit + default: + return c, nil + } + } + switch msg.String() { case "ctrl+c", "/exit": c.quitting = true @@ -99,43 +126,75 @@ func (c *ChatInterface) Update(msg tea.Msg) (tea.Model, tea.Cmd) { case userMessageMsg: c.addMessage("You", string(msg)) c.waiting = true + c.input.Blur() cmds = append(cmds, c.handleUserMessage(msg), c.spinner.Tick) case aiResponseMsg: c.waiting = false c.addMessage(c.model, string(msg)) + c.input.Focus() } if c.waiting { var cmd tea.Cmd c.spinner, cmd = c.spinner.Update(msg) cmds = append(cmds, cmd) + } else { + var cmd tea.Cmd + c.input, cmd = c.input.Update(msg) + cmds = append(cmds, cmd) } - c.input, _ = c.input.Update(msg) + // c.input, _ = c.input.Update(msg) + + var cmd tea.Cmd + c.viewport, cmd = c.viewport.Update(msg) + cmds = append(cmds, cmd) return c, tea.Batch(cmds...) } +// func (c *ChatInterface) View() string { +// var status string +// var inputView string + +// if c.waiting { +// status = fmt.Sprintf("%s AI is thinking...", c.spinner.View()) +// inputView = disabledInputStyle.Render(c.input.View()) +// } else { +// status = "Ready for your message" +// inputView = activeInputStyle.Render(c.input.View()) +// } + +// maxInputWidth := c.width - 4 // Adjust this value as needed +// if len(inputView) > maxInputWidth && maxInputWidth > 0 { +// inputView = inputView[:maxInputWidth] + "..." +// } + +// return fmt.Sprintf( +// "%s\n%s\n%s", +// c.viewport.View(), +// inputView, +// lipgloss.NewStyle().Foreground(lipgloss.Color("#98FB98")).Render(status), +// ) +// } + func (c *ChatInterface) View() string { var status string + var inputView string + if c.waiting { status = fmt.Sprintf("%s AI is thinking...", c.spinner.View()) + inputView = disabledInputStyle.Render(c.input.View()) } else { status = "Ready for your message" + inputView = activeInputStyle.Render(c.input.View()) } - inputView := c.input.View() - maxInputWidth := c.width - 4 // Adjust this value as needed - if len(inputView) > maxInputWidth && maxInputWidth > 0 { - inputView = inputView[:maxInputWidth] + "..." - } - - return fmt.Sprintf( - "%s\n%s\n%s", + return lipgloss.JoinVertical(lipgloss.Left, c.viewport.View(), - inputStyle.Render(inputView), - lipgloss.NewStyle().Foreground(lipgloss.Color("#98FB98")).Render(status), + inputView, + statusStyle.Render(status), ) } @@ -161,52 +220,65 @@ func (c *ChatInterface) handleUserMessage(msg userMessageMsg) tea.Cmd { } } +// func (c *ChatInterface) addMessage(sender, content string) { +// style := userStyle +// if sender != "You" { +// style = aiStyle +// } +// formattedMsg := style.Render(sender+":") + " " + content +// wrappedMsg := c.wrapText(formattedMsg, c.width) +// c.messages = append(c.messages, wrappedMsg) +// c.updateViewportContent() +// } + func (c *ChatInterface) addMessage(sender, content string) { - style := userStyle - if sender != "You" { - style = aiStyle + var formattedMsg string + if sender == "You" { + formattedMsg = userStyle.Render(sender+":") + " " + content + } else { + rendered, _ := c.renderer.Render(content) + formattedMsg = aiStyle.Render(sender+":") + "\n" + rendered } - formattedMsg := style.Render(sender+":") + " " + content - wrappedMsg := c.wrapText(formattedMsg, c.width) - c.messages = append(c.messages, wrappedMsg) + c.messages = append(c.messages, formattedMsg) c.updateViewportContent() } func (c *ChatInterface) updateViewportContent() { - c.viewport.SetContent(strings.Join(c.messages, "\n\n")) + content := strings.Join(c.messages, "\n\n") + c.viewport.SetContent(content) c.viewport.GotoBottom() } -func (c *ChatInterface) wrapText(text string, width int) string { - if width <= 0 { - return text - } - words := strings.Fields(text) - if len(words) == 0 { - return text - } - - var lines []string - var currentLine string - - for _, word := range words { - if len(currentLine)+len(word)+1 > width { - lines = append(lines, strings.TrimSpace(currentLine)) - currentLine = word - } else { - if currentLine != "" { - currentLine += " " - } - currentLine += word - } - } - - if currentLine != "" { - lines = append(lines, strings.TrimSpace(currentLine)) - } - - return strings.Join(lines, "\n") -} +// func (c *ChatInterface) wrapText(text string, width int) string { +// if width <= 0 { +// return text +// } +// words := strings.Fields(text) +// if len(words) == 0 { +// return text +// } + +// var lines []string +// var currentLine string + +// for _, word := range words { +// if len(currentLine)+len(word)+1 > width { +// lines = append(lines, strings.TrimSpace(currentLine)) +// currentLine = word +// } else { +// if currentLine != "" { +// currentLine += " " +// } +// currentLine += word +// } +// } + +// if currentLine != "" { +// lines = append(lines, strings.TrimSpace(currentLine)) +// } + +// return strings.Join(lines, "\n") +// } func (c *ChatInterface) Run() (bool, error) { p := tea.NewProgram(c, tea.WithAltScreen())