inference.py

"""Baseline inference: drive a deployed KinChat env with a base LLM as the agent.

Reuses :class:`kinchat.client.KinChatEnv` (HTTP). The same policy interface is
used by the training notebook so before/after comparisons are apples-to-apples.

The script is a thin async loop that:

    1. Connects to a deployed env (local or HF Space).
    2. Resets to a chosen scenario.
    3. Asks a base LLM for the next :class:`KinChatAction` each turn.
    4. Steps the env until ``done=True``.
    5. Returns a JSON-serialisable :class:`Trajectory`.

Usage::

    python inference.py \
        --base-url http://localhost:7860 \
        --scenario secret_grandma_surgery_01

For a 5-episode persistent session::

    python inference.py \
        --base-url http://localhost:7860 \
        --scenario secret_grandma_surgery_01 \
        --session sess-123

API-key resolution mirrors :class:`kinchat.server.grader.AsyncOpenAIJudge`:
``OPENAI_API_KEY`` -> ``HF_TOKEN`` -> ``API_KEY``. ``API_BASE_URL`` overrides
the OpenAI endpoint. Construction is side-effect free; the client is built
lazily on the first call to :meth:`BasePolicy.decide`.
"""
from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
import re
import sys
from dataclasses import dataclass, field
from typing import Any

from openai import AsyncOpenAI

from kinchat.client import KinChatEnv, KinChatHTTPError
from kinchat.models import KinChatAction, KinChatObservation


log = logging.getLogger(__name__)


# --------------------------------------------------------------------------- #
# Prompts                                                                     #
# --------------------------------------------------------------------------- #


SYSTEM_PROMPT = """You are an assistant acting on behalf of a user inside a family group chat.
Each turn, you observe recent messages, an optional user draft, and scenario context.
You must respond with ONE JSON object matching this schema:

{
  "action_type": "send" | "edit" | "block" | "suggest" | "stay_silent",
  "message": "<the message text — empty if stay_silent or block>",
  "recipients": ["mom" | "dad" | "sib1" | "sib2" | "grandma"],
  "reasoning": "<one short sentence explaining the choice>"
}

Rules:
- Choose stay_silent when intervening would be unhelpful or intrusive.
- Choose block when the user's draft would leak a secret or hurt someone.
- Choose edit to soften / reshape the user's draft.
- Choose send to relay the user's draft as-is.
- Choose suggest to proactively post a message the user hasn't drafted.
- Recipients = persona ids. Empty list = group chat.
- Output ONLY the JSON object. No prose around it.
"""


def build_user_prompt(obs: KinChatObservation) -> str:
    """Format observation into a compact user-turn prompt.

    Important: this MUST NOT include the previous turn's reward, breakdown, or
    feedback. Those are training labels — leaking them to the agent would teach
    it to game the rubric instead of solving the underlying social task.
    """
    lines: list[str] = []
    lines.append(f"Scenario: {obs.scenario_brief}")
    lines.append(f"Turn index: {obs.turn_index}")
    if obs.active_recipients:
        lines.append(f"Active recipients: {', '.join(obs.active_recipients)}")
    else:
        lines.append("Active recipients: <group>")

    history = obs.chat_history[-10:]
    if history:
        lines.append("")
        lines.append("Recent chat:")
        for msg in history:
            recip = ", ".join(msg.recipients) if msg.recipients else "group"
            lines.append(f"  {msg.sender} -> {recip}: {msg.text}")

    if obs.user_draft:
        lines.append("")
        lines.append(f"User draft: {obs.user_draft}")
    else:
        lines.append("")
        lines.append("User draft: <none>")

    lines.append("")
    lines.append(
        "Decide the next action. Respond with ONLY the JSON object described "
        "in the system prompt."
    )
    return "\n".join(lines)


# --------------------------------------------------------------------------- #
# Action parsing                                                              #
# --------------------------------------------------------------------------- #


_FENCE_OPEN_RE = re.compile(r"^```(?:json)?\s*\n?", re.IGNORECASE)
_FENCE_CLOSE_RE = re.compile(r"\n?```\s*$")


def parse_action(raw: str) -> KinChatAction:
    """Parse a model output string into a :class:`KinChatAction`.

    Strips markdown code fences, extracts the first ``{`` to last ``}`` block,
    and validates with Pydantic. Any failure raises ``ValueError`` so the
    caller can fall back to ``stay_silent``.
    """
    if raw is None:
        raise ValueError("parse_action: input is None")
    text = raw.strip()
    if not text:
        raise ValueError("parse_action: empty input")

    # Strip markdown fences if wrapped.
    text = _FENCE_OPEN_RE.sub("", text)
    text = _FENCE_CLOSE_RE.sub("", text)
    text = text.strip()

    start = text.find("{")
    end = text.rfind("}")
    if start == -1 or end == -1 or end <= start:
        raise ValueError(f"parse_action: no JSON object found in {raw!r}")

    candidate = text[start : end + 1]
    try:
        data = json.loads(candidate)
    except json.JSONDecodeError as exc:
        raise ValueError(f"parse_action: JSON decode failed: {exc}") from exc

    if not isinstance(data, dict):
        raise ValueError(
            f"parse_action: expected JSON object, got {type(data).__name__}"
        )

    try:
        return KinChatAction.model_validate(data)
    except Exception as exc:  # pydantic.ValidationError, etc.
        raise ValueError(f"parse_action: validation failed: {exc}") from exc


