rony-chat-bot/internal/agent/runner.go

// Package agent wraps the LLM client + RAG pipeline behind a single
// streaming call the HTTP handler can drive.
//
// We use llm.LLMClient directly (not agent.Loop) because the bot is a
// straight Q&A flow: no tools, no multi-iteration reasoning. Calling the
// underlying provider keeps the prompt under our full control so we can
// inject RAG context into the system message exactly where we want it.
//
// Auto-compaction (see CompactionConfig / Compact) folds the older portion
// of a long conversation into a single summary system message when the
// previous turn's input tokens approach the model's reported context
// window. It is opt-in via WithCompaction and short-circuits silently when
// the window is unknown or the history is too short to bother.
package agent

import (
	"context"
	"fmt"
	"iter"
	"log/slog"
	"strings"

	"github.com/VictorVargas/rony-llm-agent/pkg/llm"
	llmpersona "github.com/VictorVargas/rony-llm-agent/pkg/persona"

	"github.com/VictorVargas/rony-chat-bot/internal/i18n"
	botpersona "github.com/VictorVargas/rony-chat-bot/internal/persona"
	"github.com/VictorVargas/rony-chat-bot/internal/portfolio"
)

// Message aliases keep the HTTP handler decoupled from the upstream types.
type Message = llm.Message
type Role = llm.Role

const (
	RoleSystem    = llm.RoleSystem
	RoleUser      = llm.RoleUser
	RoleAssistant = llm.RoleAssistant
)

// Runner ties together an LLM client, the system prompt, and the RAG store.
// The persona struct is kept only for the UI greeting (its name + intro);
// the system prompt itself lives in the YAML and is passed in directly.
type Runner struct {
	client       llm.LLMClient
	persona      llmpersona.Persona
	systemPrompt string
	store        *portfolio.Store
	topK         int
	usage        *Usage

	compaction  CompactionConfig
	lastCompact CompactionStats

	// includeCatalog injects the full project list into the system prompt.
	// See portfolio.Store.Catalog for why retrieval alone isn't enough.
	includeCatalog bool

	// localizedPrompts maps a language code to a full rendition of the
	// system prompt. See WithLocalizedPrompt.
	localizedPrompts map[string]string

	// embedder enables the semantic half of retrieval. Nil → keyword only.
	embedder portfolio.Embedder
}

type Usage struct {
	InputTokens  int
	OutputTokens int
}

// CompactionConfig mirrors the YAML `compaction:` block. Zero value means
// compaction is disabled.
type CompactionConfig struct {
	Enabled             bool
	ThresholdRatio      float64 // 0–1, default 0.75 when Enabled
	KeepRecentTurns     int     // default 4 when Enabled
	SummarySystemPrompt string  // empty → built-in bilingual default
}

// CompactionStats reports the outcome of the most recent Compact call, so
// the HTTP handler can surface a "context compacted" event to the client.
type CompactionStats struct {
	Happened      bool // true when older turns were folded into a summary
	OlderTurns    int  // user turns that got summarized away
	KeptTurns     int  // user turns kept verbatim after the summary
	SummaryTokens int  // approx. token count of the summary text
	WindowTokens  int  // model's reported context window at compaction time
	UsedTokens    int  // input tokens reported by the previous turn
}

// New returns a Runner. The store may be nil (RAG disabled).
// systemPrompt is the full hand-written prompt from configs/...yaml.
func New(client llm.LLMClient, p llmpersona.Persona, systemPrompt string, store *portfolio.Store, topK int) *Runner {
	return &Runner{client: client, persona: p, systemPrompt: systemPrompt, store: store, topK: topK, usage: &Usage{}}
}

// WithCompaction enables auto-compaction with the given config. Returns
// the receiver for chaining. Pass a zero-value CompactionConfig to keep
// compaction disabled.
func (r *Runner) WithCompaction(cfg CompactionConfig) *Runner {
	r.compaction = cfg
	if cfg.Enabled {
		if r.compaction.ThresholdRatio <= 0 || r.compaction.ThresholdRatio > 1 {
			r.compaction.ThresholdRatio = 0.75
		}
		if r.compaction.KeepRecentTurns <= 0 {
			r.compaction.KeepRecentTurns = 4
		}
	}
	return r
}

