Improve message wrapping behavior

A better/more natural indication that the model is doing something

pkg/cmd/retry.go

@@ -66,7 +66,7 @@ func RetryCmd(ctx *lmcli.Context) *cobra.Command {
 		},
 	}
 
-	cmd.Flags().Int("offset", 1, "Offset from the last message retry from.")
+	cmd.Flags().Int("offset", 0, "Offset from the last message to retry from.")
 
 	applyPromptFlags(ctx, cmd)
 	return cmd

pkg/tui/chat.go

@@ -9,12 +9,14 @@ import (
 
 	cmdutil "git.mlow.ca/mlow/lmcli/pkg/cmd/util"
 	models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
+	"github.com/charmbracelet/bubbles/cursor"
 	"github.com/charmbracelet/bubbles/spinner"
 	"github.com/charmbracelet/bubbles/textarea"
 	"github.com/charmbracelet/bubbles/viewport"
 	tea "github.com/charmbracelet/bubbletea"
 	"github.com/charmbracelet/lipgloss"
 	"github.com/muesli/reflow/wordwrap"
+	"github.com/muesli/reflow/wrap"
 	"gopkg.in/yaml.v2"
 )
 
@@ -79,9 +81,10 @@ type chatModel struct {
 	messageOffsets  []int
 
 	// ui elements
-	content viewport.Model
+	content     viewport.Model
-	input   textarea.Model
+	input       textarea.Model
-	spinner spinner.Model
+	spinner     spinner.Model
+	replyCursor cursor.Model // cursor to indicate incoming response
 }
 
 func newChatModel(tui *model) chatModel {
@@ -117,8 +120,12 @@ func newChatModel(tui *model) chatModel {
 				FPS: time.Second / 3,
 			},
 		)),
+		replyCursor: cursor.New(),
 	}
 