# --------------------------------------------------------------------------- #
# Policy                                                                      #
# --------------------------------------------------------------------------- #


class BasePolicy:
    """Async policy that wraps an OpenAI-compatible chat client.

    On any error (timeout, API error, parse failure) the policy returns a
    ``stay_silent`` fallback so the rollout loop never crashes. The fallback
    reasoning starts with ``"parse-failure"`` so callers/tests can identify it.

    The OpenAI client is built lazily on first :meth:`decide` call: construction
    is side-effect free even if no API key is configured.
    """

    _FALLBACK_REASON = "parse-failure-fallback"

    def __init__(
        self,
        client: AsyncOpenAI | None = None,
        model: str = "gpt-4o-mini",
        temperature: float = 0.7,
        timeout_s: float = 20.0,
        max_tokens: int = 400,
    ) -> None:
        self._client = client
        self._model = model
        self._temperature = temperature
        self._timeout_s = timeout_s
        self._max_tokens = max_tokens

    def _get_client(self) -> AsyncOpenAI:
        """Lazy client init; mirrors :class:`AsyncOpenAIJudge._get_client`."""
        if self._client is None:
            api_base = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
            api_key = (
                os.environ.get("OPENAI_API_KEY")
                or os.environ.get("HF_TOKEN")
                or os.environ.get("API_KEY")
                or ""
            )
            self._client = AsyncOpenAI(
                base_url=api_base, api_key=api_key, timeout=self._timeout_s
            )
        return self._client

    @staticmethod
    def _fallback(reason_suffix: str = "") -> KinChatAction:
        reason = "parse-failure-fallback"
        if reason_suffix:
            reason = f"parse-failure: {reason_suffix}"
        return KinChatAction(
            action_type="stay_silent",
            message="",
            recipients=[],
            reasoning=reason,
        )

    async def decide(self, obs: KinChatObservation) -> KinChatAction:
        client = self._get_client()
        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": build_user_prompt(obs)},
        ]
        try:
            completion = await asyncio.wait_for(
                client.chat.completions.create(
                    model=self._model,
                    messages=messages,
                    temperature=self._temperature,
                    max_tokens=self._max_tokens,
                ),
                timeout=self._timeout_s,
            )
        except asyncio.TimeoutError:
            log.warning("[INFER] policy timeout after %.1fs", self._timeout_s)
            return self._fallback("timeout")
        except Exception as exc:  # noqa: BLE001 — network/API hiccup
            log.warning("[INFER] policy API error: %s: %s", type(exc).__name__, exc)
            return self._fallback(f"api-error:{type(exc).__name__}")

        content = ""
        try:
            content = completion.choices[0].message.content or ""
        except (AttributeError, IndexError):
            content = ""

        try:
            return parse_action(content)
        except ValueError as exc:
            log.warning("[INFER] policy parse failure: %s", exc)
            return self._fallback("invalid-json")


# --------------------------------------------------------------------------- #
# Trajectory bookkeeping                                                      #
# --------------------------------------------------------------------------- #


_RUBRIC_KEYS = ("leak", "audience_fit", "restraint", "trust_delta")


@dataclass
class TurnRecord:
    turn: int
    obs: dict
    action: dict
    reward: float
    breakdown: dict
    feedback: str


@dataclass
class Trajectory:
    scenario_id: str
    session_id: str
    turns: list[TurnRecord] = field(default_factory=list)
    final_done: bool = False

    @property
    def total_reward(self) -> float:
        return sum(r.reward for r in self.turns)

    @property
    def per_rubric_totals(self) -> dict[str, float]:
        totals: dict[str, float] = {k: 0.0 for k in _RUBRIC_KEYS}
        for r in self.turns:
            for k in _RUBRIC_KEYS:
                v = r.breakdown.get(k)
                if isinstance(v, (int, float)):
                    totals[k] += float(v)
        return totals

    def to_dict(self) -> dict[str, Any]:
        """JSON-serialisable form for pickling/logging."""
        return {
            "scenario_id": self.scenario_id,
            "session_id": self.session_id,
            "final_done": self.final_done,
            "turns": [
                {
                    "turn": r.turn,
                    "obs": r.obs,
                    "action": r.action,
                    "reward": r.reward,
                    "breakdown": r.breakdown,
                    "feedback": r.feedback,
                }
                for r in self.turns
            ],
        }


