lmcli/pkg/tui/views/chat/conversation.go

package chat

import (
	"context"
	"fmt"
	"time"

	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) {
	if i >= len(m.messages) {
		panic("i out of range")
	}
	m.messages[i] = msg
	m.messageCache[i] = m.renderMessage(i)
}

func (m *Model) addMessage(msg models.Message) {
	m.messages = append(m.messages, msg)
	m.messageCache = append(m.messageCache, m.renderMessage(len(m.messages)-1))
}

func (m *Model) setMessageContents(i int, content string) {
	if i >= len(m.messages) {
		panic("i out of range")
	}
	m.messages[i].Content = content
	m.messageCache[i] = m.renderMessage(i)
}

func (m *Model) rebuildMessageCache() {
	m.messageCache = make([]string, len(m.messages))
	for i := range m.messages {
		m.messageCache[i] = m.renderMessage(i)
	}
}

func (m *Model) updateContent() {
	atBottom := m.content.AtBottom()
	m.content.SetContent(m.conversationMessagesView())
	if atBottom {
		// if we were at bottom before the update, scroll with the output
		m.content.GotoBottom()
	}
}

func (m *Model) loadConversation(shortname string) tea.Cmd {
	return func() tea.Msg {
		if shortname == "" {
			return nil
		}
		c, err := m.Shared.Ctx.Store.ConversationByShortName(shortname)
		if err != nil {
			return shared.MsgError(fmt.Errorf("Could not lookup conversation: %v", err))
		}
		if c.ID == 0 {
			return shared.MsgError(fmt.Errorf("Conversation not found: %s", shortname))
		}
		rootMessages, err := m.Shared.Ctx.Store.RootMessages(c.ID)
		if err != nil {
			return shared.MsgError(fmt.Errorf("Could not load conversation root messages: %v\n", err))
		}
		return msgConversationLoaded{c, rootMessages}
	}
}

func (m *Model) loadConversationMessages() tea.Cmd {
	return func() tea.Msg {
		messages, err := m.Shared.Ctx.Store.PathToLeaf(m.conversation.SelectedRoot)
		if err != nil {
			return shared.MsgError(fmt.Errorf("Could not load conversation messages: %v\n", err))
		}
		return msgMessagesLoaded(messages)
	}
}

func (m *Model) generateConversationTitle() tea.Cmd {
	return func() tea.Msg {
		title, err := cmdutil.GenerateTitle(m.Shared.Ctx, m.messages)
		if err != nil {
			return shared.MsgError(err)
		}
		return msgConversationTitleGenerated(title)
	}
}

func (m *Model) updateConversationTitle(conversation *models.Conversation) tea.Cmd {
	return func() tea.Msg {
		err := m.Shared.Ctx.Store.UpdateConversation(conversation)
		if err != nil {
			return shared.WrapError(err)
		}
		return nil
	}
}

// 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 {
	return func() tea.Msg {
		msg, _, err := m.Ctx.Store.CloneBranch(message)
		if err != nil {
			return shared.WrapError(fmt.Errorf("Could not clone message: %v", err))
		}
		if selected {
			if msg.Parent == nil {
				msg.Conversation.SelectedRoot = msg
				err = m.Shared.Ctx.Store.UpdateConversation(&msg.Conversation)
			} else {
				msg.Parent.SelectedReply = msg
				err = m.Shared.Ctx.Store.UpdateMessage(msg.Parent)
			}
			if err != nil {
				return shared.WrapError(fmt.Errorf("Could not update selected message: %v", err))
			}
		}
		return msgMessageCloned(msg)
	}
}

func (m *Model) updateMessageContent(message *models.Message) tea.Cmd {
	return func() tea.Msg {
		err := m.Shared.Ctx.Store.UpdateMessage(message)
		if err != nil {
			return shared.WrapError(fmt.Errorf("Could not update message: %v", err))
		}
		return msgMessageUpdated(message)
	}
}

func cycleSelectedMessage(selected *models.Message, choices []models.Message, dir MessageCycleDirection) (*models.Message, error) {
	currentIndex := -1
	for i, reply := range choices {
		if reply.ID == selected.ID {
			currentIndex = i
			break
		}
	}

	if currentIndex < 0 {
		// this should probably be an assert
		return nil, fmt.Errorf("Selected message %d not found in choices, this is a bug", selected.ID)
	}

	var next int
	if dir == CyclePrev {
		// Wrap around to the last reply if at the beginning
		next = (currentIndex - 1 + len(choices)) % len(choices)
	} else {
		// Wrap around to the first reply if at the end
		next = (currentIndex + 1) % len(choices)
	}
	return &choices[next], nil
}

func (m *Model) cycleSelectedRoot(conv *models.Conversation, dir MessageCycleDirection) tea.Cmd {
	if len(m.rootMessages) < 2 {
		return nil
	}

	return func() tea.Msg {
		nextRoot, err := cycleSelectedMessage(conv.SelectedRoot, m.rootMessages, dir)
		if err != nil {
			return shared.WrapError(err)
		}

		conv.SelectedRoot = nextRoot
		err = m.Shared.Ctx.Store.UpdateConversation(conv)
		if err != nil {
			return shared.WrapError(fmt.Errorf("Could not update conversation SelectedRoot: %v", err))
		}
		return msgSelectedRootCycled(nextRoot)
	}
}

func (m *Model) cycleSelectedReply(message *models.Message, dir MessageCycleDirection) tea.Cmd {
	if len(message.Replies) < 2 {
		return nil
	}

	return func() tea.Msg {
		nextReply, err := cycleSelectedMessage(message.SelectedReply, message.Replies, dir)
		if err != nil {
			return shared.WrapError(err)
		}

		message.SelectedReply = nextReply
		err = m.Shared.Ctx.Store.UpdateMessage(message)
		if err != nil {
			return shared.WrapError(fmt.Errorf("Could not update message SelectedReply: %v", err))
		}
		return msgSelectedReplyCycled(nextReply)
	}
}

func (m *Model) persistConversation() tea.Cmd {
	conversation := m.conversation
	messages := m.messages

	var err error
	if m.conversation.ID == 0 {
		return func() tea.Msg {
			// Start a new conversation with all messages so far
			conversation, messages, err = m.Shared.Ctx.Store.StartConversation(messages...)
			if err != nil {
				return shared.MsgError(fmt.Errorf("Could not start new conversation: %v", err))
			}
			return msgConversationPersisted{conversation, messages}
		}
	}

	return func() tea.Msg {
		// 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
				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 {
					return shared.MsgError(err)
				}
				messages[i] = saved[0]
			} else {
				// message has no id and no previous messages to add it to
				// this shouldn't happen?
				return fmt.Errorf("Error: no messages to reply to")
			}
		}
		return msgConversationPersisted{conversation, messages}
	}
}

func (m *Model) promptLLM() tea.Cmd {
	m.state = pendingResponse
	m.replyCursor.Blink = false

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

	return func() tea.Msg {
		model, provider, err := m.Shared.Ctx.GetModelProvider(*m.Shared.Ctx.Config.Defaults.Model)
		if err != nil {
			return shared.MsgError(err)
		}

		requestParams := models.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,
		)

		if err != nil && !canceled {
			return msgResponseError(err)
		}

		return msgResponseEnd(resp)
	}
}