src/ragforge/context_budget.py

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Iterable

from .schemas import PipelineConfig, QueryPlan, SearchHit


@dataclass(slots=True)
class ContextBudgetDecision:
    hits: list[SearchHit]
    used: bool
    reason: str
    policy: str
    target_chunks: int
    corpus_scale: str
    retrieval_confidence: float
    score_gap: float
    chunks_before: int
    chunks_after: int
    sources_before: int
    sources_after: int
    chars_before: int
    chars_after: int
    tokens_est_before: int
    tokens_est_after: int

    @property
    def reduction_ratio(self) -> float:
        if self.chars_before <= 0:
            return 0.0
        return max(0.0, min(1.0, 1.0 - (self.chars_after / self.chars_before)))

    def trace_fields(self) -> dict[str, object]:
        return {
            "context_pruning_used": self.used,
            "context_pruning_reason": self.reason,
            "context_budget_policy": self.policy,
            "context_budget_target_chunks": self.target_chunks,
            "corpus_scale": self.corpus_scale,
            "retrieval_confidence": round(self.retrieval_confidence, 3),
            "retrieval_score_gap": round(self.score_gap, 3),
            "context_chunks_before": self.chunks_before,
            "context_chunks_after": self.chunks_after,
            "context_sources_before": self.sources_before,
            "context_sources_after": self.sources_after,
            "context_chars_before": self.chars_before,
            "context_chars_after": self.chars_after,
            "context_tokens_est_before": self.tokens_est_before,
            "context_tokens_est_after": self.tokens_est_after,
            "context_reduction_pct": round(self.reduction_ratio * 100.0, 1),
        }


def estimate_tokens_from_chars(chars: int) -> int:
    """Cheap deterministic token estimate for observability, not billing."""
    return int(math.ceil(max(0, int(chars)) / 4.0))


def corpus_scale_label(chunks: int, sources: int) -> str:
    chunks = max(0, int(chunks))
    sources = max(0, int(sources))
    if chunks >= 3000 or sources >= 80:
        return "very_large"
    if chunks >= 1000 or sources >= 30:
        return "large"
    if chunks >= 250 or sources >= 10:
        return "medium"
    return "small"


def adaptive_retrieval_top_k(
    config: PipelineConfig,
    plan: QueryPlan,
    *,
    corpus_chunks: int,
    corpus_sources: int,
) -> int:
    """Choose initial retrieval breadth from corpus scale and task breadth.

    ``top_k`` remains the small-corpus baseline configured by the user. Larger
    corpora can retrieve a wider candidate set before the focused context
    budget trims the generation context. Broad tasks retain wider evidence.
    """
    baseline = max(2, min(12, int(config.top_k)))
    if not getattr(config, "use_adaptive_top_k", True):
        return baseline

    scale = corpus_scale_label(corpus_chunks, corpus_sources)
    scale_floor = {"small": baseline, "medium": max(baseline, 8), "large": max(baseline, 10), "very_large": 12}[scale]

    if plan.task_type in {"overview", "insight_synthesis", "cross_document_synthesis", "comparison"}:
        # Broad tasks need source breadth. Keep the absolute cap bounded by the
        # schema and avoid forcing more chunks than sources can usefully supply.
        broad_floor = min(12, max(scale_floor, min(max(6, corpus_sources), 12)))
        return broad_floor
    return min(12, scale_floor)


def _context_chars(hits: Iterable[SearchHit]) -> int:
    # Include a small deterministic metadata allowance per chunk because the
    # generation context contains source/page labels in addition to raw text.
    return sum(len(hit.chunk.text or "") + len(hit.chunk.source or "") + 24 for hit in hits)


def _source_count(hits: Iterable[SearchHit]) -> int:
    return len({hit.chunk.source for hit in hits})


def _hit_signal(hit: SearchHit) -> float:
    dense = max(0.0, min(1.0, float(hit.dense_score or 0.0)))
    sparse = max(0.0, min(1.0, float(hit.sparse_score or 0.0))) * 0.9
    fused = max(0.0, min(1.0, float(hit.score or 0.0)))
    return max(dense, sparse, fused)


def _retrieval_confidence(hits: list[SearchHit]) -> tuple[float, float]:
    if not hits:
        return 0.0, 0.0
    signals = [_hit_signal(hit) for hit in hits[:4]]
    top = signals[0]
    second = signals[1] if len(signals) > 1 else 0.0
    gap = max(0.0, top - second)
    # Agreement between dense and sparse on the first result adds confidence,
    # while a small top-two gap indicates ambiguity even with a strong top hit.
    first = hits[0]
    method_agreement = float(
        float(first.dense_score or 0.0) >= 0.20 and float(first.sparse_score or 0.0) >= 0.05
    )
    confidence = 0.65 * top + 0.25 * min(1.0, gap / 0.20) + 0.10 * method_agreement
    return max(0.0, min(1.0, confidence)), gap