// WithCatalog enables injecting the full project catalogue into the system
// prompt. Returns the receiver for chaining. No-op when RAG is disabled
// (there is no store to read the catalogue from).
func (r *Runner) WithCatalog(enabled bool) *Runner {
	r.includeCatalog = enabled
	return r
}

// WithLocalizedPrompt registers a full translation of the system prompt for
// a language code ("es"). When the visitor writes in that language, this text
// replaces the default prompt wholesale.
//
// Translating the prompt beats appending a "reply in Spanish" line to an
// English one: an instruction is a weak signal next to a thousand tokens of
// English telling the model, implicitly, what language it is working in.
// Empty prompts are ignored so an absent YAML key is a no-op.
func (r *Runner) WithLocalizedPrompt(lang, prompt string) *Runner {
	if strings.TrimSpace(prompt) == "" {
		return r
	}
	if r.localizedPrompts == nil {
		r.localizedPrompts = map[string]string{}
	}
	r.localizedPrompts[lang] = prompt
	return r
}

// WithEmbedder turns on hybrid retrieval: keyword search fused with vector
// search. Nil is accepted and leaves the runner on keyword-only search, which
// is what makes the embeddings endpoint optional rather than a hard dependency.
func (r *Runner) WithEmbedder(e portfolio.Embedder) *Runner {
	r.embedder = e
	return r
}

// promptFor returns the system prompt to use for a detected language,
// falling back to the default when no translation is registered.
func (r *Runner) promptFor(lang string) string {
	if p, ok := r.localizedPrompts[lang]; ok {
		return p
	}
	return r.systemPrompt
}

// detectLanguage returns the language code of the most recent user message,
// or "" when there is nothing to go on.
func detectLanguage(history []Message) string {
	for i := len(history) - 1; i >= 0; i-- {
		if history[i].Role == RoleUser {
			return i18n.Detect(history[i].Content)
		}
	}
	return ""
}

// languageDirective returns a one-line instruction pinning the reply language,
// or "" when the persona pins no language and the history is empty.
//
// Asking the model to "reply in the user's language" does not survive a 1B
// parameter budget: the rest of the system prompt is English, so English wins
// and Spanish questions come back in English. Detecting the language in Go and
// stating it outright is deterministic and costs one line.
//
// The directive is written *in* the target language on purpose — an
// instruction in Spanish is a much stronger prior for answering in Spanish
// than the same sentence in English.
//
// persona.language acts as the override: the default "the user's language"
// (or empty) means auto-detect, anything else pins that language verbatim.
func (r *Runner) languageDirective(lang string) string {
	if pinned := strings.TrimSpace(r.persona.Language); pinned != "" &&
		!strings.EqualFold(pinned, "the user's language") {
		return fmt.Sprintf("Write your entire reply in %s.", pinned)
	}
	if lang == "" {
		return ""
	}
	if lang == "es" {
		return "El visitante escribió en español. Responde ÍNTEGRAMENTE en español, incluido el saludo."
	}
	return "The visitor wrote in English. Write your entire reply in English."
}

// catalogBlock renders the project catalogue as a markdown list, or "" when
// the feature is off, RAG is disabled, or the index is empty. A failure to
// read it is logged and swallowed: the catalogue improves grounding but a
// chat turn should never fail because of it.
func (r *Runner) catalogBlock(ctx context.Context) string {
	if !r.includeCatalog || r.store == nil {
		return ""
	}
	entries, err := r.store.Catalog(ctx)
	if err != nil {
		slog.Warn("catalog lookup failed, continuing without it", "err", err)
		return ""
	}
	var b strings.Builder
	for _, e := range entries {
		fmt.Fprintf(&b, "- **%s** — %s\n", e.ProjectID, e.Title)
	}
	return strings.TrimRight(b.String(), "\n")
}

