diff --git a/db/history/history.go b/db/history/history.go index 2b5f338..624c44b 100644 --- a/db/history/history.go +++ b/db/history/history.go @@ -1,49 +1,55 @@ package history import ( + "errors" "log/slog" "net" "net/url" + "strconv" + "strings" "git.estatecloud.org/radumaco/souvenir/config" "git.estatecloud.org/radumaco/souvenir/model" + "github.com/gopsql/db" "github.com/gopsql/pgx" ) type DbClient struct { - Cfg config.DbConfig + Cfg config.DbConfig Logger slog.Logger } -type Conversation struct { - Id string - Messages []model.Message -} - -func (client DbClient) init() error { +func (client DbClient) Init() error { const createConversationTable = ` CREATE TABLE IF NOT EXISTS conversations ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + name TEXT, + summary TEXT, created_at TIMESTAMPTZ NOT NULL DEFAULT now() )` const createMessageTable = ` CREATE TABLE IF NOT EXISTS messages ( id UUID PRIMARY KEY DEFAULT gen_random_uuid(), role TEXT NOT NULL, + seq INTEGER NOT NULL, content TEXT NOT NULL, - conversation_id UUID, - created_at TIMESTAMPTZ NOT NULL DEFAULT now() + conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + UNIQUE (conversation_id, seq) )` const createToolCallTable = ` CREATE TABLE IF NOT EXISTS tool_calls ( - id UUID PRIMARY KEY, + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), type TEXT NOT NULL, functionName TEXT NOT NULL, functionArgs TEXT NOT NULL, - message_id UUID, + message_id UUID NOT NULL REFERENCES messages(id) ON DELETE CASCADE, created_at TIMESTAMPTZ NOT NULL DEFAULT now() )` + if err := client.ensureDbExists(); err != nil { + return err + } conn := pgx.MustOpen(client.connectionString()) defer conn.Close() tables := []struct { @@ -61,29 +67,202 @@ func (client DbClient) init() error { (SELECT 1 FROM information_schema.tables WHERE table_schema = 'public' AND table_name =$1)`, table.tableName).Scan(&exists); err != nil { - client.Logger.Error("Error while checking table", "table", table.tableName) - return err + client.Logger.Error("Error while checking table", "table", table.tableName, "error", err) + return err } if exists { client.Logger.Debug("Table found", "table", table.tableName) } else { client.Logger.Debug("Table not found. Creating", "table", table.tableName) if _, err := conn.Exec(table.script); err != nil { - client.Logger.Error("Could not create table", "table", table.tableName) + client.Logger.Error("Could not create table", "table", table.tableName, "error", err) return err } } } - if _, err := conn.Exec(createToolCallTable); err != nil { - return err - } - if _, err := conn.Exec(createToolCallTable); err != nil { - return err - } return nil } +func (client DbClient) GetConversation(id string) (model.Conversation, error) { + conn := pgx.MustOpen(client.connectionString()) + defer conn.Close() + + client.Logger.Debug("Searching for conversation", "id", id) + var conv model.Conversation + if err := conn.QueryRow("SELECT id, name, summary FROM conversations WHERE id = $1", id).Scan(&conv.Id, &conv.Name, &conv.Summary); err != nil { + client.Logger.Error("Could not fetch conversation", "id", id, "error", err) + return conv, err + } + rows, err := conn.Query(`SELECT t.id, t.type, t.functionName, t.functionArgs, m.id, m.role, m.content, m.seq + FROM messages m + LEFT JOIN tool_calls t + ON m.id = t.message_id + WHERE m.conversation_id = $1 ORDER BY m.seq`, id) + if err != nil { + client.Logger.Error("Could not fetch messages", "conversation_id", id, "error", err) + return conv, err + } + defer rows.Close() + var messages []model.Message + for rows.Next() { + var ( + tcId, tcType, tcName, tcArgs *string + msgId, role, content string + seq int + ) + err = rows.Scan(&tcId, &tcType, &tcName, &tcArgs, &msgId, &role, &content, &seq) + if err != nil { + client.Logger.Error("Could not get message", "error", err) + } + if len(messages) == 0 || messages[len(messages) - 1].Id != msgId { + messages = append(messages, model.Message{ + Id: msgId, Role: role, Content: content, Seq: seq, + }) + } + if tcId != nil { + deref := func(s *string) string { + if s == nil { + return "" + } + return *s + } + messages[len(messages) - 1].ToolCalls = append(messages[len(messages)-1].ToolCalls, model.ToolCall{ + Id: *tcId, + Type: deref(tcType), + FunctionName: deref(tcName), + FunctionArguments: deref(tcArgs), + }) + } + } + conv.Messages = messages + client.Logger.Debug("Found conversation", "conversation", conv) + return conv, nil +} + +func (client DbClient) GetConversations() ([]model.Conversation, error) { + client.Logger.Debug("Fetching conversations") + conn := pgx.MustOpen(client.connectionString()) + defer conn.Close() + rows, err := conn.Query("SELECT id, name, summary FROM conversations") + if err != nil { + client.Logger.Error("Could not fetch conversations", "error", err) + return []model.Conversation{}, err + } + defer rows.Close() + conversations := []model.Conversation{} + for rows.Next() { + var conv model.Conversation + rows.Scan(&conv.Id, &conv.Name, &conv.Summary) + conversations = append(conversations, conv) + } + client.Logger.Debug("Found conversations", "count", len(conversations)) + return conversations, nil +} + +func (client DbClient) SaveConversation(conv model.Conversation) (model.Conversation, error) { + client.Logger.Debug("Saving conversation", "conversation", conv) + conn := pgx.MustOpen(client.connectionString()) + defer conn.Close() + delta := 0 + if conv.Id == "" { + if err := conn.QueryRow(`INSERT INTO conversations(name, summary) VALUES($1, $2) RETURNING id`, + conv.Name, + conv.Summary).Scan(&conv.Id); err != nil { + client.Logger.Error("Could not create conversation", "name", conv.Name, "error", err) + return conv, err + } + } + if err := conn.QueryRow(`SELECT COALESCE(MAX(seq) + 1, 1) FROM messages WHERE conversation_id = $1`, conv.Id).Scan(&delta); err != nil { + client.Logger.Warn("Could not get delta", "conversation_id", conv.Id, "error", err) + return conv, errors.New("Could not get delta") + } + + client.Logger.Debug("Found delta.", "delta", delta, "Unsaved messages", len(conv.Messages[delta - 1:])) + + if delta >= len(conv.Messages) { + return conv, nil + } + for _, message := range conv.Messages[delta - 1:] { + id, err := client.saveMessage(conv.Id, message, conn) + if err != nil { + client.Logger.Error("Could not save message", "message_id", message.Id, "seq", message.Seq, "error", err) + return conv, errors.New("Could not save message") + } + message.Id = id + } + return conv, nil +} + +func (client DbClient) saveMessage(conversation_id string, message model.Message, conn db.DB) (string, error) { + var exists bool + if message.Id != "" { + if err := conn.QueryRow(`SELECT EXISTS + (SELECT 1 FROM messages + WHERE id = $1)`, message.Id).Scan(&exists); err != nil { + client.Logger.Warn("Error while scanning for message", "id", message.Id, "error", err) + return message.Id, err + } + } + if exists { + client.Logger.Warn("Message found. Skipping", "id", message.Id) + return message.Id, nil + } + + if err := conn.QueryRow(`INSERT INTO messages (seq, role, content, conversation_id) VALUES ($1, $2, $3, $4) RETURNING id`, + message.Seq, message.Role, message.Content, conversation_id).Scan(&message.Id); err != nil { + client.Logger.Error("Could not save message", "seq", message.Seq, "error", err) + return message.Id, errors.New("Could not save message") + } + + placeholder := 0 + var placeholders strings.Builder + args := []any{} + const cols = 4 + for _, call := range message.ToolCalls { + placeholders.WriteString("($") + placeholders.WriteString(strconv.Itoa(placeholder + 1)) + placeholders.WriteString(",$") + placeholders.WriteString(strconv.Itoa(placeholder + 2)) + placeholders.WriteString(",$") + placeholders.WriteString(strconv.Itoa(placeholder + 3)) + placeholders.WriteString(",$") + placeholders.WriteString(strconv.Itoa(placeholder + 4)) + placeholders.WriteString(")") + args = append(args, call.Type, call.FunctionName, call.FunctionArguments, message.Id) + placeholder += cols + } + if _, err := conn.Exec(`INSERT INTO tool_calls(type, functionName, functionArgs, message_id) VALUES `+placeholders.String(), + args...); err != nil && placeholder > 0 { + client.Logger.Warn("Could not insert tool_calls", "message_id", message.Id, "error", err) + } + return message.Id, nil +} + +func (client DbClient) ensureDbExists() error { + maintenanceUrl := &url.URL{ + Scheme: "postgres", + User: url.UserPassword(client.Cfg.User, client.Cfg.Password), + Host: net.JoinHostPort(client.Cfg.Url, client.Cfg.Port), + Path: "postgres", + } + conn := pgx.MustOpen(maintenanceUrl.String()) + defer conn.Close() + + var exists bool + if err := conn.QueryRow(`SELECT EXISTS(SELECT 1 FROM pg_database WHERE datname = $1)`, client.Cfg.DbName).Scan(&exists); err != nil { + return err + } + if exists { + client.Logger.Debug("Database already exists", "name", client.Cfg.DbName) + return nil + } + + client.Logger.Info("Creating database", "name", client.Cfg.DbName) + _, err := conn.Exec(`CREATE DATABASE "` + client.Cfg.DbName + `"`) + return err +} + func (client DbClient) connectionString() string { u := &url.URL{ Scheme: "postgres", diff --git a/hello.go b/hello.go index 0c353ee..8002993 100644 --- a/hello.go +++ b/hello.go @@ -6,6 +6,7 @@ import ( tea "charm.land/bubbletea/v2" "git.estatecloud.org/radumaco/souvenir/config" + "git.estatecloud.org/radumaco/souvenir/db/history" ui "git.estatecloud.org/radumaco/souvenir/ui" ) @@ -18,7 +19,15 @@ func main() { } var logger = slog.Default().With("Compoent", "Main") logger.Debug("Config initialized") - p := tea.NewProgram(ui.InitialModel(*config)) + + db := history.DbClient{Cfg: config.Db, Logger: *slog.Default().With("Component", "History")}; + err = db.Init() + if err != nil { + logger.Error("Error while initializing database", "message", err) + os.Exit(1) + } + + p := tea.NewProgram(ui.InitialModel(*config, db)) if _, err := p.Run(); err != nil { logger.Error("Alas, there's been an error:", "message", err) os.Exit(1) diff --git a/llm/client.go b/llm/client.go index 30f5ac2..51c5513 100644 --- a/llm/client.go +++ b/llm/client.go @@ -152,7 +152,7 @@ func (client LLMClient) Call(query string, messages []model.Message) ([]model.Me return []model.Message{}, errors.New("The call returned an empty response") } for _, choice := range response.Choices { - client.Logger.Error("Appending response message", "message", choice.Message) + client.Logger.Debug("Appending response message", "message", choice.Message) result = append(result, choice.Message) } diff --git a/model/Conversation.go b/model/Conversation.go new file mode 100644 index 0000000..e4670ee --- /dev/null +++ b/model/Conversation.go @@ -0,0 +1,8 @@ +package model + +type Conversation struct { + Id string + Name string + Summary string + Messages []Message +} diff --git a/model/Message.go b/model/Message.go index 5afa771..5b83225 100644 --- a/model/Message.go +++ b/model/Message.go @@ -1,14 +1,9 @@ package model type Message struct { + Id string `json:"-"` + Seq int `json:"-"` Role string `json:"role"` Content string `json:"content"` ToolCalls []ToolCall `json:"tool_calls,omitempty"` } - -type ToolCall struct { - ID string `json:"id"` - Type string `json:"type"` - FunctionName string `json:"function"` - FunctionArguments string `json:"arguments"` -} diff --git a/model/ToolCall.go b/model/ToolCall.go new file mode 100644 index 0000000..0a35176 --- /dev/null +++ b/model/ToolCall.go @@ -0,0 +1,8 @@ +package model + +type ToolCall struct { + Id string `json:"-"` + Type string `json:"type"` + FunctionName string `json:"function"` + FunctionArguments string `json:"arguments"` +} diff --git a/ui/chat.go b/ui/chat.go index 500c314..accf27f 100644 --- a/ui/chat.go +++ b/ui/chat.go @@ -11,6 +11,7 @@ import ( 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" ) @@ -19,24 +20,27 @@ type focusState int const ( focusChat focusState = iota - focusOverlay + focusModels + focusHistory ) 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 + logger slog.Logger + viewport viewport.Model + conversation model.Conversation + textarea textarea.Model + focus focusState + picker picker + senderStyle lipgloss.Style + agentStyle lipgloss.Style + errorStyle lipgloss.Style + client llm.LLMClient + errorMessage string + history history.DbClient + width int + height int + waiting bool + err error } type agentResponseMessage struct { @@ -44,12 +48,12 @@ type agentResponseMessage struct { err error } -type modelsResponseMessage struct { - models []string - err error +type conversationSavedMessage struct { + conversation model.Conversation + err error } -func InitialModel(config config.Config) uiModel { +func InitialModel(config config.Config, client history.DbClient) uiModel { ta := textarea.New() ta.Placeholder = "Send a message..." ta.SetVirtualCursor(false) @@ -75,18 +79,19 @@ func InitialModel(config config.Config) uiModel { 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, + textarea: ta, + conversation: model.Conversation{}, + 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, + history: client, + logger: *slog.Default().With("Component", "TUI"), + err: nil, } - model.picker = newModelPicker(&model.focus) + model.picker = newPicker() return model } @@ -108,7 +113,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { var messages = m.renderMessages() - if len(m.messages) > 0 { + 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())) @@ -119,14 +124,14 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.viewport.GotoBottom() case tea.KeyPressMsg: - if m.focus == focusOverlay { + if m.focus != focusChat { picker, cmd := m.picker.Update(msg) m.picker = picker return m, cmd } switch msg.String() { case "esc": - if m.focus == focusOverlay { + if m.focus != focusChat { m.focus = focusChat } return m, nil @@ -137,28 +142,35 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.textarea.InsertRune('\n') return m, nil case "enter": + m.clearErrorMessage() input := m.textarea.Value() - m.messages = append(m.messages, model.Message{ - Content: input, - Role: "user", - }) + m.textarea.SetValue("") m.logger.Debug("Received enter. Creating new message", "message", input) switch input { + case "/history": + m.logger.Debug("Received history call") + m.focus = focusHistory + return m, m.getHistory() case "/models": m.logger.Debug("Received models sequence") - m.focus = focusOverlay + m.focus = focusModels return m, m.getModels() case "/exit": m.logger.Debug("Received exit sequence. Quiting") fmt.Println(m.textarea.Value()) return m, tea.Quit default: + 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.waiting = true - return m, m.callAgent(input, m.messages) + return m, tea.Batch(m.callAgent(input, m.conversation.Messages), m.saveConversation()) } default: var cmd tea.Cmd @@ -178,29 +190,56 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { 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(), - }} + m.setErrorMessage(msg.err.Error()) } for _, r := range resp { + r.Seq = len(m.conversation.Messages) + 1 m.logger.Debug("Appending response message", "message", r) - m.messages = append(m.messages, 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, nil + return m, m.saveConversation() - case modelChosenMsg: - m.logger.Debug("Received modelChosenMsg", "model", msg.id) - m.client.Cfg.Llm.Model = msg.id + case pickerChosenMsg: + switch m.focus { + case focusModels: + m.logger.Debug("Received pickerChosenMsg", "model", msg.id) + m.client.Cfg.Llm.Model = msg.id + m.waiting = false + m.focus = focusChat + case focusHistory: + m.logger.Debug("Received pickerChosenMsg", "history", msg.id) + conv, err := m.history.GetConversation(msg.id) + if err != nil { + m.logger.Error("Could not retrieve conversation", "id", msg.id) + } + m.conversation = conv + m.waiting = false + 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.waiting = false m.focus = focusChat + + case conversationSavedMessage: + m.logger.Debug("Received conversationSavedMessage", "conversation", msg.conversation) + if msg.err != nil { + m.logger.Error("Message is error", "error", msg.err) + m.setErrorMessage(msg.err.Error()) + return m, nil + } + m.conversation = msg.conversation } - if m.focus == focusOverlay { + if m.focus != focusChat { picker, cmd := m.picker.Update(msg) m.picker = picker return m, cmd @@ -210,19 +249,28 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } func (m uiModel) View() tea.View { - if m.focus == focusOverlay { + 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.errorMessage != "" { + parts = append(parts, m.errorStyle.Render(m.errorMessage)) + } + parts = append(parts, "model: "+m.client.Cfg.Llm.Model+"; url: "+m.client.Cfg.Api.Url) + parts = append(parts, m.textarea.View()) ui := lipgloss.JoinVertical( lipgloss.Left, - m.viewport.View(), - "model: " + m.client.Cfg.Llm.Model + "; url: " + m.client.Cfg.Api.Url, - m.textarea.View()) + parts...) + c := m.textarea.Cursor() if c != nil { c.Y += lipgloss.Height(m.viewport.View()) + 1 + if m.errorMessage != "" { + c.Y++ + } gap := m.width - lipgloss.Width(ui) c.X += max(gap/2, 0) } @@ -240,41 +288,96 @@ func (m uiModel) callAgent(input string, messages []model.Message) tea.Cmd { } } -func (m uiModel) getModels() tea.Cmd { - m.logger.Debug("Querying models") +func (m uiModel) saveConversation() tea.Cmd { + m.logger.Debug("Saving conversation", "conversation", m.conversation) return func() tea.Msg { - resp, _ := m.client.Models() - m.logger.Debug("Received models from llm", "models", resp) - return modelsLoadedMsg{models: resp} + conv, err := m.history.SaveConversation(m.conversation) + return conversationSavedMessage{conversation: conv, err: err} } } -func (m uiModel) renderMessages() string { +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() + 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.Name != "" { + name = conv.Name + } + 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.messages { + 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) - case "Error": - result.WriteString(m.errorStyle.Render("Error:")) - result.WriteString(message.Content) - } - result.WriteString("\n") + 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") + 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) setErrorMessage(message string) { + if m.errorMessage == "" { + m.viewport.SetHeight(m.viewport.Height() - 1) + } + m.errorMessage = message +} + +func (m *uiModel) clearErrorMessage() { + if m.errorMessage != "" { + m.viewport.SetHeight(m.viewport.Height() + 1) + m.errorMessage = "" + } +} diff --git a/ui/landing.go b/ui/landing.go index a8057a7..f70e360 100644 --- a/ui/landing.go +++ b/ui/landing.go @@ -39,6 +39,7 @@ const Wordmark = ` |___/\___/ \__,_| \_/ \___|_| |_|_|_|` const Commands = ` +/history /models /exit ` diff --git a/ui/modelPicker.go b/ui/modelPicker.go deleted file mode 100644 index 25140d8..0000000 --- a/ui/modelPicker.go +++ /dev/null @@ -1,61 +0,0 @@ -package ui - -import ( - "log/slog" - - "charm.land/bubbles/v2/list" - tea "charm.land/bubbletea/v2" -) - -type modelPicker struct { - focus *focusState - logger slog.Logger - list list.Model -} - -type modelsLoadedMsg struct{ models []string } -type modelsErrMsg struct{ err error } -type modelChosenMsg struct{ id string } - -type modelItem struct{ name string } -func (i modelItem) FilterValue() string { return i.name } -func (i modelItem) Title() string { return i.name } -func (i modelItem) Description() string { return i.name } - -func newModelPicker(focus *focusState) modelPicker { - l := list.New(nil, list.NewDefaultDelegate(), 0, 0) - l.Title = "Choose a model" - return modelPicker{list: l, logger: *slog.Default().With("Component", "ModelPicker"), focus: focus} -} - -func (p *modelPicker) SetSize(w, h int) { - p.list.SetSize(w, h) -} - -func (p modelPicker) Update(msg tea.Msg) (modelPicker, tea.Cmd) { - switch msg := msg.(type) { - case modelsLoadedMsg: - p.logger.Debug("Received modelsLoadedMsg", "models", msg.models) - items := make([]list.Item, len(msg.models)) - for i, m := range msg.models { - items[i] = modelItem{ name: m} - } - return p, p.list.SetItems(items) - case tea.KeyPressMsg: - if msg.String() == "enter" { - p.logger.Debug("Selecting model", "model", p.list.SelectedItem()) - if it, ok := p.list.SelectedItem().(modelItem); ok { - *p.focus = focusChat - p.logger.Debug("Sending modelChosenMsg") - return p, func() tea.Msg { return modelChosenMsg{it.name} } - } - } - } - var cmd tea.Cmd - p.list, cmd = p.list.Update(msg) - return p, cmd -} - -func (p modelPicker) View() string { - return p.list.View() -} diff --git a/ui/picker.go b/ui/picker.go new file mode 100644 index 0000000..8659a7f --- /dev/null +++ b/ui/picker.go @@ -0,0 +1,73 @@ +package ui + +import ( + "log/slog" + + "charm.land/bubbles/v2/list" + tea "charm.land/bubbletea/v2" +) + +type picker struct { + logger slog.Logger + list list.Model + title string +} + +type modelsLoadedMsg struct { + models []modelItem + title string +} +type pickerChosenMsg struct{ id string } +type pickerDismissMsg struct {} + +type modelItem struct { + name string + description string +} + +func (i modelItem) FilterValue() string { return i.name } +func (i modelItem) Title() string { return i.name } +func (i modelItem) Description() string { return i.description } + +func newPicker() picker { + l := list.New(nil, list.NewDefaultDelegate(), 0, 0) + l.DisableQuitKeybindings() + return picker{list: l, logger: *slog.Default().With("Component", "Picker")} +} + +func (p *picker) SetSize(w, h int) { + p.list.SetSize(w, h) +} + +func (p picker) Update(msg tea.Msg) (picker, tea.Cmd) { + switch msg := msg.(type) { + case modelsLoadedMsg: + p.logger.Debug("Received modelsLoadedMsg", "models", msg.models) + p.list.Title = msg.title + items := make([]list.Item, len(msg.models)) + for i, m := range msg.models { + items[i] = modelItem{name: m.name, description: m.description} + } + return p, p.list.SetItems(items) + case tea.KeyPressMsg: + if msg.String() == "q" || msg.String() == "esc" { + if p.list.FilterState() != list.Filtering { + return p, func() tea.Msg { return pickerDismissMsg{} } + } + } + if msg.String() == "enter" { + p.logger.Debug("Selecting model", "model", p.list.SelectedItem()) + if it, ok := p.list.SelectedItem().(modelItem); ok { + p.logger.Debug("Sending modelChosenMsg") + return p, func() tea.Msg { return pickerChosenMsg{it.name} } + } + } + } + var cmd tea.Cmd + p.list, cmd = p.list.Update(msg) + return p, cmd +} + +func (p picker) View() string { + return p.list.View() +}