rony-chat-bot/internal/server/handlers.go

package server

import (
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"net/http"
	"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"`
}

type ChatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type ChatResponse struct {
	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
	}
	history := toAgentMessages(req.Messages)
	if stream {
		h.streamChat(w, r, history)
		return
	}
	h.completeChat(w, r, history)
}

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.
func (h *Handlers) streamChat(w http.ResponseWriter, r *http.Request, history []agent.Message) {
	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()
	convID := newConvID()
	if err := streaming.WriteStart(w, convID); err != nil {
		slog.Error("sse start failed", "err", err)
		return
	}
	flusher.Flush()

	// 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
	}
	if ragContext != "" {
		sources := extractSources(ragContext)
		_ = streaming.WriteSources(w, sources)
	}

	// Stream from the model.
	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 != "" {
			if err := streaming.WriteChunk(w, chunk.Delta); err != nil {
				return
			}
		}
		if chunk.FinishReason != "" && chunk.Usage.TotalTokens > 0 {
			_ = 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()
	_ = streaming.WriteDone(w, streaming.Usage{
		InputTokens:  usage.InputTokens,
		OutputTokens: usage.OutputTokens,
	})
}

func (h *Handlers) completeChat(w http.ResponseWriter, r *http.Request, history []agent.Message) {
	w.Header().Set("Content-Type", "application/json")
	var full strings.Builder
	for chunk, err := range h.runner.Stream(r.Context(), 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(r.Context(), history)
	resp := ChatResponse{
		Content: full.String(),
		Usage: streaming.Usage{
			InputTokens:  usage.InputTokens,
			OutputTokens: usage.OutputTokens,
		},
	}
	if ragContext != "" {
		resp.Sources = extractSources(ragContext)
	}
	_ = 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 newConvID() string {
	var b [8]byte
	_, _ = rand.Read(b[:])
	return hex.EncodeToString(b[:])
}

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,
	})
}

func firstNonEmpty(vals ...string) string {
	for _, v := range vals {
		if v != "" {
			return v
		}
	}
	return ""
}