// LastUsage returns the token usage recorded on the most recent call.
func (r *Runner) LastUsage() Usage { return *r.usage }

// LastCompaction returns the outcome of the most recent Compact call.
// Useful for the HTTP handler to emit a "compaction" SSE event after a
// stream finishes.
func (r *Runner) LastCompaction() CompactionStats { return r.lastCompact }

// BuildMessages prepares the system prompt and turns the chat history into
// the upstream message list. The system prompt includes RAG context for the
// user's last message (if RAG is enabled).
//
// Compaction is NOT applied here — callers invoke Compact explicitly before
// BuildMessages so the same Compact call doesn't run twice per request
// (Stream calls BuildMessages internally, and handlers sometimes call it
// first to extract the RAG context for the sources event).
func (r *Runner) BuildMessages(ctx context.Context, history []Message) ([]Message, string, error) {
	// Pull any system-role notes (the compactor's "Earlier conversation
	// summary") out of the history before anything else: they are context,
	// not a conversational turn, and several chat templates reject them as
	// one. See foldSystemNotes.
	notes, history := foldSystemNotes(history)

	// The visitor's language selects which rendition of the system prompt we
	// build on, so it has to be resolved before anything else.
	lang := detectLanguage(history)
	systemPrompt := r.promptFor(lang)

	// Resolved before retrieval so the excerpt budget accounts for it.
	catalog := r.catalogBlock(ctx)

	ragContext := ""
	if r.store != nil && len(history) > 0 {
		last := history[len(history)-1]
		if last.Role == RoleUser {
			hits, err := r.store.HybridSearch(ctx, r.embedder, last.Content, r.topK)
			if err != nil {
				return nil, "", fmt.Errorf("rag search: %w", err)
			}
			if len(hits) > 0 {
				ragContext = r.limitRAGContext(systemPrompt, catalog, formatHits(hits), history)
			}
		}
	}

	// The system prompt comes from the YAML, not from the persona struct.
	// Keep the persona around only for the UI greeting.
	system := botpersona.BuildSystemPrompt(systemPrompt, catalog, ragContext)
	if len(notes) > 0 {
		system += "\n\n" + strings.Join(notes, "\n\n")
	}
	// Last line of the prompt, deliberately: it's the instruction a small
	// model is most likely to still be holding when it starts generating.
	if d := r.languageDirective(lang); d != "" {
		system += "\n\n" + d
	}
	history, err := r.fitHistory(system, history)
	if err != nil {
		return nil, "", err
	}
	msgs := make([]Message, 0, len(history)+1)
	msgs = append(msgs, Message{Role: RoleSystem, Content: system})
	msgs = append(msgs, history...)
	return msgs, ragContext, nil
}

// foldSystemNotes separates system-role messages from the conversational
// turns. Everything the runner puts in the history as a system message is
// really prompt context — today that is only the compactor's summary — so it
// belongs inside the single leading system message rather than in the turn
// list.
//
// This is not cosmetic. Gemma 3's chat template raises "Conversation roles
// must alternate user/assistant/..." on any system message after the first,
// which llama-server surfaces as HTTP 400: before this, enabling compaction
// killed the conversation outright the first time it fired. Anthropic's API
// rejects mid-conversation system turns too, so folding is the portable
// behaviour rather than a Gemma workaround.
//
// The returned slice never aliases the caller's array.
func foldSystemNotes(history []Message) (notes []string, rest []Message) {
	hasSystem := false
	for _, m := range history {
		if m.Role == RoleSystem {
			hasSystem = true
			break
		}
	}
	if !hasSystem {
		return nil, history
	}
	rest = make([]Message, 0, len(history))
	for _, m := range history {
		if m.Role == RoleSystem {
			if s := strings.TrimSpace(m.Content); s != "" {
				notes = append(notes, s)
			}
			continue
		}
		rest = append(rest, m)
	}
	return notes, rest
}

