Radu Macocian (admac)
577 lines
15 KiB
Go
577 lines
15 KiB
Go
package ui
|
||
|
||
import (
|
||
"context"
|
||
"fmt"
|
||
"log/slog"
|
||
"strings"
|
||
|
||
"charm.land/bubbles/v2/cursor"
|
||
"charm.land/bubbles/v2/spinner"
|
||
"charm.land/bubbles/v2/textarea"
|
||
"charm.land/bubbles/v2/viewport"
|
||
tea "charm.land/bubbletea/v2"
|
||
"charm.land/lipgloss/v2"
|
||
"git.estatecloud.org/radumaco/souvenir/config"
|
||
"git.estatecloud.org/radumaco/souvenir/db/history"
|
||
llm "git.estatecloud.org/radumaco/souvenir/llm"
|
||
"git.estatecloud.org/radumaco/souvenir/model"
|
||
)
|
||
|
||
type focusState int
|
||
|
||
const (
|
||
focusChat focusState = iota
|
||
focusModels
|
||
focusHistory
|
||
)
|
||
|
||
type uiModel struct {
|
||
logger slog.Logger
|
||
viewport viewport.Model
|
||
spinner spinner.Model
|
||
conversation model.Conversation
|
||
textarea textarea.Model
|
||
focus focusState
|
||
picker picker
|
||
thinkingBuffer strings.Builder
|
||
answerBuffer strings.Builder
|
||
senderStyle lipgloss.Style
|
||
agentStyle lipgloss.Style
|
||
errorStyle lipgloss.Style
|
||
statusStyle lipgloss.Style
|
||
client llm.ChatClient
|
||
statusMessage string
|
||
commands commandList
|
||
history history.DbClient
|
||
streamCh <-chan llm.StreamEvent
|
||
ctx context.Context
|
||
width int
|
||
height int
|
||
waiting bool
|
||
err error
|
||
}
|
||
|
||
type streamEventMessage struct {
|
||
event llm.StreamEvent
|
||
}
|
||
|
||
type streamClosedMessage struct {
|
||
}
|
||
|
||
type conversationSavedMessage struct {
|
||
conversation model.Conversation
|
||
err error
|
||
}
|
||
|
||
type conversationRenamedMessage struct {
|
||
title string
|
||
summary string
|
||
err error
|
||
}
|
||
|
||
func InitialModel(ctx context.Context, config config.Config, client history.DbClient) uiModel {
|
||
ta := textarea.New()
|
||
ta.Placeholder = "Send a message..."
|
||
ta.SetVirtualCursor(false)
|
||
ta.Focus()
|
||
|
||
ta.Prompt = "｜"
|
||
ta.CharLimit = 280
|
||
|
||
ta.SetWidth(30)
|
||
ta.SetHeight(3)
|
||
|
||
s := ta.Styles()
|
||
s.Focused.CursorLine = lipgloss.NewStyle()
|
||
ta.SetStyles(s)
|
||
|
||
ta.ShowLineNumbers = false
|
||
|
||
vp := viewport.New(viewport.WithWidth(30), viewport.WithHeight(5))
|
||
vp.SetContent(renderLanding(vp.Width(), vp.Height()))
|
||
vp.KeyMap.Left.SetEnabled(false)
|
||
vp.KeyMap.Right.SetEnabled(false)
|
||
|
||
ta.KeyMap.InsertNewline.SetEnabled(false)
|
||
|
||
sp := spinner.New()
|
||
sp.Spinner = spinner.Monkey
|
||
|
||
ui := uiModel{
|
||
textarea: ta,
|
||
conversation: model.Conversation{},
|
||
viewport: vp,
|
||
spinner: sp,
|
||
senderStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("5")),
|
||
errorStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("9")),
|
||
agentStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("86")),
|
||
client: *llm.NewLLMClient(config),
|
||
focus: focusChat,
|
||
history: client,
|
||
ctx: ctx,
|
||
logger: *slog.Default().With("Component", "TUI"),
|
||
err: nil,
|
||
}
|
||
ui.picker = newPicker()
|
||
ui.commands = newCommandList(30)
|
||
ui.commands.setAvailable(ui.buildCommands())
|
||
return ui
|
||
}
|
||
|
||
func (m uiModel) Init() tea.Cmd {
|
||
return textarea.Blink
|
||
}
|
||
|
||
func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
|
||
m.logger.Debug("Received message", "msg", msg)
|
||
if m.waiting {
|
||
m.spinner, _ = m.spinner.Update(msg)
|
||
}
|
||
|
||
switch msg := msg.(type) {
|
||
case tea.WindowSizeMsg:
|
||
m.viewport.SetWidth(msg.Width * 80 / 100)
|
||
m.textarea.SetWidth(msg.Width * 80 / 100)
|
||
m.viewport.SetHeight(msg.Height - lipgloss.Height(m.textarea.View()) - 1)
|
||
m.width = msg.Width
|
||
m.height = msg.Height
|
||
m.picker.SetSize(msg.Width, msg.Height)
|
||
m.commands.list.SetSize(msg.Width*80/100, commandDropdownHeight)
|
||
if m.commands.open {
|
||
m.viewport.SetHeight(m.viewport.Height() - commandDropdownHeight)
|
||
}
|
||
|
||
var messages = m.renderMessages()
|
||
|
||
if len(m.conversation.Messages) > 0 {
|
||
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
|
||
} else {
|
||
m.viewport.SetContent(renderLanding(m.viewport.Width(), m.viewport.Height()))
|
||
}
|
||
m.logger.Debug("New sizes",
|
||
"vp width", m.viewport.Width(), "vp height", m.viewport.Height(),
|
||
"ta width", m.textarea.Width(), "ta height", m.textarea.Height())
|
||
m.viewport.GotoBottom()
|
||
|
||
case tea.KeyPressMsg:
|
||
if m.focus != focusChat {
|
||
picker, cmd := m.picker.Update(msg)
|
||
m.picker = picker
|
||
return m, cmd
|
||
}
|
||
switch msg.String() {
|
||
case "esc":
|
||
if m.commands.open {
|
||
m.closeCommands()
|
||
return m, nil
|
||
}
|
||
if m.focus != focusChat {
|
||
m.focus = focusChat
|
||
}
|
||
return m, nil
|
||
case "ctrl+c":
|
||
fmt.Println(m.textarea.Value())
|
||
return m, tea.Quit
|
||
case "shift+enter", "ctrl+j", "alt+enter":
|
||
m.textarea.InsertRune('\n')
|
||
return m, nil
|
||
case "up", "down":
|
||
if m.commands.open {
|
||
var cmd tea.Cmd
|
||
m.commands.list, cmd = m.commands.list.Update(msg)
|
||
return m, cmd
|
||
}
|
||
var cmd tea.Cmd
|
||
m.textarea, cmd = m.textarea.Update(msg)
|
||
return m, cmd
|
||
case "tab":
|
||
if m.commands.open {
|
||
if sel, ok := m.commands.selectedItem(); ok {
|
||
m.textarea.SetValue("/" + sel.name + " ")
|
||
}
|
||
m.closeCommands()
|
||
return m, nil
|
||
}
|
||
var cmd tea.Cmd
|
||
m.textarea, cmd = m.textarea.Update(msg)
|
||
return m, cmd
|
||
case "enter":
|
||
m.setStatusMessage("")
|
||
input := m.textarea.Value()
|
||
|
||
if m.commands.open {
|
||
if sel, ok := m.commands.selectedItem(); ok {
|
||
m.textarea.SetValue("")
|
||
m.closeCommands()
|
||
mm, cmd := sel.handler(m)
|
||
return mm, cmd
|
||
}
|
||
}
|
||
|
||
if after, ok := strings.CutPrefix(input, "/"); ok {
|
||
name := after
|
||
for _, c := range m.commands.available {
|
||
if c.name == name {
|
||
m.textarea.SetValue("")
|
||
m.closeCommands()
|
||
mm, cmd := c.handler(m)
|
||
return mm, cmd
|
||
}
|
||
}
|
||
}
|
||
|
||
m.textarea.SetValue("")
|
||
m.closeCommands()
|
||
m.logger.Debug("Received enter. Creating new message", "message", input)
|
||
|
||
m.conversation.Messages = append(m.conversation.Messages, model.Message{
|
||
Content: input,
|
||
Role: "user",
|
||
Seq: len(m.conversation.Messages) + 1,
|
||
})
|
||
var messages = m.renderMessages()
|
||
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
|
||
m.textarea.Reset()
|
||
m.viewport.GotoBottom()
|
||
m.startWait()
|
||
return m, tea.Batch(m.callAgent(m.conversation.Messages), m.saveConversation(), m.spinner.Tick)
|
||
default:
|
||
var cmd tea.Cmd
|
||
m.textarea, cmd = m.textarea.Update(msg)
|
||
v := m.textarea.Value()
|
||
if strings.HasPrefix(v, "/") {
|
||
m.openCommands()
|
||
} else {
|
||
m.closeCommands()
|
||
}
|
||
return m, cmd
|
||
}
|
||
|
||
case cursor.BlinkMsg:
|
||
m.logger.Debug("Received BlinkMsg")
|
||
var cmd tea.Cmd
|
||
m.textarea, cmd = m.textarea.Update(msg)
|
||
return m, cmd
|
||
|
||
case conversationRenamedMessage:
|
||
m.logger.Debug("Received messageRenamedMessage")
|
||
m.startWait()
|
||
if msg.err != nil {
|
||
m.logger.Error("There was an error", "error", msg.err)
|
||
m.setErrorMessage(msg.err.Error())
|
||
}
|
||
m.setStatusMessage("Conversation renamed to: " + msg.title)
|
||
m.conversation.Title = msg.title
|
||
m.conversation.Summary = msg.summary
|
||
return m, nil
|
||
|
||
case streamEventMessage:
|
||
m.logger.Debug("Received streamEventMessage")
|
||
if msg.event.Err != nil {
|
||
m.logger.Error("There was an error mid stream", "error", msg.event.Err)
|
||
}
|
||
|
||
if msg.event.Reasoning != "" {
|
||
m.thinkingBuffer.WriteString(msg.event.Reasoning)
|
||
}
|
||
|
||
if msg.event.Content != "" {
|
||
m.answerBuffer.WriteString(msg.event.Content)
|
||
}
|
||
|
||
if msg.event.Done {
|
||
for _, r := range msg.event.Messages {
|
||
r.Seq = len(m.conversation.Messages) + 1
|
||
m.logger.Debug("Appending response message", "message", r)
|
||
m.conversation.Messages = append(m.conversation.Messages, r)
|
||
}
|
||
|
||
var messages = m.renderMessages()
|
||
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
|
||
m.viewport.GotoBottom()
|
||
return m, m.saveConversation()
|
||
}
|
||
|
||
return m, waitForEvent(m.streamCh)
|
||
|
||
case streamClosedMessage:
|
||
m.logger.Debug("Received streamClosedMessage")
|
||
m.stopWait()
|
||
m.answerBuffer.Reset()
|
||
m.thinkingBuffer.Reset()
|
||
|
||
return m, nil
|
||
|
||
case pickerChosenMsg:
|
||
switch m.focus {
|
||
case focusModels:
|
||
m.logger.Debug("Received pickerChosenMsg", "model", msg.id)
|
||
m.client.Cfg.Llm.Model = msg.id
|
||
m.stopWait()
|
||
m.focus = focusChat
|
||
case focusHistory:
|
||
m.logger.Debug("Received pickerChosenMsg", "history", msg.id)
|
||
conv, err := m.history.GetConversation(m.ctx, msg.id)
|
||
if err != nil {
|
||
m.logger.Error("Could not retrieve conversation", "id", msg.id)
|
||
}
|
||
m.conversation = conv
|
||
m.stopWait()
|
||
m.focus = focusChat
|
||
var messages = m.renderMessages()
|
||
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
|
||
m.viewport.GotoBottom()
|
||
}
|
||
|
||
case pickerDismissMsg:
|
||
m.logger.Debug("Received pickerDismissMsg")
|
||
m.stopWait()
|
||
m.focus = focusChat
|
||
|
||
case conversationSavedMessage:
|
||
m.logger.Debug("Received conversationSavedMessage", "conversation", msg.conversation)
|
||
if msg.err != nil {
|
||
m.logger.Error("Could not save message", "error", msg.err)
|
||
m.setErrorMessage(msg.err.Error())
|
||
return m, nil
|
||
}
|
||
m.conversation = msg.conversation
|
||
}
|
||
|
||
if m.focus != focusChat {
|
||
picker, cmd := m.picker.Update(msg)
|
||
m.picker = picker
|
||
return m, cmd
|
||
}
|
||
|
||
return m, nil
|
||
}
|
||
|
||
func (m uiModel) View() tea.View {
|
||
if m.focus != focusChat {
|
||
return tea.NewView(lipgloss.Place(m.width, m.height,
|
||
lipgloss.Center, lipgloss.Center,
|
||
m.picker.View()))
|
||
}
|
||
|
||
parts := []string{m.viewport.View()}
|
||
if m.statusMessage != "" {
|
||
parts = append(parts, m.statusMessage)
|
||
}
|
||
if m.waiting {
|
||
parts = append(parts, m.spinner.View())
|
||
}
|
||
if m.commands.open {
|
||
parts = append(parts, m.commands.view())
|
||
}
|
||
parts = append(parts, m.textarea.View())
|
||
ui := lipgloss.JoinVertical(
|
||
lipgloss.Left,
|
||
parts...)
|
||
|
||
c := m.textarea.Cursor()
|
||
if c != nil {
|
||
c.Y += lipgloss.Height(m.viewport.View())
|
||
if m.statusMessage != "" {
|
||
c.Y++
|
||
}
|
||
if m.commands.open {
|
||
c.Y += lipgloss.Height(m.commands.view())
|
||
}
|
||
if m.waiting {
|
||
c.Y++
|
||
}
|
||
gap := m.width - lipgloss.Width(ui)
|
||
c.X += max(gap/2, 0)
|
||
}
|
||
view := tea.NewView(lipgloss.PlaceHorizontal(m.width, lipgloss.Center, ui))
|
||
view.Cursor = c
|
||
view.AltScreen = true
|
||
return view
|
||
}
|
||
|
||
func (m uiModel) callAgent(messages []model.Message) tea.Cmd {
|
||
m.logger.Debug("Calling agent", "messages", messages)
|
||
return func() tea.Msg {
|
||
resp, err := m.client.QueryStream(messages)
|
||
if err != nil {
|
||
m.logger.Error("Error while calling stream query", "error", err)
|
||
return streamClosedMessage{}
|
||
}
|
||
return waitForEvent(resp)
|
||
}
|
||
}
|
||
|
||
func waitForEvent(ch <-chan llm.StreamEvent) tea.Cmd {
|
||
return func() tea.Msg {
|
||
ev, ok := <-ch
|
||
if !ok {
|
||
return streamClosedMessage{}
|
||
}
|
||
return streamEventMessage{ev}
|
||
}
|
||
}
|
||
|
||
func (m uiModel) saveConversation() tea.Cmd {
|
||
m.logger.Debug("Saving conversation", "conversation", m.conversation)
|
||
return func() tea.Msg {
|
||
conv, err := m.history.SaveConversation(m.ctx, m.conversation)
|
||
return conversationSavedMessage{conversation: conv, err: err}
|
||
}
|
||
}
|
||
|
||
func (m uiModel) renameConversation() tea.Cmd {
|
||
m.logger.Debug("Renaming conversation", "conversation", m.conversation)
|
||
return func() tea.Msg {
|
||
meta, err := m.client.Rename(m.conversation.Messages)
|
||
if err != nil {
|
||
return conversationRenamedMessage{title: meta.Title, summary: meta.Summary, err: err}
|
||
}
|
||
m.conversation.Title = meta.Title
|
||
m.conversation.Summary = meta.Summary
|
||
m.conversation, err = m.history.SaveConversation(m.ctx, m.conversation)
|
||
|
||
return conversationRenamedMessage{title: m.conversation.Title, summary: m.conversation.Summary, err: err}
|
||
}
|
||
}
|
||
|
||
func (m uiModel) getModels() tea.Cmd {
|
||
m.logger.Debug("Querying models")
|
||
return func() tea.Msg {
|
||
resp, err := m.client.Models()
|
||
if err != nil {
|
||
m.logger.Error("Could not fetch models", "error", err)
|
||
}
|
||
m.logger.Debug("Received models from llm", "models", resp)
|
||
items := []modelItem{}
|
||
for _, model := range resp {
|
||
items = append(items, modelItem{name: model, description: model})
|
||
}
|
||
return modelsLoadedMsg{models: items, title: "Choose a model"}
|
||
}
|
||
}
|
||
|
||
func (m uiModel) getHistory() tea.Cmd {
|
||
m.logger.Debug("Querying history")
|
||
return func() tea.Msg {
|
||
resp, err := m.history.GetConversations(m.ctx)
|
||
if err != nil {
|
||
m.logger.Error("Could not fetch history", "error", err)
|
||
}
|
||
m.logger.Debug("Received conversations from psql", "conversations", resp)
|
||
items := []modelItem{}
|
||
for _, conv := range resp {
|
||
name := conv.Id
|
||
if conv.Title != "" {
|
||
name = conv.Title
|
||
}
|
||
items = append(items, modelItem{name: name, description: conv.Summary})
|
||
}
|
||
return modelsLoadedMsg{models: items, title: "Select a conversation to continue from where you left off"}
|
||
}
|
||
}
|
||
|
||
func (m *uiModel) renderMessages() string {
|
||
var result strings.Builder
|
||
var width = m.viewport.Width()
|
||
for _, message := range m.conversation.Messages {
|
||
switch message.Role {
|
||
case "user":
|
||
var line strings.Builder
|
||
line.WriteString(m.senderStyle.Render("You: "))
|
||
line.WriteString(message.Content)
|
||
result.WriteString(lipgloss.PlaceHorizontal(width, lipgloss.Right, line.String()))
|
||
result.WriteString("\n")
|
||
|
||
var sepStyle lipgloss.Style = lipgloss.NewStyle().
|
||
Foreground(lipgloss.Color("240")).
|
||
MarginTop(1).
|
||
MarginBottom(1)
|
||
sep := sepStyle.Render(strings.Repeat("┈", width))
|
||
result.WriteString(sep)
|
||
result.WriteString("\n")
|
||
case "assistant":
|
||
result.WriteString(m.agentStyle.Render("Agent: "))
|
||
result.WriteString(message.Content)
|
||
result.WriteString("\n")
|
||
|
||
var sepStyle lipgloss.Style = lipgloss.NewStyle().
|
||
Foreground(lipgloss.Color("240")).
|
||
MarginTop(1).
|
||
MarginBottom(1)
|
||
sep := sepStyle.Render(strings.Repeat("┈", width))
|
||
result.WriteString(sep)
|
||
result.WriteString("\n")
|
||
}
|
||
}
|
||
if m.answerBuffer.String() != "" {
|
||
result.WriteString(m.agentStyle.Render("Agent: "))
|
||
result.WriteString(m.answerBuffer.String())
|
||
|
||
result.WriteString("\n")
|
||
var sepStyle lipgloss.Style = lipgloss.NewStyle().
|
||
Foreground(lipgloss.Color("240")).
|
||
MarginTop(1).
|
||
MarginBottom(1)
|
||
sep := sepStyle.Render(strings.Repeat("┈", width))
|
||
result.WriteString(sep)
|
||
result.WriteString("\n")
|
||
}
|
||
if m.thinkingBuffer.String() != "" && m.answerBuffer.String() == "" {
|
||
result.WriteString(m.agentStyle.Render("Thinking: "))
|
||
result.WriteString(m.answerBuffer.String())
|
||
|
||
result.WriteString("\n")
|
||
var sepStyle lipgloss.Style = lipgloss.NewStyle().
|
||
Foreground(lipgloss.Color("240")).
|
||
MarginTop(1).
|
||
MarginBottom(1)
|
||
sep := sepStyle.Render(strings.Repeat("┈", width))
|
||
result.WriteString(sep)
|
||
result.WriteString("\n")
|
||
}
|
||
return result.String()
|
||
}
|
||
|
||
func (m *uiModel) startWait() {
|
||
if !m.waiting {
|
||
m.viewport.SetHeight(m.viewport.Height() - 1)
|
||
m.waiting = true
|
||
}
|
||
}
|
||
|
||
func (m *uiModel) stopWait() {
|
||
if m.waiting {
|
||
m.waiting = false
|
||
m.viewport.SetHeight(m.viewport.Height() - 1)
|
||
}
|
||
}
|
||
|
||
func (m *uiModel) setErrorMessage(message string) {
|
||
m.statusMessage = m.errorStyle.Render(message)
|
||
}
|
||
|
||
func (m *uiModel) setStatusMessage(message string) {
|
||
if message == "" {
|
||
m.statusMessage = m.statusStyle.Render("model: " + m.client.Cfg.Llm.Model + " url: " + m.client.Cfg.Api.Url)
|
||
return
|
||
}
|
||
m.statusMessage = m.statusStyle.Render(message)
|
||
}
|
||
|
||
func (m *uiModel) openCommands() {
|
||
if !m.commands.open {
|
||
m.commands.open = true
|
||
m.viewport.SetHeight(m.viewport.Height() - commandDropdownHeight)
|
||
}
|
||
m.commands.filter(strings.TrimPrefix(m.textarea.Value(), "/"))
|
||
}
|
||
|
||
func (m *uiModel) closeCommands() {
|
||
if m.commands.open {
|
||
m.commands.open = false
|
||
m.viewport.SetHeight(m.viewport.Height() + commandDropdownHeight)
|
||
}
|
||
}
|