Package restructure and API changes, several fixes

- More emphasis on `api` package. It now holds database model structs from `lmcli/models` (which is now gone) as well as the tool spec, call, and result types. `tools.Tool` is now `api.ToolSpec`. `api.ChatCompletionClient` was renamed to `api.ChatCompletionProvider`. - Change ChatCompletion interface and implementations to no longer do automatic tool call recursion - they simply return a ToolCall message which the caller can decide what to do with (e.g. prompt for user confirmation before executing) - `api.ChatCompletionProvider` functions have had their ReplyCallback parameter removed, as now they only return a single reply. - Added a top-level `agent` package, moved the current built-in tools implementations under `agent/toolbox`. `tools.ExecuteToolCalls` is now `agent.ExecuteToolCalls`. - Fixed request context handling in openai, google, ollama (use `NewRequestWithContext`), cleaned up request cancellation in TUI - Fix tool call tui persistence bug (we were skipping message with empty content) - Now handle tool calling from TUI layer TODO: - Prompt users before executing tool calls - Automatically send tool results to the model (or make this toggleable)
2024-06-12 08:35:07 +00:00
parent 85a2abbbf3
commit 3fde58b77d
35 changed files with 608 additions and 749 deletions
--- a/pkg/tui/views/chat/chat.go
+++ b/pkg/tui/views/chat/chat.go
@@ -4,7 +4,6 @@ import (
 	"time"

 	"git.mlow.ca/mlow/lmcli/pkg/api"
-	models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
 	"git.mlow.ca/mlow/lmcli/pkg/tui/shared"
 	"github.com/charmbracelet/bubbles/cursor"
 	"github.com/charmbracelet/bubbles/spinner"
@@ -16,37 +15,39 @@ import (

 // custom tea.Msg types
 type (
-	// sent on each chunk received from LLM
-	msgResponseChunk api.Chunk
-	// sent when response is finished being received
-	msgResponseEnd string
-	// a special case of common.MsgError that stops the response waiting animation
-	msgResponseError error
-	// sent on each completed reply
-	msgResponse models.Message
 	// sent when a conversation is (re)loaded
 	msgConversationLoaded struct {
-		conversation *models.Conversation
-		rootMessages []models.Message
+		conversation *api.Conversation
+		rootMessages []api.Message
 	}
 	// sent when a new conversation title generated
 	msgConversationTitleGenerated string
-	// sent when a conversation's messages are laoded
-	msgMessagesLoaded []models.Message
 	// sent when the conversation has been persisted, triggers a reload of contents
 	msgConversationPersisted struct {
 		isNew        bool
-		conversation *models.Conversation
-		messages     []models.Message
+		conversation *api.Conversation
+		messages     []api.Message
 	}
+	// sent when a conversation's messages are laoded
+	msgMessagesLoaded []api.Message
+	// a special case of common.MsgError that stops the response waiting animation
+	msgChatResponseError error
+	// sent on each chunk received from LLM
+	msgChatResponseChunk api.Chunk
+	// sent on each completed reply
+	msgChatResponse *api.Message
+	// sent when the response is canceled
+	msgChatResponseCanceled struct{}
+	// sent when results from a tool call are returned
+	msgToolResults []api.ToolResult
 	// sent when the given message is made the new selected reply of its parent
-	msgSelectedReplyCycled *models.Message
+	msgSelectedReplyCycled *api.Message
 	// sent when the given message is made the new selected root of the current conversation
-	msgSelectedRootCycled *models.Message
+	msgSelectedRootCycled *api.Message
 	// sent when a message's contents are updated and saved
-	msgMessageUpdated *models.Message
+	msgMessageUpdated *api.Message
 	// sent when a message is cloned, with the cloned message
-	msgMessageCloned *models.Message
+	msgMessageCloned *api.Message
 )

 type focusState int
@@ -77,14 +78,14 @@ type Model struct {

 	// app state
 	state           state // current overall status of the view
-	conversation    *models.Conversation
-	rootMessages    []models.Message
-	messages        []models.Message
+	conversation    *api.Conversation
+	rootMessages    []api.Message
+	messages        []api.Message
 	selectedMessage int
 	editorTarget    editorTarget
 	stopSignal      chan struct{}
-	replyChan       chan models.Message
-	replyChunkChan  chan api.Chunk
+	replyChan       chan api.Message
+	chatReplyChunks chan api.Chunk
 	persistence     bool // whether we will save new messages in the conversation

 	// ui state
@@ -111,12 +112,12 @@ func Chat(shared shared.Shared) Model {
 		Shared: shared,

 		state:        idle,
-		conversation: &models.Conversation{},
+		conversation: &api.Conversation{},
 		persistence:  true,

-		stopSignal:     make(chan struct{}),
-		replyChan:      make(chan models.Message),
-		replyChunkChan: make(chan api.Chunk),
+		stopSignal:      make(chan struct{}),
+		replyChan:       make(chan api.Message),
+		chatReplyChunks: make(chan api.Chunk),

 		wrap:            true,
 		selectedMessage: -1,
@@ -144,8 +145,8 @@ func Chat(shared shared.Shared) Model {

 	system := shared.Ctx.GetSystemPrompt()
 	if system != "" {
-		m.messages = []models.Message{{
-			Role:    models.MessageRoleSystem,
+		m.messages = []api.Message{{
+			Role:    api.MessageRoleSystem,
 			Content: system,
 		}}
 	}
@@ -166,6 +167,5 @@ func Chat(shared shared.Shared) Model {
 func (m Model) Init() tea.Cmd {
 	return tea.Batch(
 		m.waitForResponseChunk(),
-		m.waitForResponse(),
 	)
 }
--- a/pkg/tui/views/chat/conversation.go
+++ b/pkg/tui/views/chat/conversation.go
@@ -2,16 +2,18 @@ package chat

 import (
 	"context"
+	"errors"
 	"fmt"
 	"time"

+	"git.mlow.ca/mlow/lmcli/pkg/agent"
+	"git.mlow.ca/mlow/lmcli/pkg/api"
 	cmdutil "git.mlow.ca/mlow/lmcli/pkg/cmd/util"
-	models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
 	"git.mlow.ca/mlow/lmcli/pkg/tui/shared"
 	tea "github.com/charmbracelet/bubbletea"
 )

-func (m *Model) setMessage(i int, msg models.Message) {
+func (m *Model) setMessage(i int, msg api.Message) {
 	if i >= len(m.messages) {
 		panic("i out of range")
 	}
@@ -19,7 +21,7 @@ func (m *Model) setMessage(i int, msg models.Message) {
 	m.messageCache[i] = m.renderMessage(i)
 }

-func (m *Model) addMessage(msg models.Message) {
+func (m *Model) addMessage(msg api.Message) {
 	m.messages = append(m.messages, msg)
 	m.messageCache = append(m.messageCache, m.renderMessage(len(m.messages)-1))
 }
@@ -88,7 +90,7 @@ func (m *Model) generateConversationTitle() tea.Cmd {
 	}
 }

-func (m *Model) updateConversationTitle(conversation *models.Conversation) tea.Cmd {
+func (m *Model) updateConversationTitle(conversation *api.Conversation) tea.Cmd {
 	return func() tea.Msg {
 		err := m.Shared.Ctx.Store.UpdateConversation(conversation)
 		if err != nil {
@@ -101,7 +103,7 @@ func (m *Model) updateConversationTitle(conversation *models.Conversation) tea.C
 // Clones the given message (and its descendents). If selected is true, updates
 // either its parent's SelectedReply or its conversation's SelectedRoot to
 // point to the new clone
-func (m *Model) cloneMessage(message models.Message, selected bool) tea.Cmd {
+func (m *Model) cloneMessage(message api.Message, selected bool) tea.Cmd {
 	return func() tea.Msg {
 		msg, _, err := m.Ctx.Store.CloneBranch(message)
 		if err != nil {
@@ -123,7 +125,7 @@ func (m *Model) cloneMessage(message models.Message, selected bool) tea.Cmd {
 	}
 }

-func (m *Model) updateMessageContent(message *models.Message) tea.Cmd {
+func (m *Model) updateMessageContent(message *api.Message) tea.Cmd {
 	return func() tea.Msg {
 		err := m.Shared.Ctx.Store.UpdateMessage(message)
 		if err != nil {
@@ -133,7 +135,7 @@ func (m *Model) updateMessageContent(message *models.Message) tea.Cmd {
 	}
 }

-func cycleSelectedMessage(selected *models.Message, choices []models.Message, dir MessageCycleDirection) (*models.Message, error) {
+func cycleSelectedMessage(selected *api.Message, choices []api.Message, dir MessageCycleDirection) (*api.Message, error) {
 	currentIndex := -1
 	for i, reply := range choices {
 		if reply.ID == selected.ID {
@@ -158,7 +160,7 @@ func cycleSelectedMessage(selected *models.Message, choices []models.Message, di
 	return &choices[next], nil
 }

-func (m *Model) cycleSelectedRoot(conv *models.Conversation, dir MessageCycleDirection) tea.Cmd {
+func (m *Model) cycleSelectedRoot(conv *api.Conversation, dir MessageCycleDirection) tea.Cmd {
 	if len(m.rootMessages) < 2 {
 		return nil
 	}
@@ -178,7 +180,7 @@ func (m *Model) cycleSelectedRoot(conv *models.Conversation, dir MessageCycleDir
 	}
 }

-func (m *Model) cycleSelectedReply(message *models.Message, dir MessageCycleDirection) tea.Cmd {
+func (m *Model) cycleSelectedReply(message *api.Message, dir MessageCycleDirection) tea.Cmd {
 	if len(message.Replies) < 2 {
 		return nil
 	}
@@ -218,15 +220,12 @@ func (m *Model) persistConversation() tea.Cmd {
 		// else, we'll handle updating an existing conversation's messages
 		for i := range messages {
 			if messages[i].ID > 0 {
-				// message has an ID, update its contents
+				// message has an ID, update it
 				err := m.Shared.Ctx.Store.UpdateMessage(&messages[i])
 				if err != nil {
 					return shared.MsgError(err)
 				}
 			} else if i > 0 {
-				if messages[i].Content == "" {
-					continue
-				}
 				// messages is new, so add it as a reply to previous message
 				saved, err := m.Shared.Ctx.Store.Reply(&messages[i-1], messages[i])
 				if err != nil {
@@ -243,13 +242,23 @@ func (m *Model) persistConversation() tea.Cmd {
 	}
 }

+func (m *Model) executeToolCalls(toolCalls []api.ToolCall) tea.Cmd {
+	return func() tea.Msg {
+		results, err := agent.ExecuteToolCalls(toolCalls, m.Ctx.EnabledTools)
+		if err != nil {
+			return shared.MsgError(err)
+		}
+		return msgToolResults(results)
+	}
+}
+
 func (m *Model) promptLLM() tea.Cmd {
 	m.state = pendingResponse
 	m.replyCursor.Blink = false

-	m.tokenCount = 0
 	m.startTime = time.Now()
 	m.elapsed = 0
+	m.tokenCount = 0

 	return func() tea.Msg {
 		model, provider, err := m.Shared.Ctx.GetModelProvider(*m.Shared.Ctx.Config.Defaults.Model)
@@ -257,36 +266,34 @@ func (m *Model) promptLLM() tea.Cmd {
 			return shared.MsgError(err)
 		}

-		requestParams := models.RequestParameters{
+		requestParams := api.RequestParameters{
 			Model:       model,
 			MaxTokens:   *m.Shared.Ctx.Config.Defaults.MaxTokens,
 			Temperature: *m.Shared.Ctx.Config.Defaults.Temperature,
 			ToolBag:     m.Shared.Ctx.EnabledTools,
 		}

-		replyHandler := func(msg models.Message) {
-			m.replyChan <- msg
-		}
-
 		ctx, cancel := context.WithCancel(context.Background())

-		canceled := false
 		go func() {
 			select {
 			case <-m.stopSignal:
-				canceled = true
 				cancel()
 			}
 		}()

 		resp, err := provider.CreateChatCompletionStream(
-			ctx, requestParams, m.messages, replyHandler, m.replyChunkChan,
+			ctx, requestParams, m.messages, m.chatReplyChunks,
 		)

-		if err != nil && !canceled {
-			return msgResponseError(err)
+		if errors.Is(err, context.Canceled) {
+			return msgChatResponseCanceled(struct{}{})
 		}

-		return msgResponseEnd(resp)
+		if err != nil {
+			return msgChatResponseError(err)
+		}
+
+		return msgChatResponse(resp)
 	}
 }
--- a/pkg/tui/views/chat/input.go
+++ b/pkg/tui/views/chat/input.go
@@ -4,7 +4,7 @@ import (
 	"fmt"
 	"strings"

-	models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
+	"git.mlow.ca/mlow/lmcli/pkg/api"
 	"git.mlow.ca/mlow/lmcli/pkg/tui/shared"
 	tuiutil "git.mlow.ca/mlow/lmcli/pkg/tui/util"
 	tea "github.com/charmbracelet/bubbletea"
@@ -150,12 +150,12 @@ func (m *Model) handleInputKey(msg tea.KeyMsg) (bool, tea.Cmd) {
 			return true, nil
 		}

-		if len(m.messages) > 0 && m.messages[len(m.messages)-1].Role == models.MessageRoleUser {
+		if len(m.messages) > 0 && m.messages[len(m.messages)-1].Role == api.MessageRoleUser {
 			return true, shared.WrapError(fmt.Errorf("Can't reply to a user message"))
 		}

-		m.addMessage(models.Message{
-			Role:    models.MessageRoleUser,
+		m.addMessage(api.Message{
+			Role:    api.MessageRoleUser,
 			Content: input,
 		})

--- a/pkg/tui/views/chat/update.go
+++ b/pkg/tui/views/chat/update.go
@@ -4,7 +4,7 @@ import (
 	"strings"
 	"time"

-	models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
+	"git.mlow.ca/mlow/lmcli/pkg/api"
 	"git.mlow.ca/mlow/lmcli/pkg/tui/shared"
 	tuiutil "git.mlow.ca/mlow/lmcli/pkg/tui/util"
 	"github.com/charmbracelet/bubbles/cursor"
@@ -21,15 +21,9 @@ func (m *Model) HandleResize(width, height int) {
 	}
 }

-func (m *Model) waitForResponse() tea.Cmd {
-	return func() tea.Msg {
-		return msgResponse(<-m.replyChan)
-	}
-}
-
 func (m *Model) waitForResponseChunk() tea.Cmd {
 	return func() tea.Msg {
-		return msgResponseChunk(<-m.replyChunkChan)
+		return msgChatResponseChunk(<-m.chatReplyChunks)
 	}
 }

@@ -48,7 +42,7 @@ func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) {

 			if m.conversation.ShortName.String != m.Shared.Values.ConvShortname {
 				// clear existing messages if we're loading a new conversation
-				m.messages = []models.Message{}
+				m.messages = []api.Message{}
 				m.selectedMessage = 0
 			}
 		}
@@ -87,7 +81,7 @@ func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) {
 		}
 		m.rebuildMessageCache()
 		m.updateContent()
-	case msgResponseChunk:
+	case msgChatResponseChunk:
 		cmds = append(cmds, m.waitForResponseChunk()) // wait for the next chunk

 		if msg.Content == "" {
@@ -100,8 +94,8 @@ func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) {
 			m.setMessageContents(last, m.messages[last].Content+msg.Content)
 		} else {
 			// use chunk in new message
-			m.addMessage(models.Message{
-				Role:    models.MessageRoleAssistant,
+			m.addMessage(api.Message{
+				Role:    api.MessageRoleAssistant,
 				Content: msg.Content,
 			})
 		}
@@ -113,10 +107,10 @@ func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) {

 		m.tokenCount += msg.TokenCount
 		m.elapsed = time.Now().Sub(m.startTime)
-	case msgResponse:
-		cmds = append(cmds, m.waitForResponse()) // wait for the next response
+	case msgChatResponse:
+		m.state = idle

-		reply := models.Message(msg)
+		reply := (*api.Message)(msg)
 		reply.Content = strings.TrimSpace(reply.Content)

 		last := len(m.messages) - 1
@@ -124,11 +118,18 @@ func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) {
 			panic("Unexpected empty messages handling msgAssistantReply")
 		}

-		if reply.Role.IsAssistant() && m.messages[last].Role.IsAssistant() {
-			// this was a continuation, so replace the previous message with the completed reply
-			m.setMessage(last, reply)
+		if m.messages[last].Role.IsAssistant() {
+			// TODO: handle continuations gracefully - some models support them well, others fail horribly.
+			m.setMessage(last, *reply)
 		} else {
-			m.addMessage(reply)
+			m.addMessage(*reply)
+		}
+
+		switch reply.Role {
+		case api.MessageRoleToolCall:
+			// TODO: user confirmation before execution
+			// m.state = waitingForConfirmation
+			cmds = append(cmds, m.executeToolCalls(reply.ToolCalls))
 		}

 		if m.persistence {
@@ -140,17 +141,32 @@ func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) {
 		}

 		m.updateContent()
-	case msgResponseEnd:
+	case msgChatResponseCanceled:
 		m.state = idle
-		last := len(m.messages) - 1
-		if last < 0 {
-			panic("Unexpected empty messages handling msgResponseEnd")
-		}
-		m.setMessageContents(last, strings.TrimSpace(m.messages[last].Content))
 		m.updateContent()
-	case msgResponseError:
+	case msgChatResponseError:
 		m.state = idle
 		m.Shared.Err = error(msg)
+		m.updateContent()
+	case msgToolResults:
+		last := len(m.messages) - 1
+		if last < 0 {
+			panic("Unexpected empty messages handling msgAssistantReply")
+		}
+
+		if m.messages[last].Role != api.MessageRoleToolCall {
+			panic("Previous message not a tool call, unexpected")
+		}
+
+		m.addMessage(api.Message{
+			Role:        api.MessageRoleToolResult,
+			ToolResults: api.ToolResults(msg),
+		})
+
+		if m.persistence {
+			cmds = append(cmds, m.persistConversation())
+		}
+
 		m.updateContent()
 	case msgConversationTitleGenerated:
 		title := string(msg)
@@ -167,7 +183,7 @@ func (m Model) Update(msg tea.Msg) (Model, tea.Cmd) {
 		m.conversation = msg.conversation
 		m.messages = msg.messages
 		if msg.isNew {
-			m.rootMessages = []models.Message{m.messages[0]}
+			m.rootMessages = []api.Message{m.messages[0]}
 		}
 		m.rebuildMessageCache()
 		m.updateContent()
--- a/pkg/tui/views/chat/view.go
+++ b/pkg/tui/views/chat/view.go
@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"strings"

-	models "git.mlow.ca/mlow/lmcli/pkg/lmcli/model"
+	"git.mlow.ca/mlow/lmcli/pkg/api"
 	"git.mlow.ca/mlow/lmcli/pkg/tui/styles"
 	tuiutil "git.mlow.ca/mlow/lmcli/pkg/tui/util"
 	"github.com/charmbracelet/lipgloss"
@@ -63,22 +63,22 @@ func (m Model) View() string {
 	return lipgloss.JoinVertical(lipgloss.Left, sections...)
 }

-func (m *Model) renderMessageHeading(i int, message *models.Message) string {
+func (m *Model) renderMessageHeading(i int, message *api.Message) string {
 	icon := ""
 	friendly := message.Role.FriendlyRole()
 	style := lipgloss.NewStyle().Faint(true).Bold(true)

 	switch message.Role {
-	case models.MessageRoleSystem:
+	case api.MessageRoleSystem:
 		icon = "⚙️"
-	case models.MessageRoleUser:
+	case api.MessageRoleUser:
 		style = userStyle
-	case models.MessageRoleAssistant:
+	case api.MessageRoleAssistant:
 		style = assistantStyle
-	case models.MessageRoleToolCall:
+	case api.MessageRoleToolCall:
 		style = assistantStyle
-		friendly = models.MessageRoleAssistant.FriendlyRole()
-	case models.MessageRoleToolResult:
+		friendly = api.MessageRoleAssistant.FriendlyRole()
+	case api.MessageRoleToolResult:
 		icon = "🔧"
 	}

@@ -139,21 +139,21 @@ func (m *Model) renderMessage(i int) string {
 	}

 	// Show the assistant's cursor
-	if m.state == pendingResponse && i == len(m.messages)-1 && msg.Role == models.MessageRoleAssistant {
+	if m.state == pendingResponse && i == len(m.messages)-1 && msg.Role == api.MessageRoleAssistant {
 		sb.WriteString(m.replyCursor.View())
 	}

 	// Write tool call info
 	var toolString string
 	switch msg.Role {
-	case models.MessageRoleToolCall:
+	case api.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:
+	case api.MessageRoleToolResult:
 		if !m.showToolResults {
 			break
 		}
@@ -221,11 +221,11 @@ func (m *Model) conversationMessagesView() string {
 		m.messageOffsets[i] = lineCnt

 		switch message.Role {
-		case models.MessageRoleToolCall:
+		case api.MessageRoleToolCall:
 			if !m.showToolResults && message.Content == "" {
 				continue
 			}
-		case models.MessageRoleToolResult:
+		case api.MessageRoleToolResult:
 			if !m.showToolResults {
 				continue
 			}
@@ -251,9 +251,9 @@ func (m *Model) conversationMessagesView() string {
 	}

 	// Render a placeholder for the incoming assistant reply
-	if m.state == pendingResponse && (len(m.messages) == 0 || m.messages[len(m.messages)-1].Role != models.MessageRoleAssistant) {
-		heading := m.renderMessageHeading(-1, &models.Message{
-			Role: models.MessageRoleAssistant,
+	if m.state == pendingResponse && (len(m.messages) == 0 || m.messages[len(m.messages)-1].Role != api.MessageRoleAssistant) {
+		heading := m.renderMessageHeading(-1, &api.Message{
+			Role: api.MessageRoleAssistant,
 		})
 		sb.WriteString(heading)
 		sb.WriteString("\n")