cleanup db and added embedding coroutine
This commit is contained in:
Radu Macocian (admac)
@@ -0,0 +1,576 @@
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user