const maxContextSafetyMargin = 256

func contextBudget(window int) int {
	margin := int(float64(window) * 0.125)
	if margin > maxContextSafetyMargin {
		margin = maxContextSafetyMargin
	}
	if margin < 1 {
		margin = 1
	}
	return window - margin
}

func (r *Runner) limitRAGContext(systemPrompt, catalog, ragContext string, history []Message) string {
	if strings.TrimSpace(ragContext) == "" {
		return ""
	}
	budget := contextBudget(r.client.Capabilities().MaxContextWindow)
	if budget <= 0 {
		return ""
	}

	historyTokens := totalPromptTokens(history)
	fitted := ""
	for _, block := range strings.Split(strings.TrimSpace(ragContext), "\n\n") {
		block = strings.TrimSpace(block)
		if block == "" {
			continue
		}
		candidate := block
		if fitted != "" {
			candidate = fitted + "\n\n" + block
		}
		if estimatePromptTokens(botpersona.BuildSystemPrompt(systemPrompt, catalog, candidate))+historyTokens > budget {
			break
		}
		fitted = candidate
	}
	return fitted
}

func (r *Runner) fitHistory(system string, history []Message) ([]Message, error) {
	window := r.client.Capabilities().MaxContextWindow
	if window <= 0 {
		return history, nil
	}
	budget := contextBudget(window)
	if budget <= 0 {
		return nil, fmt.Errorf("context window is too small")
	}

	history = dropLeadingAssistantMessages(history)
	for {
		if estimatePromptTokens(system)+totalPromptTokens(history) <= budget {
			return history, nil
		}
		turnStarts := []int{}
		for i, message := range history {
			if message.Role == RoleUser {
				turnStarts = append(turnStarts, i)
			}
		}
		if len(turnStarts) < 2 {
			return nil, fmt.Errorf("prompt exceeds context window (%d tokens)", window)
		}
		first, next := turnStarts[0], turnStarts[1]
		kept := make([]Message, 0, len(history)-(next-first))
		kept = append(kept, history[:first]...)
		kept = append(kept, history[next:]...)
		history = kept
	}
}

func dropLeadingAssistantMessages(history []Message) []Message {
	firstUser := -1
	for i, message := range history {
		if message.Role == RoleUser {
			firstUser = i
			break
		}
	}
	if firstUser <= 0 {
		return history
	}
	out := make([]Message, 0, len(history))
	for _, message := range history[:firstUser] {
		if message.Role != RoleAssistant {
			out = append(out, message)
		}
	}
	return append(out, history[firstUser:]...)
}

// Compact reduces the older portion of history to a single system-role
// summary message when the previous turn's input tokens exceed the
// configured threshold. Returns the (possibly compacted) message slice —
// unchanged when compaction is disabled, history is too short, the model
// has no reported context window, or summarization fails.
//
// On summarization failure the runner falls back to dropping the oldest
// turns until the remaining slice fits a conservative budget, so a flaky
// summarize call never fails the user's request.
func (r *Runner) Compact(ctx context.Context, history []Message) ([]Message, error) {
	r.lastCompact = CompactionStats{}
	if !r.compaction.Enabled || len(history) == 0 {
		return history, nil
	}

	caps := r.client.Capabilities()
	if caps.MaxContextWindow <= 0 {
		return history, nil
	}

	used := r.usage.InputTokens
	estimated := estimatePromptTokens(r.systemPrompt) + totalPromptTokens(history)
	if estimated > used {
		used = estimated
	}
	window := caps.MaxContextWindow
	threshold := int(float64(window) * r.compaction.ThresholdRatio)
	if used < threshold {
		return history, nil
	}

	older, recent := splitByTurns(history, r.compaction.KeepRecentTurns)
	if len(older) == 0 {
		return history, nil
	}

	summary, err := r.summarize(ctx, older)
	if err != nil {
		slog.Warn("compaction: summarize failed, falling back to truncation",
			"err", err, "older_turns", userTurnCount(older))
		// Fall back: drop oldest turns until the slice is small enough to
		// fit in (threshold) tokens, using a 4-chars-per-token heuristic.
		dropped, kept := truncateToBudget(history, threshold)
		r.lastCompact = CompactionStats{
			Happened:     true,
			OlderTurns:   userTurnCount(dropped),
			KeptTurns:    userTurnCount(kept),
			WindowTokens: window,
			UsedTokens:   used,
		}
		return kept, nil
	}

	out := make([]Message, 0, 1+len(recent))
	out = append(out, Message{
		Role:    RoleSystem,
		Content: "Earlier conversation summary:\n" + summary,
	})
	out = append(out, recent...)
	r.lastCompact = CompactionStats{
		Happened:      true,
		OlderTurns:    userTurnCount(older),
		KeptTurns:     userTurnCount(recent),
		SummaryTokens: estimateTokens(summary),
		WindowTokens:  window,
		UsedTokens:    used,
	}
	return out, nil
}

