app/agentic/speculative_executor.py

"""Speculative tool execution for the inspector agent loop (Phase 2)."""

from __future__ import annotations

import asyncio
import json
import logging
from collections.abc import Awaitable, Callable
from dataclasses import dataclass
from typing import Any

from app.config import settings

logger = logging.getLogger(__name__)

DispatchFn = Callable[..., Awaitable[tuple[str, list[Any]]]]


def speculative_execution_active() -> bool:
    return bool(settings.enable_async_pipeline and settings.enable_speculative_executor)


@dataclass(frozen=True)
class ToolPattern:
    """Stored tuple: context suffix → predicted tool with empirical probability."""

    context_suffix: tuple[str, ...]
    predicted_tool: str
    param_mapping_fn: Callable[[dict[str, Any]], dict[str, Any]]
    empirical_probability: float


@dataclass
class _InflightSpeculation:
    tool_name: str
    args: dict[str, Any]
    task: asyncio.Task[tuple[str, list[Any]]]


def _default_param_mapping(_ctx: dict[str, Any]) -> dict[str, Any]:
    return {}


def _mapping_from_last_query(ctx: dict[str, Any]) -> dict[str, Any]:
    q = str(ctx.get("last_query") or "").strip()
    if not q:
        return {"query": "RICS inspection evidence", "k": 14, "rerank_top_n": 7}
    return {"query": q, "k": 14, "rerank_top_n": 7}


def _mapping_section_plan(_ctx: dict[str, Any]) -> dict[str, Any]:
    return {
        "section_code": _ctx.get("section_code"),
        "outline": _ctx.get("outline") or "Follow extraction audit.",
    }


def _mapping_for_learned_tool(tool_name: str) -> Callable[[dict[str, Any]], dict[str, Any]]:
    """Pick a param mapper for auto-registered patterns."""
    if tool_name == "retrieve_survey_rag":
        return _mapping_from_last_query
    if tool_name == "submit_section_plan":
        return _mapping_section_plan
    return _default_param_mapping


class PatternRegistry:
    """Exact-sequence pattern store with promotion/cancel speculative dispatch."""

    def __init__(
        self,
        *,
        context_window: int | None = None,
        probability_threshold: float | None = None,
    ) -> None:
        self._window = int(context_window or settings.speculative_context_window)
        self._threshold = float(
            probability_threshold or settings.speculative_probability_threshold
        )
        self._patterns: list[ToolPattern] = []
        self._sequence_counts: dict[tuple[str, ...], dict[str, int]] = {}
        self._register_builtins()

    def _register_builtins(self) -> None:
        self.register(
            ToolPattern(
                context_suffix=("submit_extraction_audit",),
                predicted_tool="submit_section_plan",
                param_mapping_fn=_mapping_section_plan,
                empirical_probability=0.85,
            )
        )
        self.register(
            ToolPattern(
                context_suffix=("retrieve_survey_rag",),
                predicted_tool="retrieve_survey_rag",
                param_mapping_fn=_mapping_from_last_query,
                empirical_probability=0.78,
            )
        )

    def register(self, pattern: ToolPattern) -> None:
        self._patterns.append(pattern)

    def match(self, trace_tools: list[str]) -> ToolPattern | None:
        if len(trace_tools) < 1:
            return None
        suffix = tuple(trace_tools[-self._window :])
        best: ToolPattern | None = None
        for pat in self._patterns:
            n = len(pat.context_suffix)
            if len(suffix) < n or suffix[-n:] != pat.context_suffix:
                continue
            if pat.empirical_probability >= self._threshold:
                if best is None or pat.empirical_probability > best.empirical_probability:
                    best = pat
        return best

    def record_outcome(self, trace_tools: list[str], actual_tool: str, _args: dict[str, Any]) -> None:
        """Promotion mechanism: reinforce sequences that led to ``actual_tool``."""
        if len(trace_tools) < 1:
            return
        prefix = tuple(trace_tools[:-1]) if len(trace_tools) > 1 else tuple()
        key = prefix[-self._window :] if prefix else tuple()
        bucket = self._sequence_counts.setdefault(key, {})
        bucket[actual_tool] = bucket.get(actual_tool, 0) + 1
        self._maybe_register_learned_pattern(key, actual_tool)

    def _maybe_register_learned_pattern(
        self,
        context_key: tuple[str, ...],
        predicted_tool: str,
    ) -> None:
        if not context_key:
            return
        min_obs = int(getattr(settings, "speculative_learn_min_observations", 5))
        bucket = self._sequence_counts.get(context_key, {})
        count = int(bucket.get(predicted_tool, 0))
        if count < min_obs:
            return
        total = sum(bucket.values())
        if total < min_obs:
            return
        probability = count / total
        if probability < self._threshold:
            return
        suffix = context_key[-self._window :]
        for pat in self._patterns:
            if pat.context_suffix == suffix and pat.predicted_tool == predicted_tool:
                return
        self.register(
            ToolPattern(
                context_suffix=suffix,
                predicted_tool=predicted_tool,
                param_mapping_fn=_mapping_for_learned_tool(predicted_tool),
                empirical_probability=probability,
            )
        )
        logger.info(
            "speculative learned pattern suffix=%s -> %s p=%.2f",
            suffix,
            predicted_tool,
            probability,
        )

    @staticmethod
    def _args_compatible(expected: dict[str, Any], actual: dict[str, Any]) -> bool:
        for k, v in expected.items():
            if k not in actual:
                continue
            if json.dumps(actual[k], sort_keys=True, default=str) != json.dumps(
                v, sort_keys=True, default=str
            ):
                return False
        return True


