evaluation/model_policy.py

"""Model-backed evaluation policy for SQLEnv.

This is the missing piece the success gate ("beat 28-32% on the N=50 Spider
subset") depends on: a `Policy` that drives the TRAINED LLM through the
environment, so `evaluate(env, ModelPolicy(...), n_episodes=...)` produces a
real `success_rate`. The same agentic-rollout primitive is what the Gradio Space
needs to serve a question, so build/validate it once here.

IMPORTANT — this reconstructs the multi-turn tool-use transcript INDEPENDENTLY of
TRL. During training, TRL's `environment_factory` owns the chat rendering and
tool dispatch; that machinery is internal and not reusable at eval time. So this
policy maintains its own message history and parses the model's emitted tool call
itself. That introduces a small train/serve-skew risk in prompt formatting —
**validate it** by running the smoke eval and confirming the parse rate is
~95-100% (matching training) before trusting any accuracy number.

Heavy deps (torch/transformers) are imported lazily so this module lint-imports
without the training extras. `ModelPolicy` satisfies the `Policy` protocol in
``evaluation.policies`` (``select_action(observation) -> SQLAction``).
"""

from __future__ import annotations

import logging
from typing import Any

try:
    from sql_env.models import SQLAction, SQLObservation
except ImportError:  # pragma: no cover - flat-layout / Docker fallback
    from models import SQLAction, SQLObservation  # type: ignore[no-redef]

try:
    from sql_env.server.tooling import (
        get_system_prompt,
        get_tool_definitions,
        parse_action,
    )
except ImportError:  # pragma: no cover - flat-layout / Docker fallback
    from server.tooling import (  # type: ignore[no-redef]
        get_system_prompt,
        get_tool_definitions,
        parse_action,
    )

logger = logging.getLogger(__name__)


class ModelPolicy:
    """Drive a trained causal-LM through SQLEnv as an agentic tool-use policy.

    Parameters
    ----------
    model, tokenizer
        A loaded causal LM + tokenizer (base+adapter already merged/attached).
    enable_thinking
        Must match how the model was trained (default False -> /no_think).
    temperature, top_p, top_k
        Inference sampling. Defaults are Qwen3 non-thinking inference settings
        (T=0.7, top_p=0.8, top_k=20) — NOT the RL training temperature (1.0).
    max_new_tokens
        Per-turn generation cap.
    """

    def __init__(
        self,
        model: Any,
        tokenizer: Any,
        *,
        enable_thinking: bool = False,
        temperature: float = 0.7,
        top_p: float = 0.8,
        top_k: int = 20,
        max_new_tokens: int = 512,
    ) -> None:
        self._model = model
        self._tokenizer = tokenizer
        self._enable_thinking = enable_thinking
        self._temperature = temperature
        self._top_p = top_p
        self._top_k = top_k
        self._max_new_tokens = max_new_tokens

        self._system_prompt = get_system_prompt(enable_thinking=enable_thinking)
        self._tools = get_tool_definitions()
        # Per-episode conversation state, (re)initialised when step_count == 0.
        self._messages: list[dict[str, str]] = []
        self._started = False

    def select_action(self, observation: SQLObservation) -> SQLAction:
        """Choose one tool call for the current observation."""
        # A fresh episode: evaluate() calls env.reset() (step_count == 0) before
        # the first select_action. Detect that and reset our transcript.
        if observation.step_count == 0 or not self._started:
            self._messages = [
                {"role": "system", "content": self._system_prompt},
                {
                    "role": "user",
                    "content": (
                        f"{observation.question}\n\n"
                        f"Available tables:\n{observation.schema_info}"
                    ),
                },
            ]
            self._started = True
        else:
            # The observation carries the result/error of OUR previous tool call.
            tool_response = observation.result or (
                f"Error: {observation.error}" if observation.error else "No output."
            )
            self._messages.append({"role": "tool", "content": tool_response})

        completion = self._generate(self._messages)
        self._messages.append({"role": "assistant", "content": completion})

        action = self._parse_action(completion)
        if action is None:
            # Unparseable output near the budget: submit a best-effort answer so
            # the episode terminates instead of burning steps on malformed calls.
            if observation.budget_remaining <= 1:
                return SQLAction(action_type="ANSWER", argument="[]")
            # Otherwise emit a harmless exploration step; the parse failure is
            # itself the signal we watch (parse rate).
            return SQLAction(action_type="QUERY", argument="SELECT 1")
        return action

    def _generate(self, messages: list[dict[str, str]]) -> str:
        import torch  # noqa: PLC0415

        inputs = self._tokenizer.apply_chat_template(
            messages,
            tools=self._tools,
            add_generation_prompt=True,
            return_tensors="pt",
            tokenize=True,
            # transformers v5 flipped the default to return a BatchEncoding
            # dict; this loop (and the trained format) expects the bare
            # input_ids tensor — pin the pre-v5 behavior explicitly.
            return_dict=False,
        )
        inputs = inputs.to(self._model.device)
        # Pass an explicit all-ones attention mask. We generate one unpadded
        # sequence at a time, so this is behaviorally a no-op today — but because
        # Qwen3's pad_token == eos_token, transformers cannot auto-infer the mask
        # and warns ("attention mask is not set ... you may observe unexpected
        # behavior"). Passing it silences the warning AND prevents a real
        # correctness bug if we ever batch eval generation (padded sequences need
        # a mask, or the model attends to pad tokens).
        attention_mask = torch.ones_like(inputs)
        with torch.no_grad():
            out = self._model.generate(
                inputs,
                attention_mask=attention_mask,
                max_new_tokens=self._max_new_tokens,
                do_sample=self._temperature > 0,
                temperature=self._temperature,
                top_p=self._top_p,
                top_k=self._top_k,
                pad_token_id=self._tokenizer.eos_token_id,
            )
        generated = out[0][inputs.shape[-1] :]
        return self._tokenizer.decode(generated, skip_special_tokens=False)

    @staticmethod
    def _parse_action(completion: str) -> SQLAction | None:
        """Extract the first <tool_call> JSON and map it to an SQLAction.

        Thin back-compat alias delegating to ``sql_env.server.tooling.parse_action``.
        """
        return parse_action(completion)