// summarize calls the LLM (non-streaming) to condense older turns. We use
// a small max_tokens cap so the compaction step itself stays cheap.
func (r *Runner) summarize(ctx context.Context, older []Message) (string, error) {
	transcript := renderTranscript(older)
	sys := r.compaction.SummarySystemPrompt
	if sys == "" {
		sys = defaultSummaryPrompt
	}
	maxTokens := 512
	resp, err := r.client.Generate(ctx, llm.CompletionRequest{
		Messages: []Message{
			{Role: RoleSystem, Content: sys},
			{Role: RoleUser, Content: transcript},
		},
		MaxTokens: &maxTokens,
	})
	if err != nil {
		return "", err
	}
	summary := strings.TrimSpace(resp.Content)
	if summary == "" {
		return "", fmt.Errorf("empty summary")
	}
	return summary, nil
}

// Stream runs the model and yields each streamed chunk. The caller
// forwards chunk.Delta to the SSE stream; chunk.Usage on the last chunk
// carries token counts.
func (r *Runner) Stream(ctx context.Context, history []Message) iter.Seq2[llm.StreamChunk, error] {
	return func(yield func(llm.StreamChunk, error) bool) {
		msgs, _, err := r.BuildMessages(ctx, history)
		if err != nil {
			yield(llm.StreamChunk{}, err)
			return
		}
		req := llm.CompletionRequest{
			Messages: msgs,
			// No tools: this is a Q&A bot, not an agent.
		}
		for chunk, err := range r.client.Stream(ctx, req) {
			if chunk.Usage.TotalTokens > 0 || chunk.Usage.InputTokens > 0 || chunk.Usage.OutputTokens > 0 {
				r.usage = &Usage{
					InputTokens:  chunk.Usage.InputTokens,
					OutputTokens: chunk.Usage.OutputTokens,
				}
			}
			if !yield(chunk, err) {
				return
			}
			if err != nil {
				return
			}
		}
	}
}

// formatHits renders retrieved chunks for the prompt. Non-project documents
// are labelled as such: the CV is prime evidence for "does he know X?" but it
// is not a portfolio project, and without the label a small model happily
// announces "cv" as one of Victor's projects.
func formatHits(hits []portfolio.SearchResult) string {
	var b strings.Builder
	for i, h := range hits {
		// The marker goes after the section, never in the name slot: the
		// SSE "sources" event parses everything before " — " as the source
		// id, and a decorated name would leak into the UI chips.
		section := h.Section
		if h.Kind == portfolio.KindDoc {
			section += " (reference document, not one of Victor's projects)"
		}
		fmt.Fprintf(&b, "### [%d] %s — %s\n", i+1, h.ProjectID, section)
		b.WriteString(h.Content)
		b.WriteString("\n\n")
	}
	return b.String()
}

