Agentic

"""Bounded iterative retrieval as an inspectable agentic-RAG control surface.

The planner in this module is deterministic so examples remain reproducible.
Its callable interface is the seam where a supervised or reinforcement-learned
query/stop policy can be inserted.  Hard budgets remain outside that policy.
"""

import re
from collections import defaultdict
from dataclasses import dataclass
from typing import Callable, DefaultDict, Dict, List, Optional, Sequence, Tuple

from .models import Chunk, SearchResult
from .retrievers import Retriever


@dataclass(frozen=True)
class RetrievalStep:
    """One search action and the evidence-state transition it caused."""

    step: int
    query: str
    returned: int
    new_chunks: int
    accumulated_chunks: int
    stopped: bool
    reason: str


Planner = Callable[[str], Sequence[str]]
StopCondition = Callable[[str, Sequence[SearchResult]], bool]


def comparison_query_plan(query: str) -> Sequence[str]:
    """Create entity-focused follow-up searches for acronym comparisons."""

    entities = []
    for entity in re.findall(r"\b[A-Z][A-Z0-9-]{1,}\b", query):
        if entity not in entities:
            entities.append(entity)
    return tuple(f"{entity} architecture retrieval method" for entity in entities)


def named_entities_covered(query: str, results: Sequence[SearchResult]) -> bool:
    """Stop after every named acronym is represented in retrieved identity metadata."""

    required = set(re.findall(r"\b[A-Z][A-Z0-9-]{1,}\b", query))
    if not required:
        return False
    identity = set()
    for result in results:
        identity.update(re.findall(r"\b[A-Z][A-Z0-9-]{1,}\b", result.chunk.title))
        entities = result.chunk.metadata.get("entities", ())
        if isinstance(entities, str):
            entities = (entities,)
        for entity in entities:
            identity.update(re.findall(r"\b[A-Z][A-Z0-9-]{1,}\b", str(entity)))
    return required <= identity


class BudgetedIterativeRetriever:
    """Plan, retrieve, accumulate, and stop under a non-negotiable call budget.

    Results from each retrieval action are combined with reciprocal-rank fusion.
    ``last_trace`` makes the policy trajectory observable.  A learned planner
    or stop classifier may be injected, but cannot exceed ``max_steps``.
    """

    name = "budgeted-iterative"

    def __init__(
        self,
        base_retriever: Retriever,
        planner: Planner = comparison_query_plan,
        stop_when: Optional[StopCondition] = named_entities_covered,
        max_steps: int = 3,
        rrf_constant: int = 30,
        candidate_multiplier: int = 3,
    ) -> None:
        if max_steps <= 0:
            raise ValueError("max_steps must be positive")
        self.base_retriever = base_retriever
        self.planner = planner
        self.stop_when = stop_when
        self.max_steps = max_steps
        self.rrf_constant = rrf_constant
        self.candidate_multiplier = candidate_multiplier
        self.last_trace: Tuple[RetrievalStep, ...] = ()

    def search(self, query: str, k: int = 5) -> List[SearchResult]:
        if k <= 0:
            self.last_trace = ()
            return []
        plan = []
        for planned_query in (query, *self.planner(query)):
            normalized = planned_query.strip()
            if normalized and normalized not in plan:
                plan.append(normalized)
        plan = plan[: self.max_steps]

        chunks: Dict[str, Chunk] = {}
        scores: DefaultDict[str, float] = defaultdict(float)
        components: DefaultDict[str, Dict[str, float]] = defaultdict(dict)
        trace: List[RetrievalStep] = []
        candidate_k = max(k * self.candidate_multiplier, k)

        for step, planned_query in enumerate(plan, start=1):
            before = len(chunks)
            returned = self.base_retriever.search(planned_query, candidate_k)
            for result in returned:
                chunk_id = result.chunk.id
                chunks[chunk_id] = result.chunk
                contribution = 1.0 / (self.rrf_constant + result.rank)
                scores[chunk_id] += contribution
                components[chunk_id][f"step_{step}_rrf"] = contribution
            interim_ids = sorted(scores, key=lambda item: (-scores[item], item))
            interim = [
                SearchResult(
                    chunk=chunks[chunk_id],
                    score=scores[chunk_id],
                    rank=rank,
                    retriever=self.name,
                    component_scores=components[chunk_id],
                )
                for rank, chunk_id in enumerate(interim_ids[:k], start=1)
            ]
            covered = bool(self.stop_when and self.stop_when(query, interim))
            no_new_evidence = len(chunks) == before and step > 1
            is_last = step == len(plan)
            stopped = covered or no_new_evidence or is_last
            reason = (
                "coverage condition met"
                if covered
                else "no new evidence"
                if no_new_evidence
                else "plan exhausted"
                if is_last
                else "continue"
            )
            trace.append(
                RetrievalStep(
                    step=step,
                    query=planned_query,
                    returned=len(returned),
                    new_chunks=len(chunks) - before,
                    accumulated_chunks=len(chunks),
                    stopped=stopped,
                    reason=reason,
                )
            )
            if stopped:
                break

        self.last_trace = tuple(trace)
        ranked_ids = sorted(scores, key=lambda item: (-scores[item], item))[:k]
        return [
            SearchResult(
                chunk=chunks[chunk_id],
                score=scores[chunk_id],
                rank=rank,
                retriever=self.name,
                component_scores=components[chunk_id],
            )
            for rank, chunk_id in enumerate(ranked_ids, start=1)
        ]