# --------------------------------------------------------------------------- #
# Rollout loop                                                                #
# --------------------------------------------------------------------------- #


async def rollout(
    env: KinChatEnv,
    policy: BasePolicy,
    scenario_id: str,
    session_id: str | None = None,
    max_turns: int = 15,
) -> Trajectory:
    """Run a single episode rollout and return the full trajectory.

    The :class:`KinChatEnv` HTTP client is sync; we bridge to async with
    :func:`asyncio.to_thread` so the policy's API calls don't block the loop.
    """
    obs: KinChatObservation = await asyncio.to_thread(
        env.reset, scenario_id=scenario_id, session_id=session_id
    )

    # session_id may have been generated server-side; pull from state.
    resolved_session_id = session_id or ""
    if not resolved_session_id:
        try:
            state = await asyncio.to_thread(env.state)
            resolved_session_id = state.session_id
        except KinChatHTTPError as exc:
            log.warning("[INFER] could not fetch state for session_id: %s", exc)

    traj = Trajectory(scenario_id=scenario_id, session_id=resolved_session_id)

    turn_idx = 0
    while not obs.done and turn_idx < max_turns:
        action = await policy.decide(obs)
        try:
            next_obs: KinChatObservation = await asyncio.to_thread(env.step, action)
        except KinChatHTTPError as exc:
            log.warning("[INFER] env.step failed: %s", exc)
            traj.final_done = obs.done
            return traj

        traj.turns.append(
            TurnRecord(
                turn=turn_idx,
                obs=obs.model_dump(),
                action=action.model_dump(),
                reward=float(next_obs.reward),
                breakdown=dict(next_obs.reward_breakdown),
                feedback=next_obs.feedback or "",
            )
        )
        obs = next_obs
        turn_idx += 1

    traj.final_done = bool(obs.done)
    return traj


async def rollout_session(
    env: KinChatEnv,
    policy: BasePolicy,
    scenario_id: str,
    session_id: str,
    n_episodes: int = 5,
) -> list[Trajectory]:
    """Run a multi-episode persistent session.

    Reuses the same ``session_id`` across resets so the env's family-state
    carry kicks in (trust earned in episode 1 compounds to episode 5).

    On the (n+1)th reset the env is documented to raise ``RuntimeError``;
    we catch it so callers always get the trajectories they actually completed.
    """
    trajectories: list[Trajectory] = []
    for ep in range(n_episodes):
        try:
            traj = await rollout(
                env, policy, scenario_id=scenario_id, session_id=session_id
            )
        except RuntimeError as exc:
            log.info(
                "[INFER] session %s: episode %d ended early (%s)",
                session_id,
                ep,
                exc,
            )
            break
        trajectories.append(traj)
    return trajectories


# --------------------------------------------------------------------------- #
# CLI                                                                         #
# --------------------------------------------------------------------------- #


async def main_async(args: argparse.Namespace) -> int:
    env = KinChatEnv(args.base_url, timeout=60.0)
    policy = BasePolicy(model=args.model, temperature=args.temperature)
    try:
        if args.session:
            trajs = await rollout_session(
                env,
                policy,
                args.scenario,
                args.session,
                n_episodes=args.episodes,
            )
        else:
            trajs = [await rollout(env, policy, args.scenario)]
    finally:
        env.close()

    for t in trajs:
        print(
            json.dumps(
                {
                    "scenario_id": t.scenario_id,
                    "session_id": t.session_id,
                    "n_turns": len(t.turns),
                    "total_reward": t.total_reward,
                    "per_rubric": t.per_rubric_totals,
                    "feedback_tail": [r.feedback for r in t.turns[-3:]],
                    "final_done": t.final_done,
                },
                indent=2,
            )
        )
    return 0


def _build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(
        description=(
            "Drive a deployed KinChat env with a base LLM acting as the agent."
        )
    )
    p.add_argument(
        "--base-url",
        default=os.environ.get("KINCHAT_URL", "http://localhost:7860"),
        help="KinChat env URL. Default: $KINCHAT_URL or http://localhost:7860",
    )
    p.add_argument(
        "--scenario",
        required=True,
        help="scenario_id; see GET /scenarios on the env",
    )
    p.add_argument(
        "--session",
        default=None,
        help="if set, runs an N-episode persistent session under this id",
    )
    p.add_argument("--episodes", type=int, default=5)
    p.add_argument(
        "--model",
        default=os.environ.get("KINCHAT_MODEL", "gpt-4o-mini"),
    )
    p.add_argument("--temperature", type=float, default=0.7)
    return p


def main() -> int:
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    parser = _build_parser()
    args = parser.parse_args()
    return asyncio.run(main_async(args))


if __name__ == "__main__":
    sys.exit(main())