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"

	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
}

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
}

// 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) {
	ragContext := ""
	if r.store != nil && len(history) > 0 {
		last := history[len(history)-1]
		if last.Role == RoleUser {
			hits, err := r.store.Search(ctx, last.Content, r.topK)
			if err != nil {
				return nil, "", fmt.Errorf("rag search: %w", err)
			}
			if len(hits) > 0 {
				ragContext = r.limitRAGContext(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(r.systemPrompt, ragContext)
	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
}

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(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(r.systemPrompt, 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
			}
		}
	}
}

func formatHits(hits []portfolio.SearchResult) string {
	var b strings.Builder
	for i, h := range hits {
		fmt.Fprintf(&b, "### [%d] %s — %s\n", i+1, h.ProjectID, h.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]:]
}