class SpeculativeToolDispatcher:
    """Transparent wrapper around ``_dispatch_tool`` with speculate / promote / cancel."""

    def __init__(
        self,
        *,
        registry: PatternRegistry,
        dispatch_fn: DispatchFn,
        context: dict[str, Any],
    ) -> None:
        self._registry = registry
        self._dispatch_fn = dispatch_fn
        self._context = context
        self._trace: list[str] = []
        self._inflight: dict[str, _InflightSpeculation] = {}

    @property
    def trace_tools(self) -> list[str]:
        return list(self._trace)

    def _task_key(self, tool_name: str, args: dict[str, Any]) -> str:
        blob = json.dumps({"tool": tool_name, "args": args}, sort_keys=True, default=str)
        return blob

    def _cancel_inflight(self, *, except_key: str | None = None) -> None:
        for key, spec in list(self._inflight.items()):
            if except_key is not None and key == except_key:
                continue
            if not spec.task.done():
                spec.task.cancel()
            del self._inflight[key]

    async def maybe_start_speculation(self) -> None:
        if not speculative_execution_active():
            return
        pat = self._registry.match(self._trace)
        if pat is None:
            return
        args = pat.param_mapping_fn(dict(self._context))
        key = self._task_key(pat.predicted_tool, args)
        if key in self._inflight:
            return

        async def _run() -> tuple[str, list[Any]]:
            return await self._dispatch_fn(name=pat.predicted_tool, args=args)

        task = asyncio.create_task(_run())
        self._inflight[key] = _InflightSpeculation(
            tool_name=pat.predicted_tool,
            args=args,
            task=task,
        )
        logger.debug(
            "speculative dispatch started tool=%s suffix=%s",
            pat.predicted_tool,
            pat.context_suffix,
        )

    async def dispatch(self, *, name: str, args: dict[str, Any]) -> tuple[str, list[Any]]:
        key = self._task_key(name, args)
        spec = self._inflight.pop(key, None)

        if spec is not None and spec.tool_name == name:
            if PatternRegistry._args_compatible(spec.args, args):
                if spec.task.done() and not spec.task.cancelled():
                    try:
                        result = spec.task.result()
                        self._trace.append(name)
                        self._registry.record_outcome(self._trace, name, args)
                        self._cancel_inflight()
                        await self.maybe_start_speculation()
                        logger.debug("speculative promote tool=%s", name)
                        return result
                    except Exception:  # noqa: BLE001
                        logger.debug("speculative promote failed tool=%s", name, exc_info=True)
                elif not spec.task.done():
                    self._cancel_inflight(except_key=key)
                    result = await spec.task
                    self._trace.append(name)
                    self._registry.record_outcome(self._trace, name, args)
                    await self.maybe_start_speculation()
                    return result

        # Cancel path: LLM chose a different tool or params than speculated.
        self._cancel_inflight()
        result = await self._dispatch_fn(name=name, args=args)
        self._trace.append(name)
        self._registry.record_outcome(self._trace, name, args)
        if isinstance(args.get("query"), str):
            self._context["last_query"] = args["query"]
        await self.maybe_start_speculation()
        return result