diff --git a/config/config.go b/config/config.go index 8811f23..58c3fe6 100644 --- a/config/config.go +++ b/config/config.go @@ -50,6 +50,7 @@ type EmbeddingConfig struct { Dim int Timeout int BatchSize int + Interval int } type MemoryConfig struct { @@ -152,6 +153,7 @@ func defaultConfig() *Config { Model: "", BatchSize: 64, Timeout: 60000, + Interval: 60, }, } } @@ -185,6 +187,7 @@ func overrideFromEnv(cfg *Config) { }{ {conf: &cfg.Api.Timeout, key: ENV_PREFIX + "api__timeout"}, {conf: &cfg.Embedding.Timeout, key: ENV_PREFIX + "embedding__timeout"}, + {conf: &cfg.Embedding.Interval, key: ENV_PREFIX + "embedding__interval"}, } for _, override := range durationOverrides { val := os.Getenv(override.key) diff --git a/db/embed/embed.go b/db/embed/embed.go index 83c2e44..1b1645d 100644 --- a/db/embed/embed.go +++ b/db/embed/embed.go @@ -7,6 +7,7 @@ import ( "log/slog" "git.estatecloud.org/radumaco/souvenir/config" + "git.estatecloud.org/radumaco/souvenir/db" "git.estatecloud.org/radumaco/souvenir/llm" "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" @@ -32,6 +33,35 @@ func (e EmbedStore) TableName() string { return e.tableName } +func (e EmbedStore) GetConversationsToEmbed(ctx context.Context) ([]string, error) { + rows, err := e.pool.Query(ctx, fmt.Sprintf( + ` + SELECT DISTINC c.id + FROM conversations c + JOIN message_chunks mc + ON mc.conversation_id = c.id + LEFT JOIN %s v + ON v.chunk_id = mc.id AND v.content_hash = mc.content_hash + WHERE c.updated_at < now() - interval '5 hours' + AND v.chunk_id is NULL; + `, + e.TableName())) + if err != nil { + e.logger.Error("Error while getting conversations to embed", "error", err) + } + defer rows.Close() + var ids []string + for rows.Next() { + var id string + err = rows.Scan(&id) + if err != nil { + e.logger.Error("Could not get id", "error", err) + } + ids = append(ids, id) + } + return ids, nil +} + func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedder) (*EmbedStore, error) { logger := *slog.Default().With("Component", "EmbedStore") @@ -47,17 +77,16 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd pool: pool, } - const createVectorTableTemplate = ` - CREATE TABLE IF NOT EXISTS %s ( - chunk_id UUID PRIMARY KEY REFERENCES message_chunks(id) ON DELETE CASCADE, - embedding vector(%d), - content_hash text not null - )` var exists bool - if err := pool.QueryRow(ctx, `SELECT EXISTS - (SELECT 1 FROM information_schema.tables - WHERE table_schema = 'public' - AND table_name =$1)`, store.TableName()).Scan(&exists); err != nil { + if err := pool.QueryRow(ctx, + ` + SELECT EXISTS + (SELECT 1 + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name =$1) + `, + store.TableName()).Scan(&exists); err != nil { logger.Error("Error while checking table", "table", store.TableName(), "error", err) return nil, err } @@ -65,7 +94,7 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd logger.Debug("Table found", "table", store.TableName()) } else { logger.Debug("Table not found. Creating", "table", store.TableName()) - script := fmt.Sprintf(createVectorTableTemplate, store.TableName(), store.embedder.Dim()) + script := fmt.Sprintf(db.CreateVectorTableTemplate, store.TableName(), store.embedder.Dim()) if _, err := pool.Exec(ctx, script); err != nil { logger.Error("Could not create table", "table", store.TableName(), "error", err) return nil, err @@ -75,11 +104,17 @@ func NewEmbedStore(ctx context.Context, cfg config.DbConfig, embedder llm.Embedd } func (e EmbedStore) EmbedConversation(ctx context.Context, conversationId string) error { - rows, err := e.pool.Query(ctx, fmt.Sprintf(`SELECT c.id, c.content, c.content_hash - FROM message_chunks c + rows, err := e.pool.Query(ctx, fmt.Sprintf( + ` + SELECT c.id, c.content, c.content_hash + FROM message_chunks c LEFT JOIN %s v - ON v.chunk_id = c.id AND v.content_hash = c.content_hash - WHERE c.conversation_id = $1 and v.chunk_id IS NULL`, e.TableName()), + ON v.chunk_id = c.id + AND v.content_hash = c.content_hash + WHERE c.conversation_id = $1 + AND v.chunk_id IS NULL + `, + e.TableName()), conversationId) if err != nil { @@ -131,12 +166,15 @@ func (e EmbedStore) EmbedConversation(ctx context.Context, conversationId string } func (e EmbedStore) upsert(ctx context.Context, chunks []Chunk, vecs [][]float32) error { - query := fmt.Sprintf(` + query := fmt.Sprintf( + ` INSERT INTO %s (chunk_id, embedding, content_hash) - VALUES ($1, $2, $3) + VALUES ($1, $2, $3) ON CONFLICT (chunk_id) DO UPDATE - SET embedding = EXCLUDED.embedding, - content_hash = EXCLUDED.content_hash`, e.TableName()) + SET embedding = EXCLUDED.embedding, + content_hash = EXCLUDED.content_hash + `, + e.TableName()) batch := &pgx.Batch{} for i, chunk := range chunks { @@ -152,3 +190,7 @@ func (e EmbedStore) upsert(ctx context.Context, chunks []Chunk, vecs [][]float32 } return nil } + +func (e EmbedStore) Close() { + e.pool.Close() +} diff --git a/db/history/history.go b/db/history/history.go index 351c1c8..162d207 100644 --- a/db/history/history.go +++ b/db/history/history.go @@ -11,6 +11,7 @@ import ( "strings" "git.estatecloud.org/radumaco/souvenir/config" + "git.estatecloud.org/radumaco/souvenir/db" "git.estatecloud.org/radumaco/souvenir/model" "github.com/gopsql/pgx" "github.com/jackc/pgx/v5/pgxpool" @@ -24,36 +25,6 @@ type DbClient struct { func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) { logger := *slog.Default().With("Component", "History") - const createConversationTable = ` - CREATE TABLE IF NOT EXISTS conversations ( - id UUID PRIMARY KEY DEFAULT gen_random_uuid(), - title TEXT, - title_source TEXT, - summary TEXT, - summary_through_seq int default 0, - 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 NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, - created_at TIMESTAMPTZ NOT NULL DEFAULT now(), - tsv tsvector generated always as (to_tsvector('english', content)) stored, - UNIQUE (conversation_id, seq) - ); - CREATE INDEX on messages using gin (tsv); - CREATE INDEX on messages (conversation_id, seq);` - const createMessageChunkTable = ` - CREATE TABLE IF NOT EXISTS message_chunks ( - id UUID PRIMARY KEY DEFAULT gen_random_uuid(), - conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, - message_id UUID NOT NULL REFERENCES messages(id) ON DELETE CASCADE, - content TEXT NOT NULL, - content_hash TEXT NOT NULL - )` if err := ensureDbExists(cfg, logger); err != nil { return nil, err @@ -67,17 +38,23 @@ func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) { tableName string script string }{ - {"conversations", createConversationTable}, - {"messages", createMessageTable}, - {"message_chunks", createMessageChunkTable}, + {"conversations", db.ConversationTableScript}, + {"messages", db.MessageTableScript}, + {"message_chunks", db.MessageChunkTableScript}, } for _, table := range tables { var exists bool - if err := pool.QueryRow(ctx, `SELECT EXISTS - (SELECT 1 FROM information_schema.tables - WHERE table_schema = 'public' - AND table_name =$1)`, table.tableName).Scan(&exists); err != nil { + if err := pool.QueryRow(ctx, + ` + SELECT EXISTS ( + SELECT 1 + FROM information_schema.tables + WHERE table_schema = 'public' + AND table_name =$1 + ) + `, + table.tableName).Scan(&exists); err != nil { logger.Error("Error while checking table", "table", table.tableName, "error", err) return nil, err } @@ -102,13 +79,22 @@ func Init(ctx context.Context, cfg config.DbConfig) (*DbClient, error) { func (cl DbClient) GetConversation(ctx context.Context, id string) (model.Conversation, error) { cl.logger.Debug("Searching for conversation", "id", id) var conv model.Conversation - if err := cl.pool.QueryRow(ctx, "SELECT id, title, summary FROM conversations WHERE id = $1", id).Scan(&conv.Id, &conv.Title, &conv.Summary); err != nil { + if err := cl.pool.QueryRow(ctx, + ` + SELECT id, title, summary + FROM conversations + WHERE id = $1 + `, id).Scan(&conv.Id, &conv.Title, &conv.Summary); err != nil { cl.logger.Error("Could not fetch conversation", "id", id, "error", err) return conv, err } - rows, err := cl.pool.Query(ctx, `SELECT id, role, content, seq - FROM messages - WHERE conversation_id = $1 ORDER BY seq`, id) + rows, err := cl.pool.Query(ctx, + ` + SELECT id, role, content, seq + FROM messages + WHERE conversation_id = $1 + ORDER BY seq + `, id) if err != nil { cl.logger.Error("Could not fetch messages", "conversation_id", id, "error", err) return conv, err @@ -137,7 +123,11 @@ func (cl DbClient) GetConversation(ctx context.Context, id string) (model.Conver func (cl DbClient) GetConversations(ctx context.Context) ([]model.Conversation, error) { cl.logger.Debug("Fetching conversations") - rows, err := cl.pool.Query(ctx, "SELECT id, title, summary FROM conversations") + rows, err := cl.pool.Query(ctx, + ` + SELECT id, title, summary + FROM conversations + `) if err != nil { cl.logger.Error("Could not fetch conversations", "error", err) return []model.Conversation{}, err @@ -157,19 +147,33 @@ func (cl DbClient) SaveConversation(ctx context.Context, conv model.Conversation cl.logger.Debug("Saving conversation", "conversation", conv) delta := 0 if conv.Id == "" { - if err := cl.pool.QueryRow(ctx, `INSERT INTO conversations(title, summary) VALUES($1, $2) RETURNING id`, - conv.Title, - conv.Summary).Scan(&conv.Id); err != nil { + if err := cl.pool.QueryRow(ctx, + ` + INSERT INTO conversations(title, summary) + VALUES ($1, $2) + RETURNING id + `, + conv.Title, conv.Summary).Scan(&conv.Id); err != nil { cl.logger.Error("Could not create conversation", "title", conv.Title, "error", err) return conv, err } } else { - if _, err := cl.pool.Exec(ctx, `UPDATE conversations SET title = $1, summary = $2 WHERE id = $3`, conv.Title, conv.Summary, conv.Id); err != nil { + if _, err := cl.pool.Exec(ctx, + ` + UPDATE conversations + SET title = $1, summary = $2 + WHERE id = $3`, conv.Title, conv.Summary, conv.Id); err != nil { cl.logger.Error("Could not update conversation", "id", conv.Id, "error", err) return conv, err } } - if err := cl.pool.QueryRow(ctx, `SELECT COALESCE(MAX(seq) + 1, 1) FROM messages WHERE conversation_id = $1`, conv.Id).Scan(&delta); err != nil { + if err := cl.pool.QueryRow(ctx, + ` + SELECT COALESCE(MAX(seq) + 1, 1) + FROM messages + WHERE conversation_id = $1 + `, + conv.Id).Scan(&delta); err != nil { cl.logger.Warn("Could not get delta", "conversation_id", conv.Id, "error", err) return conv, errors.New("Could not get delta") } @@ -193,16 +197,27 @@ func (cl DbClient) SaveConversation(ctx context.Context, conv model.Conversation func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, message model.Message) (string, error) { var exists bool if message.Id != "" { - if err := cl.pool.QueryRow(ctx, `SELECT EXISTS - (SELECT 1 FROM messages - WHERE id = $1)`, message.Id).Scan(&exists); err != nil { + if err := cl.pool.QueryRow(ctx, + ` + SELECT EXISTS ( + SELECT 1 + FROM messages + WHERE id = $1 + ) + `, + message.Id).Scan(&exists); err != nil { cl.logger.Warn("Error while scanning for message", "id", message.Id, "error", err) return message.Id, err } } if !exists { cl.logger.Debug("Insertting message ", "seq", message.Seq) - if err := cl.pool.QueryRow(ctx, `INSERT INTO messages (seq, role, content, conversation_id) VALUES ($1, $2, $3, $4) RETURNING id`, + if err := cl.pool.QueryRow(ctx, + ` + 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 { cl.logger.Error("Could not save message", "seq", message.Seq, "error", err) return message.Id, errors.New("Could not save message") @@ -213,7 +228,15 @@ func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, mess for _, chunk := range chunks { hash := sha256.Sum256([]byte(chunk)) - if err := cl.pool.QueryRow(ctx, `SELECT EXISTS (SELECT 1 FROM message_chunks WHERE message_id = $1 and content_hash = $2)`, + if err := cl.pool.QueryRow(ctx, + ` + SELECT EXISTS ( + SELECT 1 + FROM message_chunks + WHERE message_id = $1 + AND content_hash = $2 + ) + `, message.Id, hex.EncodeToString(hash[:])).Scan(&exists); err != nil { cl.logger.Warn("Error while scanning for message_chunk", "message id", message.Id, "error", err) return message.Id, errors.New("Could not save message chunk") @@ -223,8 +246,12 @@ func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, mess continue } - if _, err := cl.pool.Exec(ctx, `INSERT INTO message_chunks (conversation_id, message_id, content, content_hash) - VALUES ($1, $2, $3, $4)`, conversation_id, message.Id, chunk, hex.EncodeToString(hash[:])); err != nil { + if _, err := cl.pool.Exec(ctx, + ` + INSERT INTO message_chunks (conversation_id, message_id, content, content_hash) + VALUES ($1, $2, $3, $4) + `, + conversation_id, message.Id, chunk, hex.EncodeToString(hash[:])); err != nil { cl.logger.Warn("Error while creating message chunk", "message id", message.Id, "error", err) } } diff --git a/db/schema.go b/db/schema.go new file mode 100644 index 0000000..0c303f1 --- /dev/null +++ b/db/schema.go @@ -0,0 +1,44 @@ +package db + +const MessageTableScript = ` + 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 NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now(), + tsv tsvector generated always as (to_tsvector('english', content)) stored, + UNIQUE (conversation_id, seq) + ); + CREATE INDEX on messages using gin (tsv); + CREATE INDEX on messages (conversation_id, seq); + ` +const ConversationTableScript = ` + CREATE TABLE IF NOT EXISTS conversations ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + title TEXT, + title_source TEXT, + summary TEXT, + summary_through_seq int default 0, + created_at TIMESTAMPTZ NOT NULL DEFAULT now(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT now() + )` +const MessageChunkTableScript = ` + CREATE TABLE IF NOT EXISTS message_chunks ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + conversation_id UUID NOT NULL REFERENCES conversations(id) ON DELETE CASCADE, + message_id UUID NOT NULL REFERENCES messages(id) ON DELETE CASCADE, + content TEXT NOT NULL, + content_hash TEXT NOT NULL, + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + )` + +const CreateVectorTableTemplate = ` + CREATE TABLE IF NOT EXISTS %s ( + chunk_id UUID PRIMARY KEY REFERENCES message_chunks(id) ON DELETE CASCADE, + embedding vector(%d), + content_hash text not null + created_at TIMESTAMPTZ NOT NULL DEFAULT now() + )` diff --git a/hello.go b/hello.go index 48df291..4e97a50 100644 --- a/hello.go +++ b/hello.go @@ -2,16 +2,19 @@ package main import ( "context" + "errors" "log/slog" "os" + "time" tea "charm.land/bubbletea/v2" "git.estatecloud.org/radumaco/souvenir/config" + "git.estatecloud.org/radumaco/souvenir/db/embed" "git.estatecloud.org/radumaco/souvenir/db/history" + "git.estatecloud.org/radumaco/souvenir/llm" ui "git.estatecloud.org/radumaco/souvenir/ui" ) - func main() { ctx := context.Background() config, err := config.Load(".config.json") @@ -28,6 +31,10 @@ func main() { os.Exit(1) } defer db.Close() + es, _ := setupEmbedding(config, logger) + if es != nil { + defer es.Close() + } p := tea.NewProgram(ui.InitialModel(ctx, *config, *db)) if _, err := p.Run(); err != nil { @@ -35,3 +42,46 @@ func main() { os.Exit(1) } } + +func setupEmbedding(cfg *config.Config, logger *slog.Logger) (*embed.EmbedStore, error) { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + embedder := llm.NewEmbedder(cfg.Embedding) + embedStore, err := embed.NewEmbedStore(ctx, cfg.Db, *embedder) + if err != nil { + logger.Error("Could not initialize embedder", "error", err) + return nil, err + } + ticker := time.NewTicker(time.Duration(cfg.Embedding.Interval) * time.Minute) + go func() { + defer ticker.Stop() + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + runEmbedding(ctx, embedStore, logger) + } + } + }() + return embedStore, nil +} + +func runEmbedding(ctx context.Context, em *embed.EmbedStore, logger *slog.Logger) error { + convs, err := em.GetConversationsToEmbed(ctx) + if err != nil { + logger.Error("Could not get conversations to embed", "error", err) + return err + } + var errs []error + for _, conv := range convs { + if err := em.EmbedConversation(ctx, conv); err != nil { + logger.Error("Could not embed conversation", "conversation_id", conv, "error", err) + errs = append(errs, err) + } + } + if errs != nil { + return errors.Join(errs...) + } + return nil +} diff --git a/ui/chat.go b/ui/tui.go similarity index 74% rename from ui/chat.go rename to ui/tui.go index 7834cde..dbb9b8e 100644 --- a/ui/chat.go +++ b/ui/tui.go @@ -7,6 +7,7 @@ import ( "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" @@ -26,30 +27,36 @@ const ( ) type uiModel struct { - 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 - statusStyle lipgloss.Style - client llm.ChatClient - statusMessage string - commands commandList - history history.DbClient - ctx context.Context - width int - height int - waiting bool - err error + 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 agentResponseMessage struct { - response []model.Message - err error +type streamEventMessage struct { + event llm.StreamEvent +} + +type streamClosedMessage struct { } type conversationSavedMessage struct { @@ -69,7 +76,7 @@ func InitialModel(ctx context.Context, config config.Config, client history.DbCl ta.SetVirtualCursor(false) ta.Focus() - ta.Prompt = "⦀ " + ta.Prompt = "｜" ta.CharLimit = 280 ta.SetWidth(30) @@ -88,10 +95,14 @@ func InitialModel(ctx context.Context, config config.Config, client history.DbCl 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")), @@ -114,6 +125,9 @@ func (m uiModel) Init() tea.Cmd { 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: @@ -195,7 +209,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } } - if after, ok :=strings.CutPrefix(input, "/"); ok { + if after, ok := strings.CutPrefix(input, "/"); ok { name := after for _, c := range m.commands.available { if c.name == name { @@ -220,8 +234,8 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages)) m.textarea.Reset() m.viewport.GotoBottom() - m.waiting = true - return m, tea.Batch(m.callAgent(input, m.conversation.Messages), m.saveConversation()) + 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) @@ -242,7 +256,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { case conversationRenamedMessage: m.logger.Debug("Received messageRenamedMessage") - m.waiting = false + m.startWait() if msg.err != nil { m.logger.Error("There was an error", "error", msg.err) m.setErrorMessage(msg.err.Error()) @@ -252,31 +266,49 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.conversation.Summary = msg.summary return m, nil - 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) - m.setErrorMessage(msg.err.Error()) + 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) } - for _, r := range resp { - r.Seq = len(m.conversation.Messages) + 1 - m.logger.Debug("Appending response message", "message", r) - m.conversation.Messages = append(m.conversation.Messages, r) + if msg.event.Reasoning != "" { + m.thinkingBuffer.WriteString(msg.event.Reasoning) } - var messages = m.renderMessages() - m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages)) - m.viewport.GotoBottom() - return m, m.saveConversation() + + 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.waiting = false + m.stopWait() m.focus = focusChat case focusHistory: m.logger.Debug("Received pickerChosenMsg", "history", msg.id) @@ -285,7 +317,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.logger.Error("Could not retrieve conversation", "id", msg.id) } m.conversation = conv - m.waiting = false + m.stopWait() m.focus = focusChat var messages = m.renderMessages() m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages)) @@ -294,13 +326,13 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { case pickerDismissMsg: m.logger.Debug("Received pickerDismissMsg") - m.waiting = false + m.stopWait() 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.logger.Error("Could not save message", "error", msg.err) m.setErrorMessage(msg.err.Error()) return m, nil } @@ -327,6 +359,9 @@ func (m uiModel) View() tea.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()) } @@ -344,6 +379,9 @@ func (m uiModel) View() tea.View { 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) } @@ -353,11 +391,25 @@ func (m uiModel) View() tea.View { return view } -func (m uiModel) callAgent(input string, messages []model.Message) tea.Cmd { - m.logger.Debug("Calling agent", "query", input, "messages", messages) +func (m uiModel) callAgent(messages []model.Message) tea.Cmd { + m.logger.Debug("Calling agent", "messages", messages) return func() tea.Msg { - resp, err := m.client.Query(input, messages) - return agentResponseMessage{response: resp, err: err} + 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} } } @@ -453,9 +505,49 @@ func (m *uiModel) renderMessages() string { 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) }