def adaptive_context_budget(
    hits: list[SearchHit],
    plan: QueryPlan,
    config: PipelineConfig,
    *,
    corpus_chunks: int = 0,
    corpus_sources: int = 0,
) -> ContextBudgetDecision:
    """Adaptively shrink context for focused local lookups only.

    v1.8 proved that a three-chunk safety floor could cut the demo context by
    roughly half without harming source recall. v1.9 generalizes that policy:
    the target is chosen from retrieval confidence, score separation and corpus
    scale. Broad/synthesis tasks are never pruned by this function.
    """
    before = list(hits or [])
    chunks_before = len(before)
    chars_before = _context_chars(before)
    sources_before = _source_count(before)
    scale = corpus_scale_label(corpus_chunks or chunks_before, corpus_sources or sources_before)
    retrieval_confidence, score_gap = _retrieval_confidence(before)

    def decision(
        after: list[SearchHit],
        used: bool,
        reason: str,
        *,
        target: int,
        policy: str = "adaptive_focused_budget",
    ) -> ContextBudgetDecision:
        chars_after = _context_chars(after)
        return ContextBudgetDecision(
            hits=after,
            used=used,
            reason=reason,
            policy=policy,
            target_chunks=max(0, int(target)),
            corpus_scale=scale,
            retrieval_confidence=retrieval_confidence,
            score_gap=score_gap,
            chunks_before=chunks_before,
            chunks_after=len(after),
            sources_before=sources_before,
            sources_after=_source_count(after),
            chars_before=chars_before,
            chars_after=chars_after,
            tokens_est_before=estimate_tokens_from_chars(chars_before),
            tokens_est_after=estimate_tokens_from_chars(chars_after),
        )

    if not before:
        return decision(before, False, "no_document_evidence", target=0)
    if not getattr(config, "use_context_pruning", True):
        return decision(before, False, "disabled_by_user", target=chunks_before, policy="disabled")
    if plan.route not in {"documents", "hybrid"}:
        return decision(before, False, "non_local_route", target=chunks_before)
    if plan.task_type not in {"fact_lookup", "followup"}:
        return decision(before, False, "broad_or_multi_source_task", target=chunks_before)
    if plan.retrieval_strategy not in {"semantic", "hierarchical"}:
        return decision(before, False, "strategy_requires_breadth", target=chunks_before)

    # Base safety grows with corpus scale because a larger corpus increases the
    # chance that a seemingly focused question needs a runner-up source.
    base = {"small": 3, "medium": 3, "large": 4, "very_large": 5}[scale]
    target = base

    # A clearly separated high-confidence top result can safely use two chunks
    # on small/medium corpora. Avoid top-1 pruning so a second supporting chunk
    # remains available for citation and accidental under-classification.
    if scale in {"small", "medium"} and retrieval_confidence >= 0.78 and score_gap >= 0.15:
        target = 2
    # Ambiguous retrieval needs extra evidence, especially when the top results
    # already span multiple sources.
    top_sources = len({hit.chunk.source for hit in before[:3]})
    if retrieval_confidence < 0.55 or score_gap < 0.035 or top_sources >= 3:
        target = min(5, max(target, base + 1))
    if plan.task_type == "followup":
        target = min(5, max(target, 3))

    # The initial adaptive retrieval depth may exceed config.top_k on a large
    # corpus, so cap against the actual candidate count rather than config.top_k.
    target = max(2, min(target, chunks_before))
    if chunks_before <= target:
        return decision(before, False, "already_within_adaptive_budget", target=target)

    after = before[:target]
    reason = f"adaptive_{scale}_top{target}"
    return decision(after, True, reason, target=target)


def focused_context_budget(
    hits: list[SearchHit],
    plan: QueryPlan,
    config: PipelineConfig,
    *,
    safety_floor: int = 3,
) -> ContextBudgetDecision:
    """Backward-compatible v1.8 fixed safety-floor policy.

    The runtime uses :func:`adaptive_context_budget` in v1.9. Keeping this
    helper preserves the documented v1.8 ablation semantics and makes release
    deltas auditable instead of silently changing an old experiment.
    """
    before = list(hits or [])
    chunks_before = len(before)
    chars_before = _context_chars(before)
    sources_before = _source_count(before)
    confidence, gap = _retrieval_confidence(before)

    def decision(after: list[SearchHit], used: bool, reason: str) -> ContextBudgetDecision:
        chars_after = _context_chars(after)
        return ContextBudgetDecision(
            hits=after,
            used=used,
            reason=reason,
            policy="v1.8_fixed_safety_floor",
            target_chunks=len(after),
            corpus_scale=corpus_scale_label(chunks_before, sources_before),
            retrieval_confidence=confidence,
            score_gap=gap,
            chunks_before=chunks_before,
            chunks_after=len(after),
            sources_before=sources_before,
            sources_after=_source_count(after),
            chars_before=chars_before,
            chars_after=chars_after,
            tokens_est_before=estimate_tokens_from_chars(chars_before),
            tokens_est_after=estimate_tokens_from_chars(chars_after),
        )

    if not before:
        return decision(before, False, "no_document_evidence")
    if not getattr(config, "use_context_pruning", True):
        return decision(before, False, "disabled_by_user")
    if plan.route not in {"documents", "hybrid"}:
        return decision(before, False, "non_local_route")
    if plan.task_type not in {"fact_lookup", "followup"}:
        return decision(before, False, "broad_or_multi_source_task")
    if plan.retrieval_strategy not in {"semantic", "hierarchical"}:
        return decision(before, False, "strategy_requires_breadth")
    target = max(2, min(int(safety_floor), int(config.top_k)))
    if chunks_before <= target:
        return decision(before, False, "already_within_budget")
    return decision(before[:target], True, "focused_lookup_top3_safety_floor")