550ba526c4
streaming.WriteCompaction packages a 'compaction' event with the kept/older turn counts, summary tokens and provider-reported window/used tokens so the client can hint 'context optimized' to the user without parsing the stream body. streamChat runs Compact before BuildMessages and writes the event right after start, ensuring the client sees it before any chunk is emitted. Add a runner test that exercises limitRAGContext to keep the system prompt + RAG block under the configured window.
443 lines
No EOL
13 KiB
Go
443 lines
No EOL
13 KiB
Go
package server
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"log/slog"
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/VictorVargas/rony-chat-bot/internal/agent"
|
|
"github.com/VictorVargas/rony-chat-bot/internal/config"
|
|
"github.com/VictorVargas/rony-chat-bot/internal/portfolio"
|
|
"github.com/VictorVargas/rony-chat-bot/internal/streaming"
|
|
)
|
|
|
|
type Handlers struct {
|
|
cfg *config.Config
|
|
runner *agent.Runner
|
|
store *portfolio.Store
|
|
version string
|
|
}
|
|
|
|
func NewHandlers(cfg *config.Config, runner *agent.Runner, store *portfolio.Store, version string) *Handlers {
|
|
return &Handlers{cfg: cfg, runner: runner, store: store, version: version}
|
|
}
|
|
|
|
type ChatRequest struct {
|
|
Messages []ChatMessage `json:"messages"`
|
|
Stream *bool `json:"stream,omitempty"`
|
|
// ConversationID is optional. If empty, the server creates a new
|
|
// conversation and returns its ID in the response (or in the SSE
|
|
// `start` event). Pass an existing ID to continue a previous thread.
|
|
ConversationID string `json:"conversation_id,omitempty"`
|
|
}
|
|
|
|
type ChatMessage struct {
|
|
Role string `json:"role"`
|
|
Content string `json:"content"`
|
|
}
|
|
|
|
type ChatResponse struct {
|
|
ConversationID string `json:"conversation_id"`
|
|
Content string `json:"content"`
|
|
Sources []string `json:"sources,omitempty"`
|
|
Usage streaming.Usage `json:"usage"`
|
|
}
|
|
|
|
func (h *Handlers) Chat(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPost {
|
|
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
var req ChatRequest
|
|
if err := json.NewDecoder(http.MaxBytesReader(w, r.Body, 1<<20)).Decode(&req); err != nil {
|
|
http.Error(w, "invalid JSON body: "+err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
if err := req.validate(); err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
stream := true
|
|
if req.Stream != nil {
|
|
stream = *req.Stream
|
|
}
|
|
|
|
// Resolve or create the conversation. If the client passed a non-empty
|
|
// ID we use it as-is; otherwise we mint a new one.
|
|
ctx := r.Context()
|
|
convID := req.ConversationID
|
|
if convID == "" {
|
|
var err error
|
|
convID, err = h.store.CreateConversation(ctx)
|
|
if err != nil {
|
|
http.Error(w, "create conversation: "+err.Error(), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
}
|
|
|
|
// Persist the incoming user message. Even if the LLM fails afterward
|
|
// the user sees their question in the conversation history.
|
|
if userMsg := lastUserMessage(req.Messages); userMsg != "" {
|
|
if err := h.store.SaveMessage(ctx, convID, "user", userMsg, nil); err != nil {
|
|
slog.Error("save user message", "err", err)
|
|
}
|
|
}
|
|
|
|
history := toAgentMessages(req.Messages)
|
|
if stream {
|
|
h.streamChat(w, r, history, convID)
|
|
return
|
|
}
|
|
h.completeChat(w, r, history, convID)
|
|
}
|
|
|
|
func lastUserMessage(msgs []ChatMessage) string {
|
|
for i := len(msgs) - 1; i >= 0; i-- {
|
|
if msgs[i].Role == "user" {
|
|
return msgs[i].Content
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func toAgentMessages(in []ChatMessage) []agent.Message {
|
|
out := make([]agent.Message, len(in))
|
|
for i, m := range in {
|
|
out[i] = agent.Message{
|
|
Role: agent.Role(m.Role),
|
|
Content: m.Content,
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func (req *ChatRequest) validate() error {
|
|
if len(req.Messages) == 0 {
|
|
return errors.New("messages must not be empty")
|
|
}
|
|
for i, m := range req.Messages {
|
|
switch m.Role {
|
|
case "user", "assistant", "system":
|
|
default:
|
|
return fmt.Errorf("messages[%d].role %q is invalid", i, m.Role)
|
|
}
|
|
if strings.TrimSpace(m.Content) == "" {
|
|
return fmt.Errorf("messages[%d].content is empty", i)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// streamChat drives the LLM agent and forwards each chunk to the SSE
|
|
// stream using the protocol described in docs/architecture.md §3.2.
|
|
// convID is the persistent conversation ID — passed in from Chat() after
|
|
// resolve-or-create. The assistant reply is saved on success.
|
|
func (h *Handlers) streamChat(w http.ResponseWriter, r *http.Request, history []agent.Message, convID string) {
|
|
w.Header().Set("Content-Type", "text/event-stream")
|
|
w.Header().Set("Cache-Control", "no-cache")
|
|
w.Header().Set("Connection", "keep-alive")
|
|
w.Header().Set("X-Accel-Buffering", "no")
|
|
w.WriteHeader(http.StatusOK)
|
|
|
|
flusher, ok := w.(http.Flusher)
|
|
if !ok {
|
|
http.Error(w, "streaming unsupported", http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
ctx := r.Context()
|
|
if err := streaming.WriteStart(w, convID); err != nil {
|
|
slog.Error("sse start failed", "err", err)
|
|
return
|
|
}
|
|
flusher.Flush()
|
|
|
|
// Auto-compact the older part of the conversation when the previous
|
|
// turn's input tokens have crossed the configured threshold. Done
|
|
// before RAG so the LLM's compaction summary call is counted in the
|
|
// next request's reported usage, not the current one. A failure here
|
|
// is logged and swallowed: the request still proceeds with the
|
|
// original (uncompacted) history.
|
|
history, _ = h.runner.Compact(ctx, history)
|
|
if stats := h.runner.LastCompaction(); stats.Happened {
|
|
_ = streaming.WriteCompaction(w, map[string]any{
|
|
"older_turns": stats.OlderTurns,
|
|
"kept_turns": stats.KeptTurns,
|
|
"summary_tokens": stats.SummaryTokens,
|
|
"window_tokens": stats.WindowTokens,
|
|
"used_tokens": stats.UsedTokens,
|
|
})
|
|
}
|
|
|
|
// Report which RAG sources were used (search happens in BuildMessages).
|
|
_, ragContext, err := h.runner.BuildMessages(ctx, history)
|
|
if err != nil {
|
|
_ = streaming.WriteError(w, "rag: "+err.Error())
|
|
return
|
|
}
|
|
var sources []string
|
|
if ragContext != "" {
|
|
sources = extractSources(ragContext)
|
|
_ = streaming.WriteSources(w, sources)
|
|
}
|
|
|
|
// Stream from the model, collecting the full reply as we go.
|
|
var full strings.Builder
|
|
for chunk, err := range h.runner.Stream(ctx, history) {
|
|
if err != nil {
|
|
slog.Error("llm stream", "err", err)
|
|
_ = streaming.WriteError(w, "llm: "+err.Error())
|
|
return
|
|
}
|
|
if chunk.Delta != "" {
|
|
full.WriteString(chunk.Delta)
|
|
if err := streaming.WriteChunk(w, chunk.Delta); err != nil {
|
|
return
|
|
}
|
|
}
|
|
if chunk.FinishReason != "" && chunk.Usage.TotalTokens > 0 {
|
|
h.persistAssistant(ctx, convID, full.String(), sources)
|
|
_ = streaming.WriteDone(w, streaming.Usage{
|
|
InputTokens: chunk.Usage.InputTokens,
|
|
OutputTokens: chunk.Usage.OutputTokens,
|
|
})
|
|
return
|
|
}
|
|
}
|
|
// Final usage chunk may come on the last iteration; if we never saw a
|
|
// finish_reason + usage in-band, surface what we recorded.
|
|
usage := h.runner.LastUsage()
|
|
h.persistAssistant(ctx, convID, full.String(), sources)
|
|
_ = streaming.WriteDone(w, streaming.Usage{
|
|
InputTokens: usage.InputTokens,
|
|
OutputTokens: usage.OutputTokens,
|
|
})
|
|
}
|
|
|
|
// persistAssistant saves the full assistant reply. Best-effort: a failure
|
|
// here doesn't fail the user's request (the response is already streamed).
|
|
func (h *Handlers) persistAssistant(ctx context.Context, convID, content string, sources []string) {
|
|
if strings.TrimSpace(content) == "" {
|
|
return
|
|
}
|
|
if err := h.store.SaveMessage(ctx, convID, "assistant", content, sources); err != nil {
|
|
slog.Error("save assistant message", "err", err, "conv", convID)
|
|
}
|
|
}
|
|
|
|
func (h *Handlers) completeChat(w http.ResponseWriter, r *http.Request, history []agent.Message, convID string) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
ctx := r.Context()
|
|
// Apply compaction before streaming — same as streamChat. The
|
|
// compaction event isn't surfaced in the non-streaming JSON response
|
|
// (would be redundant noise), but LastCompaction is still recorded so
|
|
// any future caller can introspect it.
|
|
history, _ = h.runner.Compact(ctx, history)
|
|
var full strings.Builder
|
|
var sources []string
|
|
for chunk, err := range h.runner.Stream(ctx, history) {
|
|
if err != nil {
|
|
http.Error(w, "llm: "+err.Error(), http.StatusBadGateway)
|
|
return
|
|
}
|
|
full.WriteString(chunk.Delta)
|
|
}
|
|
usage := h.runner.LastUsage()
|
|
_, ragContext, _ := h.runner.BuildMessages(ctx, history)
|
|
if ragContext != "" {
|
|
sources = extractSources(ragContext)
|
|
}
|
|
h.persistAssistant(ctx, convID, full.String(), sources)
|
|
resp := ChatResponse{
|
|
ConversationID: convID,
|
|
Content: full.String(),
|
|
Usage: streaming.Usage{
|
|
InputTokens: usage.InputTokens,
|
|
OutputTokens: usage.OutputTokens,
|
|
},
|
|
Sources: sources,
|
|
}
|
|
_ = json.NewEncoder(w).Encode(resp)
|
|
}
|
|
|
|
func extractSources(ragContext string) []string {
|
|
// The formatted hits look like "### [N] projectID — section\n..."
|
|
// so we scan the lines for that prefix and dedupe by projectID.
|
|
seen := map[string]bool{}
|
|
var out []string
|
|
for _, line := range strings.Split(ragContext, "\n") {
|
|
if !strings.HasPrefix(line, "### [") {
|
|
continue
|
|
}
|
|
// between "### [N] " and " — "
|
|
rest := strings.TrimPrefix(line, "### [")
|
|
rest = rest[strings.Index(rest, " ")+1:]
|
|
project := rest
|
|
if i := strings.Index(rest, " — "); i > 0 {
|
|
project = rest[:i]
|
|
}
|
|
if !seen[project] {
|
|
seen[project] = true
|
|
out = append(out, project)
|
|
}
|
|
}
|
|
return out
|
|
}
|
|
|
|
func truncate(s string, n int) string {
|
|
if len(s) <= n {
|
|
return s
|
|
}
|
|
return s[:n] + "…"
|
|
}
|
|
|
|
func (h *Handlers) Info(w http.ResponseWriter, _ *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
p := h.cfg.DefaultProvider()
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"name": h.cfg.Persona.Name,
|
|
"version": h.version,
|
|
"provider": p.Type,
|
|
"model": firstNonEmpty(p.Model, p.ModelPath),
|
|
"rag": h.cfg.RAG.Enabled,
|
|
"top_k": h.cfg.RAG.TopK,
|
|
"started": time.Now().UTC().Format(time.RFC3339),
|
|
})
|
|
}
|
|
|
|
func (h *Handlers) Reindex(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodPost {
|
|
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
start := time.Now()
|
|
dbPath := h.cfg.RAG.DBPath
|
|
if dbPath == "" {
|
|
http.Error(w, "rag.db_path not configured", http.StatusBadRequest)
|
|
return
|
|
}
|
|
dur, files, chunks, err := portfolio.ReindexOnDisk(dbPath, h.cfg.RAG.DataPath, portfolio.DefaultChunkerConfig())
|
|
if err != nil {
|
|
http.Error(w, "reindex failed: "+err.Error(), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
_ = dur
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"indexed_files": files,
|
|
"total_chunks": chunks,
|
|
"duration_ms": time.Since(start).Milliseconds(),
|
|
"db_path": dbPath,
|
|
})
|
|
}
|
|
|
|
// ---- Conversation REST endpoints -------------------------------------------
|
|
|
|
// GetConversation returns the full history of one conversation. The
|
|
// conversation ID is treated as a bearer token: anyone who knows it can
|
|
// read the history. For a public bot this is fine; for private contexts
|
|
// add auth at the proxy layer.
|
|
func (h *Handlers) GetConversation(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodGet {
|
|
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
id := conversationIDFromPath(r.URL.Path)
|
|
if id == "" {
|
|
http.Error(w, "missing conversation id", http.StatusBadRequest)
|
|
return
|
|
}
|
|
conv, err := h.store.GetConversation(r.Context(), id)
|
|
if errors.Is(err, portfolio.ErrConversationNotFound) {
|
|
http.Error(w, "conversation not found", http.StatusNotFound)
|
|
return
|
|
}
|
|
if err != nil {
|
|
http.Error(w, "get conversation: "+err.Error(), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_ = json.NewEncoder(w).Encode(conv)
|
|
}
|
|
|
|
// ListConversations returns the most recent N conversation summaries.
|
|
// Useful for a "show my chats" sidebar in the widget or a custom UI.
|
|
func (h *Handlers) ListConversations(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodGet {
|
|
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
limit := 50
|
|
if s := r.URL.Query().Get("limit"); s != "" {
|
|
if n, err := strconv.Atoi(s); err == nil && n > 0 && n <= 200 {
|
|
limit = n
|
|
}
|
|
}
|
|
items, err := h.store.ListConversations(r.Context(), limit)
|
|
if err != nil {
|
|
http.Error(w, "list conversations: "+err.Error(), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
w.Header().Set("Content-Type", "application/json")
|
|
if items == nil {
|
|
items = []portfolio.ConversationSummary{}
|
|
}
|
|
_ = json.NewEncoder(w).Encode(map[string]any{
|
|
"conversations": items,
|
|
"count": len(items),
|
|
})
|
|
}
|
|
|
|
// DeleteConversation removes a conversation and all its messages.
|
|
func (h *Handlers) DeleteConversation(w http.ResponseWriter, r *http.Request) {
|
|
if r.Method != http.MethodDelete {
|
|
http.Error(w, "method not allowed", http.StatusMethodNotAllowed)
|
|
return
|
|
}
|
|
id := conversationIDFromPath(r.URL.Path)
|
|
if id == "" {
|
|
http.Error(w, "missing conversation id", http.StatusBadRequest)
|
|
return
|
|
}
|
|
err := h.store.DeleteConversation(r.Context(), id)
|
|
if errors.Is(err, portfolio.ErrConversationNotFound) {
|
|
http.Error(w, "conversation not found", http.StatusNotFound)
|
|
return
|
|
}
|
|
if err != nil {
|
|
http.Error(w, "delete conversation: "+err.Error(), http.StatusInternalServerError)
|
|
return
|
|
}
|
|
w.WriteHeader(http.StatusNoContent)
|
|
}
|
|
|
|
// conversationIDFromPath pulls the ID out of "/api/conversations/{id}".
|
|
// chi does this with chi.URLParam(r, "id"); for clarity (and so this
|
|
// handler works even without chi) we do it by hand.
|
|
func conversationIDFromPath(path string) string {
|
|
const prefix = "/api/conversations/"
|
|
if !strings.HasPrefix(path, prefix) {
|
|
return ""
|
|
}
|
|
id := strings.TrimPrefix(path, prefix)
|
|
// strip trailing slash and any further segments
|
|
if i := strings.IndexByte(id, '/'); i >= 0 {
|
|
id = id[:i]
|
|
}
|
|
return id
|
|
}
|
|
|
|
func firstNonEmpty(vals ...string) string {
|
|
for _, v := range vals {
|
|
if v != "" {
|
|
return v
|
|
}
|
|
}
|
|
return ""
|
|
} |