// defaultSummaryPrompt is used when the YAML doesn't override it. The
// instruction explicitly avoids headers / meta-commentary so the summary
// drops in cleanly as a system message and the model treats it as facts
// about the prior conversation rather than instructions from the user.
const defaultSummaryPrompt = "Summarize the following conversation between a user and the assistant concisely but completely. " +
	"Preserve decisions made, concrete facts (names, paths, IDs, preferences) and any pending task or context the assistant needs to keep helping. " +
	"Respond only with the summary, in the same language as the conversation, with no headers or extra commentary."

// splitByTurns divides messages into an older portion to summarize and a
// recent tail to keep verbatim. A turn is a user message plus everything
// up to (but not including) the next user message — so a kept turn's tool
// calls, thoughts and response always stay together. Returns older == nil
// when there are not enough turns to bother compacting.
func splitByTurns(history []Message, keepRecentTurns int) (older, recent []Message) {
	turnStarts := []int{}
	for i, m := range history {
		if m.Role == RoleUser {
			turnStarts = append(turnStarts, i)
		}
	}
	if len(turnStarts) <= keepRecentTurns {
		return nil, history
	}
	cut := turnStarts[len(turnStarts)-keepRecentTurns]
	return history[:cut], history[cut:]
}

// renderTranscript flattens chat messages into a plain User/Assistant
// transcript for summarization. Tool messages and empty assistant
// placeholders are skipped — neither contributes facts the model needs
// to remember.
func renderTranscript(messages []Message) string {
	var b strings.Builder
	for _, m := range messages {
		switch m.Role {
		case RoleUser:
			b.WriteString("User: ")
			b.WriteString(m.Content)
			b.WriteString("\n")
		case RoleAssistant:
			if m.Content == "" {
				continue
			}
			b.WriteString("Assistant: ")
			b.WriteString(m.Content)
			b.WriteString("\n")
		}
	}
	return b.String()
}

// userTurnCount is the number of user-role messages in a slice — a
// convenient proxy for "turns" in stats reporting (each turn starts
// with a user message).
func userTurnCount(messages []Message) int {
	n := 0
	for _, m := range messages {
		if m.Role == RoleUser {
			n++
		}
	}
	return n
}

// estimateTokens gives a rough token count for s. ~4 chars per token is
// the usual rule of thumb for English; close enough for Spanish / mixed
// text that this only drives a trigger threshold.
func estimateTokens(s string) int {
	if s == "" {
		return 0
	}
	return len(s)/4 + 1
}

func estimatePromptTokens(s string) int {
	if s == "" {
		return 0
	}
	return len(s)/3 + 1
}

func totalPromptTokens(messages []Message) int {
	total := 0
	for _, m := range messages {
		total += estimatePromptTokens(m.Content)
	}
	return total
}

// truncateToBudget is the fallback when summarization fails. It drops
// the oldest user-turns one at a time until the remaining slice fits in
// `budget` tokens (conservative heuristic). At least the last user turn is
// always preserved, so we never return an empty slice.
//
// Implementation note: as `cut` increases, history[turnStarts[cut]:]
// shrinks, so `totalPromptTokens(...)` is monotonically non-increasing in
// `cut`. The loop is therefore guaranteed to terminate by either
// fitting the budget or hitting the "keep at least the last turn" floor.
func truncateToBudget(history []Message, budget int) (dropped, kept []Message) {
	if len(history) == 0 {
		return nil, nil
	}
	turnStarts := []int{}
	for i, m := range history {
		if m.Role == RoleUser {
			turnStarts = append(turnStarts, i)
		}
	}
	if len(turnStarts) == 0 {
		return nil, history
	}
	// cut = number of leading turns to drop. 0 = keep everything; we
	// grow it until the remaining slice fits in budget (or we hit the
	// floor of one turn kept).
	cut := 0
	for cut < len(turnStarts)-1 {
		start := turnStarts[cut+1]
		if totalPromptTokens(history[start:]) <= budget {
			break
		}
		cut++
	}
	if cut == 0 {
		return nil, history
	}
	return history[:turnStarts[cut]], history[turnStarts[cut]:]
}