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 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{true, 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{false, 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) } }