+	m.replyCursor.SetChar(" ")
+	m.replyCursor.Focus()
+
 	system := tui.ctx.GetSystemPrompt()
 	if system != "" {
 		m.messages = []models.Message{{
@@ -220,8 +227,6 @@ func (m *chatModel) handleInput(msg tea.KeyMsg) (bool, tea.Cmd) {
 
 func (m chatModel) Init() tea.Cmd {
 	return tea.Batch(
-		textarea.Blink,
-		m.spinner.Tick,
 		m.waitForChunk(),
 		m.waitForReply(),
 	)
@@ -241,9 +246,13 @@ func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 	var cmds []tea.Cmd
 	switch msg := msg.(type) {
 	case msgStateEnter:
+		// wake up spinners and cursors
+		cmds = append(cmds, cursor.Blink, m.spinner.Tick)
+
 		if m.opts.convShortname != "" && m.conversation.ShortName.String != m.opts.convShortname {
 			cmds = append(cmds, m.loadConversation(m.opts.convShortname))
 		}
+
 		m.rebuildMessageCache()
 		m.updateContent()
 	case tea.WindowSizeMsg:
@@ -275,10 +284,16 @@ func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 		m.content.GotoBottom()
 	case msgResponseChunk:
 		chunk := string(msg)
+		if chunk == "" {
+			break
+		}
+
 		last := len(m.messages) - 1
 		if last >= 0 && m.messages[last].Role.IsAssistant() {
+			// append chunk to existing message
 			m.setMessageContents(last, m.messages[last].Content+chunk)
 		} else {
+			// use chunk in new message
 			m.addMessage(models.Message{
 				Role:    models.MessageRoleAssistant,
 				Content: chunk,
@@ -287,6 +302,10 @@ func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 		m.updateContent()
 		cmds = append(cmds, m.waitForChunk()) // wait for the next chunk
 
+		// show cursor and reset blink interval (simulate typing)
+		m.replyCursor.Blink = false
+		cmds = append(cmds, m.replyCursor.BlinkCmd())
+
 		m.tokenCount++
 		m.elapsed = time.Now().Sub(m.startTime)
 	case msgAssistantReply:
@@ -332,6 +351,7 @@ func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 		m.waitingForReply = false
 		m.status = "Press ctrl+s to send"
 		m.err = error(msg)
+		m.updateContent()
 	case msgConversationTitleChanged:
 		title := string(msg)
 		m.conversation.Title = title
@@ -341,6 +361,11 @@ func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 				cmds = append(cmds, wrapError(err))
 			}
 		}
+	case cursor.BlinkMsg:
+		if m.waitingForReply {
+			// ensure we show updated "wait for response" cursor blink state
+			m.updateContent()
+		}
 	}
 
 	var cmd tea.Cmd
@@ -348,6 +373,10 @@ func (m chatModel) Update(msg tea.Msg) (chatModel, tea.Cmd) {
 	if cmd != nil {
 		cmds = append(cmds, cmd)
 	}
+	m.replyCursor, cmd = m.replyCursor.Update(msg)
+	if cmd != nil {
+		cmds = append(cmds, cmd)
+	}
 
 	prevInputLineCnt := m.input.LineCount()
 	inputCaptured := false
@@ -442,9 +471,10 @@ func (m *chatModel) handleMessagesKey(msg tea.KeyMsg) (bool, tea.Cmd) {
 		}
 		m.messages = m.messages[:m.selectedMessage+1]
 		m.messageCache = m.messageCache[:m.selectedMessage+1]
+		cmd := m.promptLLM()
 		m.updateContent()
 		m.content.GotoBottom()
-		return true, m.promptLLM()
+		return true, cmd
 	}
 	return false, nil
 }
@@ -464,6 +494,10 @@ func (m *chatModel) handleInputKey(msg tea.KeyMsg) (bool, tea.Cmd) {
 		m.input.Blur()
 		return true, nil
 	case "ctrl+s":
+		if m.waitingForReply {
+			return false, nil
+		}
+
 		input := strings.TrimSpace(m.input.Value())
 		if input == "" {
 			return true, nil
@@ -487,9 +521,10 @@ func (m *chatModel) handleInputKey(msg tea.KeyMsg) (bool, tea.Cmd) {
 			}
 		}
 
+		cmd := m.promptLLM()
 		m.updateContent()
 		m.content.GotoBottom()
-		return true, m.promptLLM()
+		return true, cmd
 	case "ctrl+e":
 		cmd := openTempfileEditor("message.*.md", m.input.Value(), "# Edit your input below\n")
 		m.editorTarget = input
@@ -536,7 +571,10 @@ func (m *chatModel) renderMessageHeading(i int, message *models.Message) string
 	return messageHeadingStyle.Render(prefix + user + suffix)
 }
 
-func (m *chatModel) renderMessage(msg *models.Message) string {
+func (m *chatModel) renderMessage(i int) string {
+	msg := &m.messages[i]
+
+	// Write message contents
 	sb := &strings.Builder{}
 	sb.Grow(len(msg.Content) * 2)
 	if msg.Content != "" {
@@ -547,6 +585,12 @@ func (m *chatModel) renderMessage(msg *models.Message) string {
 		}
 	}
 
+	// Show the assistant's cursor
+	if m.waitingForReply && i == len(m.messages)-1 {
+		sb.WriteString(m.replyCursor.View())
+	}
+
+	// Write tool call info
 	var toolString string
 	switch msg.Role {
 	case models.MessageRoleToolCall:
@@ -604,8 +648,11 @@ func (m *chatModel) renderMessage(msg *models.Message) string {
 	content := strings.TrimRight(sb.String(), "\n")
 
 	if m.wrap {
-		wrapWidth := m.content.Width - messageStyle.GetHorizontalPadding() - 2
+		wrapWidth := m.content.Width - messageStyle.GetHorizontalPadding()
-		content = wordwrap.String(content, wrapWidth)
+		// first we word-wrap text to slightly less than desired width (since
+		// wordwrap seems to have an off-by-1 issue), then hard wrap at
+		// desired with
+		content = wrap.String(wordwrap.String(content, wrapWidth-2), wrapWidth)
 	}
 
 	return messageStyle.Width(0).Render(content)
@@ -636,10 +683,18 @@ func (m *chatModel) conversationMessagesView() string {
 		sb.WriteString("\n")
 		lineCnt += lipgloss.Height(heading)
 
-		cached := m.messageCache[i]
+		var rendered string
-		sb.WriteString(cached)
+		if m.waitingForReply && i == len(m.messages)-1 {
+			// do a direct render of final (assistant) message to handle the
+			// assistant cursor blink
+			rendered = m.renderMessage(i)
+		} else {
+			rendered = m.messageCache[i]
+		}
+
+		sb.WriteString(rendered)
 		sb.WriteString("\n")
-		lineCnt += lipgloss.Height(cached)
+		lineCnt += lipgloss.Height(rendered)
 	}
 
 	return sb.String()
@@ -712,12 +767,12 @@ func (m *chatModel) setMessage(i int, msg models.Message) {
 		panic("i out of range")
 	}
 	m.messages[i] = msg
-	m.messageCache[i] = m.renderMessage(&msg)
+	m.messageCache[i] = m.renderMessage(i)
 }
 
 func (m *chatModel) addMessage(msg models.Message) {
 	m.messages = append(m.messages, msg)
-	m.messageCache = append(m.messageCache, m.renderMessage(&msg))
+	m.messageCache = append(m.messageCache, m.renderMessage(len(m.messages)-1))
 }
 
 func (m *chatModel) setMessageContents(i int, content string) {
@@ -725,13 +780,13 @@ func (m *chatModel) setMessageContents(i int, content string) {
 		panic("i out of range")
 	}
 	m.messages[i].Content = content
-	m.messageCache[i] = m.renderMessage(&m.messages[i])
+	m.messageCache[i] = m.renderMessage(i)
 }
 
 func (m *chatModel) rebuildMessageCache() {
 	m.messageCache = make([]string, len(m.messages))
-	for i, msg := range m.messages {
+	for i := range m.messages {
-		m.messageCache[i] = m.renderMessage(&msg)
+		m.messageCache[i] = m.renderMessage(i)
 	}
 }
 
@@ -792,7 +847,7 @@ func (m *chatModel) persistConversation() error {
 			if err != nil {
 				return err
 			}
-		} else if i > 0  {
+		} else if i > 0 {
 			// messages is new, so add it as a reply to previous message
 			saved, err := m.ctx.Store.Reply(&m.messages[i-1], m.messages[i])
 			if err != nil {
@@ -833,8 +888,19 @@ func (m *chatModel) waitForChunk() tea.Cmd {
 
 func (m *chatModel) promptLLM() tea.Cmd {
 	m.waitingForReply = true
+	m.replyCursor.Blink = false
 	m.status = "Press ctrl+c to cancel"
 
+	toPrompt := m.messages
+
+	// Add response placeholder message
+	if m.messages[len(m.messages)-1].Role != models.MessageRoleAssistant {
+		m.addMessage(models.Message{
+			Role:    models.MessageRoleAssistant,
+			Content: "",
+		})
+	}
+
 	m.tokenCount = 0
 	m.startTime = time.Now()
 	m.elapsed = 0
@@ -868,7 +934,7 @@ func (m *chatModel) promptLLM() tea.Cmd {
 		}()
 
 		resp, err := completionProvider.CreateChatCompletionStream(
-			ctx, requestParams, m.messages, replyHandler, m.replyChunkChan,
+			ctx, requestParams, toPrompt, replyHandler, m.replyChunkChan,
 		)
 
 		if err != nil && !canceled {