From e57b8f428e79df6f50f3100134720070004eedd1 Mon Sep 17 00:00:00 2001 From: "Radu Macocian (admac)" Date: Sat, 27 Jun 2026 22:06:06 +0200 Subject: [PATCH] Added streaming --- db/history/history.go | 9 ++- db/schema.go | 2 +- llm/chatclient.go | 113 +++++++++++++++++++++----------- llm/streamclient.go | 146 ++++++++++++++++++++++++++++++++++++++++++ ui/tui.go | 78 +++++++++++++--------- 5 files changed, 274 insertions(+), 74 deletions(-) create mode 100644 llm/streamclient.go diff --git a/db/history/history.go b/db/history/history.go index 162d207..bd60aeb 100644 --- a/db/history/history.go +++ b/db/history/history.go @@ -161,7 +161,7 @@ func (cl DbClient) SaveConversation(ctx context.Context, conv model.Conversation if _, err := cl.pool.Exec(ctx, ` UPDATE conversations - SET title = $1, summary = $2 + SET title = $1, summary = $2, updated_at = now() 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 @@ -249,9 +249,12 @@ func (cl DbClient) saveMessage(ctx context.Context, conversation_id string, mess if _, err := cl.pool.Exec(ctx, ` INSERT INTO message_chunks (conversation_id, message_id, content, content_hash) - VALUES ($1, $2, $3, $4) + VALUES ($1, $2, $3, $4); + UPDATE messages + SET updated_at = now() + WHERE id = $5 `, - conversation_id, message.Id, chunk, hex.EncodeToString(hash[:])); err != nil { + conversation_id, message.Id, chunk, hex.EncodeToString(hash[:]), message.Id); 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 index 0c303f1..4154aa0 100644 --- a/db/schema.go +++ b/db/schema.go @@ -39,6 +39,6 @@ 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 + content_hash text not null, created_at TIMESTAMPTZ NOT NULL DEFAULT now() )` diff --git a/llm/chatclient.go b/llm/chatclient.go index a271192..8f1af3e 100644 --- a/llm/chatclient.go +++ b/llm/chatclient.go @@ -8,6 +8,7 @@ import ( "log/slog" "net/http" "strconv" + "strings" "time" config "git.estatecloud.org/radumaco/souvenir/config" @@ -29,6 +30,16 @@ type ChatRequest struct { Model string `json:"model"` Messages []model.Message `json:"messages"` Tools []config.Tool `json:"tools,omitempty"` + Stream bool `json:"stream,omitempty"` + Kwargs Kwargs `json:"chat_template_kwargs"` +} + +type Kwargs struct { + Thinking bool `json:"enable_thinking"` +} + +type Reasoning struct { + Effort string `json:"effort"` } // Response types // @@ -129,7 +140,12 @@ func (cl ChatClient) Rename(messages []model.Message) (RenameResponse, error) { Role: "system", Content: "Reply with ONLY JSON with the form: {\"title\":, \"summary\":<summary>}. Title <= 10 words. Summary 2-3 sentances.", }}, messages...), + Kwargs: Kwargs{Thinking: false}, } + request.Messages = append(request.Messages, model.Message{ + Role: "user", + Content: "Give me a title and description for this conversation. Reply with just a json and nothing else. The format is exactly {\"title\":<title>, \"summary\":<summary>}. Add nothing. Change nothing", + }) cl.logger.Debug("Calling rename", "request", request) out, err := cl.Call(request) if err != nil { @@ -137,63 +153,45 @@ func (cl ChatClient) Rename(messages []model.Message) (RenameResponse, error) { return RenameResponse{}, err } var meta RenameResponse - if err := json.Unmarshal([]byte(out[0].Content), &meta); err != nil { - cl.logger.Error("Could not parse the rename response", "content", out[0].Content) - return meta, err + cl.logger.Debug("Received rename answer", "content", out[0].Content) + if idx := strings.IndexByte(out[0].Content, '{'); idx >= 0 { + if lidx := strings.LastIndexByte(out[0].Content, '}'); lidx <= len(out[0].Content) { + trimmed := out[0].Content[idx:lidx + 1] + cl.logger.Debug("Trimmed answer", "content", trimmed) + if err := json.Unmarshal([]byte(trimmed), &meta); err != nil { + cl.logger.Error("Could not parse the rename response", "content", out[0].Content) + return meta, err + } + return meta, nil + } } - return meta, nil + cl.logger.Error("Response is not standard json", "content", out[0].Content) + return meta, errors.New("Response is not standard json") } -func (cl ChatClient) Query(query string, messages []model.Message) ([]model.Message, error) { +func (cl ChatClient) Query(messages []model.Message) ([]model.Message, error) { requestBody := ChatRequest{ Model: cl.Cfg.Llm.Model, Messages: messages, Tools: config.AvailableTools(), + Stream: true, + Kwargs: Kwargs{Thinking: false}, } return cl.Call(requestBody) } func (cl ChatClient) Call(requestBody ChatRequest) ([]model.Message, error) { - serialized, err := json.Marshal(requestBody) - if err != nil { - cl.logger.Error("There was an error marshelling the request body", - "request", requestBody, - "error", err) - return []model.Message{}, err - } + resp, err := cl.CallApi(requestBody) - body := bytes.NewBuffer(serialized) - req, err := http.NewRequest("POST", cl.Cfg.Api.Url+COMPLETIONS_API, body) if err != nil { - cl.logger.Error("There was an error creating the request", "error", err.Error()) - return []model.Message{}, err - } - req.Header.Set("Content-Type", "application/json") - if cl.Cfg.Api.Key != "" { - req.Header.Set("Authorization", "Bearer "+cl.Cfg.Api.Key) - } - cl.logger.Debug("Executing llm call") - resp, err := cl.client.Do(req) - if err != nil { - cl.logger.Error("There was an error during http call", "error", err.Error()) - return []model.Message{}, err + cl.logger.Error("There was an error during the API call", "error", err) + return nil, err } defer resp.Body.Close() - cl.logger.Debug("LLM Call returned", "response", resp) - data, err := io.ReadAll(resp.Body) if err != nil { cl.logger.Error("There was an error reading the response body", "error", err.Error()) - return []model.Message{}, err - } - if resp.StatusCode != 200 && resp.StatusCode != 202 { - cl.logger.Error("Call returned non 200 status", - "statusCode", resp.StatusCode, - "status", resp.Status) - return []model.Message{}, - errors.New("Call returned invalid status " + - strconv.Itoa(resp.StatusCode) + - resp.Status) + return nil, err } cl.logger.Debug("Received data", "data", data) @@ -215,3 +213,42 @@ func (cl ChatClient) Call(requestBody ChatRequest) ([]model.Message, error) { return result, nil } + +func (cl ChatClient) CallApi(requestBody ChatRequest) (*http.Response, error) { + serialized, err := json.Marshal(requestBody) + if err != nil { + cl.logger.Error("There was an error marshelling the request body", + "request", requestBody, + "error", err) + return nil, err + } + + body := bytes.NewBuffer(serialized) + req, err := http.NewRequest("POST", cl.Cfg.Api.Url+COMPLETIONS_API, body) + if err != nil { + cl.logger.Error("There was an error creating the request", "error", err.Error()) + return nil, err + } + req.Header.Set("Content-Type", "application/json") + if cl.Cfg.Api.Key != "" { + req.Header.Set("Authorization", "Bearer "+cl.Cfg.Api.Key) + } + cl.logger.Debug("Executing llm call") + resp, err := cl.client.Do(req) + if err != nil { + cl.logger.Error("There was an error during http call", "error", err.Error()) + return nil, err + } + cl.logger.Debug("LLM Call returned", "response", resp.Status) + + if resp.StatusCode != 200 && resp.StatusCode != 202 { + cl.logger.Error("Call returned non 200 status", + "statusCode", resp.StatusCode, + "status", resp.Status) + return nil, + errors.New("Call returned invalid status " + + strconv.Itoa(resp.StatusCode) + + resp.Status) + } + return resp, nil +} diff --git a/llm/streamclient.go b/llm/streamclient.go new file mode 100644 index 0000000..76d00f8 --- /dev/null +++ b/llm/streamclient.go @@ -0,0 +1,146 @@ +package llm + +import ( + "bufio" + "encoding/json" + "net/http" + "strings" + + config "git.estatecloud.org/radumaco/souvenir/config" + "git.estatecloud.org/radumaco/souvenir/model" +) + +type StreamEvent struct { + Reasoning string + Content string + Messages []model.Message + Done bool + Err error +} + +type ChatStreamResponse struct { + Choices []StreamChoice `json:"choices"` +} + +type StreamChoice struct { + Delta StreamDelta `json:"delta"` + FinishReason *string `json:"finish_reason"` +} + +type StreamDelta struct { + Role string `json:"role,omitempty"` + Content string `json:"content,omitempty"` + ReasoningContent string `json:"reasoning_content,omitempty"` + ToolCalls []ToolCallDelta `json:"tool_calls,omitempty"` +} + +type ToolCallDelta struct { + Index int `json:"index"` + Id string `json:"id,omitempty"` + Type string `json:"type,omitempty"` + Function struct { + Name string `json:"name,omitempty"` + Arguments string `json:"arguments,omitempty"` + } `json:"function"` +} + +func (cl ChatClient) QueryStream(messages []model.Message) (<-chan StreamEvent, error) { + requestBody := ChatRequest{ + Model: cl.Cfg.Llm.Model, + Messages: messages, + Tools: config.AvailableTools(), + Stream: true, + Kwargs: Kwargs{Thinking: false}, + } + return cl.CallStream(requestBody) +} + +func (cl ChatClient) CallStream(requestBody ChatRequest) (<-chan StreamEvent, error) { + resp, err := cl.CallApi(requestBody) + if err != nil { + cl.logger.Error("There was an error during the API call", "error", err) + return nil, err + } + events := make(chan StreamEvent) + go cl.readStream(resp, events) + return events, nil +} + +func (cl ChatClient) readStream(resp *http.Response, events chan<- StreamEvent) { + cl.logger.Debug("Started readStrea") + defer resp.Body.Close() + defer close(events) + + scanner := bufio.NewScanner(resp.Body) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + + var content strings.Builder + toolAcc := map[int]*model.ToolCall{} + var order []int + + for scanner.Scan() { + line := scanner.Text() + cl.logger.Debug("Stream got new line", "line", line) + + if !strings.HasPrefix(line, "data:") { + continue + } + + payload := strings.TrimSpace(strings.TrimPrefix(line, "data:")) + if payload == "[DONE]" { + cl.logger.Debug("Stream done") + break + } + + var chunk ChatStreamResponse + if err := json.Unmarshal([]byte(payload), &chunk); err != nil { + cl.logger.Error("Error parsing stream chunk", "payload", payload, "error", err) + events <- StreamEvent{Err: err} + return + } + + if len(chunk.Choices) == 0 { + continue + } + + delta := chunk.Choices[0].Delta + if delta.ReasoningContent != "" { + cl.logger.Debug("Stream got reasoning content", "content", delta.ReasoningContent) + events <- StreamEvent{Reasoning: delta.ReasoningContent} + } + if delta.Content != "" { + cl.logger.Debug("Stream got content", "content", delta.Content) + content.WriteString(delta.Content) + events <- StreamEvent{Content: delta.Content} + } + + for _, tc := range delta.ToolCalls { + acc, ok := toolAcc[tc.Index] + if !ok { + acc = &model.ToolCall{} + toolAcc[tc.Index] = acc + order = append(order, tc.Index) + } + if tc.Id != "" { + acc.Id = tc.Id + } + if tc.Function.Name != "" { + acc.FunctionName = tc.Function.Name + } + acc.FunctionArguments += tc.Function.Arguments + } + } + + if err := scanner.Err(); err != nil { + cl.logger.Error("error reading stream", "error", err) + events <- StreamEvent{Err: err} + return + } + + final := model.Message{Role: "assistant", Content: content.String()} + for _, idx := range order { + final.ToolCalls = append(final.ToolCalls, *toolAcc[idx]) + } + + events <- StreamEvent{Done: true, Messages: []model.Message{final}} +} diff --git a/ui/tui.go b/ui/tui.go index dbb9b8e..5610d67 100644 --- a/ui/tui.go +++ b/ui/tui.go @@ -34,8 +34,8 @@ type uiModel struct { textarea textarea.Model focus focusState picker picker - thinkingBuffer strings.Builder - answerBuffer strings.Builder + thinkingBuffer *strings.Builder + answerBuffer *strings.Builder senderStyle lipgloss.Style agentStyle lipgloss.Style errorStyle lipgloss.Style @@ -59,6 +59,10 @@ type streamEventMessage struct { type streamClosedMessage struct { } +type streamStartedMessage struct { + ch <-chan llm.StreamEvent +} + type conversationSavedMessage struct { conversation model.Conversation err error @@ -76,7 +80,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) @@ -96,22 +100,24 @@ func InitialModel(ctx context.Context, config config.Config, client history.DbCl ta.KeyMap.InsertNewline.SetEnabled(false) sp := spinner.New() - sp.Spinner = spinner.Monkey + sp.Spinner = spinner.Dot 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, + 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")), + thinkingBuffer: &strings.Builder{}, + answerBuffer: &strings.Builder{}, + 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) @@ -125,9 +131,6 @@ 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: @@ -256,7 +259,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { case conversationRenamedMessage: m.logger.Debug("Received messageRenamedMessage") - m.startWait() + m.stopWait() if msg.err != nil { m.logger.Error("There was an error", "error", msg.err) m.setErrorMessage(msg.err.Error()) @@ -266,6 +269,11 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.conversation.Summary = msg.summary return m, nil + case streamStartedMessage: + m.logger.Debug("Received streamStartedMessage") + m.streamCh = msg.ch + return m, waitForEvent(m.streamCh) + case streamEventMessage: m.logger.Debug("Received streamEventMessage") if msg.event.Err != nil { @@ -286,13 +294,12 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { 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, func() tea.Msg { return streamClosedMessage{} } } + var messages = m.renderMessages() + m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages)) + m.viewport.GotoBottom() return m, waitForEvent(m.streamCh) case streamClosedMessage: @@ -301,6 +308,10 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { m.answerBuffer.Reset() m.thinkingBuffer.Reset() + var messages = m.renderMessages() + m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages)) + m.viewport.GotoBottom() + return m, nil case pickerChosenMsg: @@ -308,7 +319,6 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { 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) @@ -317,7 +327,6 @@ 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.stopWait() m.focus = focusChat var messages = m.renderMessages() m.viewport.SetContent(lipgloss.NewStyle().Width(m.viewport.Width()).Render(messages)) @@ -338,6 +347,10 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { } m.conversation = msg.conversation } + var spinnerCmd tea.Cmd + if m.waiting { + m.spinner, spinnerCmd = m.spinner.Update(msg) + } if m.focus != focusChat { picker, cmd := m.picker.Update(msg) @@ -345,7 +358,7 @@ func (m uiModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) { return m, cmd } - return m, nil + return m, spinnerCmd } func (m uiModel) View() tea.View { @@ -399,7 +412,7 @@ func (m uiModel) callAgent(messages []model.Message) tea.Cmd { m.logger.Error("Error while calling stream query", "error", err) return streamClosedMessage{} } - return waitForEvent(resp) + return streamStartedMessage{ch: resp} } } @@ -409,7 +422,7 @@ func waitForEvent(ch <-chan llm.StreamEvent) tea.Cmd { if !ok { return streamClosedMessage{} } - return streamEventMessage{ev} + return streamEventMessage{event: ev} } } @@ -425,6 +438,7 @@ 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) + m.startWait() if err != nil { return conversationRenamedMessage{title: meta.Title, summary: meta.Summary, err: err} } @@ -544,7 +558,7 @@ func (m *uiModel) startWait() { func (m *uiModel) stopWait() { if m.waiting { m.waiting = false - m.viewport.SetHeight(m.viewport.Height() - 1) + m.viewport.SetHeight(m.viewport.Height() + 1) } }