From 377a4f1dfa1cf1b9574186509dfeed0ce272af11 Mon Sep 17 00:00:00 2001
From: Matt Low <matt@mlow.ca>
Date: Thu, 14 Mar 2024 03:07:41 +0000
Subject: [PATCH] tui: add LLM response error handling

+ various other small tweaks
---
 pkg/tui/tui.go | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/pkg/tui/tui.go b/pkg/tui/tui.go
index 56da29a..2912bdd 100644
--- a/pkg/tui/tui.go
+++ b/pkg/tui/tui.go
@@ -66,6 +66,8 @@ type (
 	msgResponseChunk string
 	// sent when response is finished being received
 	msgResponseEnd string
+	// a special case of msgError that stops the response waiting animation
+	msgResponseError error
 	// sent on each completed reply
 	msgAssistantReply models.Message
 	// sent when a conversation is (re)loaded
@@ -142,7 +144,7 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.height = msg.Height
 		m.content.Width = msg.Width
 		m.content.Height = msg.Height - m.getFixedComponentHeight()
-		m.input.SetWidth(msg.Width-1)
+		m.input.SetWidth(msg.Width - 1)
 		m.updateContent()
 	case msgConversationLoaded:
 		m.conversation = (*models.Conversation)(msg)
@@ -198,6 +200,11 @@ func (m model) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		m.replyCancelFunc = nil
 		m.waitingForReply = false
 		m.status = "Press ctrl+s to send"
+	case msgResponseError:
+		m.replyCancelFunc = nil
+		m.waitingForReply = false
+		m.status = "Press ctrl+s to send"
+		m.err = error(msg)
 	case msgConversationTitleChanged:
 		title := string(msg)
 		m.conversation.Title = title
@@ -288,7 +295,7 @@ func (m *model) headerView() string {
 	}
 	part := titleStyle.Render(title)
 
-	return headerStyle.Width(m.content.Width).Render(part)
+	return headerStyle.Width(m.width).Render(part)
 }
 
 func (m *model) contentView() string {
@@ -300,10 +307,10 @@ func (m *model) errorView() string {
 		return ""
 	}
 	return lipgloss.NewStyle().
+		Width(m.width).
+		AlignHorizontal(lipgloss.Center).
 		Bold(true).
 		Foreground(lipgloss.Color("1")).
-		Width(m.content.Width).
-		AlignHorizontal(lipgloss.Center).
 		Render(fmt.Sprintf("%s", m.err))
 }
 
@@ -312,12 +319,12 @@ func (m *model) scrollbarView() string {
 		return ""
 	}
 
-	count := int(m.content.ScrollPercent() * float64(m.content.Width-2))
+	count := int(m.content.ScrollPercent() * float64(m.width-2))
 	fill := strings.Repeat("-", count)
 	return lipgloss.NewStyle().
+		Width(m.width).
 		PaddingLeft(1).
 		PaddingRight(1).
-		Width(m.content.Width).
 		Render(fill)
 }
 
@@ -381,6 +388,7 @@ func initialModel(ctx *lmcli.Context, convShortname string) model {
 
 	m.input.FocusedStyle.CursorLine = lipgloss.NewStyle()
 	m.input.ShowLineNumbers = false
+	m.input.SetHeight(4)
 	m.input.Focus()
 
 	m.updateContent()
@@ -458,10 +466,10 @@ func (m *model) loadConversation(shortname string) tea.Cmd {
 		}
 		c, err := m.ctx.Store.ConversationByShortName(shortname)
 		if err != nil {
-			return msgError(fmt.Errorf("Could not lookup conversation: %v\n", err))
+			return msgError(fmt.Errorf("Could not lookup conversation: %v", err))
 		}
 		if c.ID == 0 {
-			return msgError(fmt.Errorf("Conversation not found with short name: %s\n", shortname))
+			return msgError(fmt.Errorf("Conversation not found: %s", shortname))
 		}
 		return msgConversationLoaded(c)
 	}
@@ -520,11 +528,13 @@ func (m *model) promptLLM() tea.Cmd {
 		ctx, replyCancelFunc := context.WithCancel(context.Background())
 		m.replyCancelFunc = replyCancelFunc
 
-		// TODO: handle error
-		resp, _ := completionProvider.CreateChatCompletionStream(
+		resp, err := completionProvider.CreateChatCompletionStream(
 			ctx, requestParams, m.messages, replyHandler, m.replyChunkChan,
 		)
 
+		if err != nil {
+			return msgResponseError(err)
+		}
 		return msgResponseEnd(resp)
 	}
 }