Radu Macocian (admac)
281 lines
7.2 KiB
Go
281 lines
7.2 KiB
Go
package ui
|
|
|
|
import (
|
|
"fmt"
|
|
"log/slog"
|
|
"strings"
|
|
|
|
"charm.land/bubbles/v2/cursor"
|
|
"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"
|
|
llm "git.estatecloud.org/radumaco/souvenir/llm"
|
|
"git.estatecloud.org/radumaco/souvenir/model"
|
|
)
|
|
|
|
type focusState int
|
|
|
|
const (
|
|
focusChat focusState = iota
|
|
focusOverlay
|
|
)
|
|
|
|
type uiModel struct {
|
|
logger slog.Logger
|
|
viewport viewport.Model
|
|
messages []model.Message
|
|
textarea textarea.Model
|
|
focus focusState
|
|
picker modelPicker
|
|
senderStyle lipgloss.Style
|
|
agentStyle lipgloss.Style
|
|
errorStyle lipgloss.Style
|
|
client llm.LLMClient
|
|
width int
|
|
height int
|
|
waiting bool
|
|
err error
|
|
}
|
|
|
|
type agentResponseMessage struct {
|
|
response []model.Message
|
|
err error
|
|
}
|
|
|
|
type modelsResponseMessage struct {
|
|
models []string
|
|
err error
|
|
}
|
|
|
|
func InitialModel(config config.Config) 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)
|
|
|
|
model := uiModel{
|
|
textarea: ta,
|
|
messages: []model.Message{},
|
|
viewport: vp,
|
|
senderStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("5")),
|
|
errorStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("9")),
|
|
agentStyle: lipgloss.NewStyle().Foreground(lipgloss.Color("86")),
|
|
client: llm.LLMClient{Cfg: config, Logger: *slog.Default().With("Component", "LLM")},
|
|
focus: focusChat,
|
|
logger: *slog.Default().With("Component", "TUI"),
|
|
err: nil,
|
|
}
|
|
model.picker = newModelPicker(&model.focus)
|
|
return model
|
|
}
|
|
|
|
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)
|
|
|
|
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)
|
|
|
|
var messages = m.renderMessages()
|
|
|
|
if len(m.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 == focusOverlay {
|
|
picker, cmd := m.picker.Update(msg)
|
|
m.picker = picker
|
|
return m, cmd
|
|
}
|
|
switch msg.String() {
|
|
case "esc":
|
|
if m.focus == focusOverlay {
|
|
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 "enter":
|
|
input := m.textarea.Value()
|
|
m.messages = append(m.messages, model.Message{
|
|
Content: input,
|
|
Role: "user",
|
|
})
|
|
m.logger.Debug("Received enter. Creating new message", "message", input)
|
|
switch input {
|
|
case "/models":
|
|
m.logger.Debug("Received models sequence")
|
|
m.focus = focusOverlay
|
|
return m, m.getModels()
|
|
case "/exit":
|
|
m.logger.Debug("Received exit sequence. Quiting")
|
|
fmt.Println(m.textarea.Value())
|
|
return m, tea.Quit
|
|
default:
|
|
var messages = m.renderMessages()
|
|
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
|
|
m.textarea.Reset()
|
|
m.viewport.GotoBottom()
|
|
m.waiting = true
|
|
return m, m.callAgent(input, m.messages)
|
|
}
|
|
default:
|
|
var cmd tea.Cmd
|
|
m.textarea, cmd = m.textarea.Update(msg)
|
|
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 agentResponseMessage:
|
|
m.logger.Debug("Received agentResponseMessage")
|
|
m.waiting = false
|
|
resp := msg.response
|
|
if msg.err != nil {
|
|
m.logger.Error("Message is error", "error", msg.err)
|
|
resp = []model.Message{{
|
|
Role: "Error",
|
|
Content: msg.err.Error(),
|
|
}}
|
|
}
|
|
|
|
for _, r := range resp {
|
|
m.logger.Debug("Appending response message", "message", r)
|
|
m.messages = append(m.messages, r)
|
|
}
|
|
var messages = m.renderMessages()
|
|
m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages))
|
|
m.viewport.GotoBottom()
|
|
return m, nil
|
|
|
|
case modelChosenMsg:
|
|
m.logger.Debug("Received modelChosenMsg", "model", msg.id)
|
|
m.client.Cfg.Llm.Model = msg.id
|
|
m.waiting = false
|
|
m.focus = focusChat
|
|
}
|
|
|
|
if m.focus == focusOverlay {
|
|
picker, cmd := m.picker.Update(msg)
|
|
m.picker = picker
|
|
return m, cmd
|
|
}
|
|
|
|
return m, nil
|
|
}
|
|
|
|
func (m uiModel) View() tea.View {
|
|
if m.focus == focusOverlay {
|
|
return tea.NewView(lipgloss.Place(m.width, m.height,
|
|
lipgloss.Center, lipgloss.Center,
|
|
m.picker.View()))
|
|
}
|
|
ui := lipgloss.JoinVertical(
|
|
lipgloss.Left,
|
|
m.viewport.View(),
|
|
"model: " + m.client.Cfg.Llm.Model + "; url: " + m.client.Cfg.Api.Url,
|
|
m.textarea.View())
|
|
c := m.textarea.Cursor()
|
|
if c != nil {
|
|
c.Y += lipgloss.Height(m.viewport.View()) + 1
|
|
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(input string, messages []model.Message) tea.Cmd {
|
|
m.logger.Debug("Calling agent", "query", input, "messages", messages)
|
|
return func() tea.Msg {
|
|
resp, err := m.client.Call(input, messages)
|
|
return agentResponseMessage{response: resp, err: err}
|
|
}
|
|
}
|
|
|
|
func (m uiModel) getModels() tea.Cmd {
|
|
m.logger.Debug("Querying models")
|
|
return func() tea.Msg {
|
|
resp, _ := m.client.Models()
|
|
m.logger.Debug("Received models from llm", "models", resp)
|
|
return modelsLoadedMsg{models: resp}
|
|
}
|
|
}
|
|
|
|
func (m uiModel) renderMessages() string {
|
|
var result strings.Builder
|
|
var width = m.viewport.Width()
|
|
for _, message := range m.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()))
|
|
case "assistant":
|
|
result.WriteString(m.agentStyle.Render("Agent: "))
|
|
result.WriteString(message.Content)
|
|
case "Error":
|
|
result.WriteString(m.errorStyle.Render("Error:"))
|
|
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")
|
|
}
|
|
return result.String()
|
|
}
|