proteus/game/agents/base.py

"""Slim agent interface for the PROTEUS arena.

An agent observes the grid and produces an action (plus the reasoning behind
it). It may optionally answer a side-channel probe. There is no forfeit,
stake, or risk machinery — the arena measures motive-reading, not
self-preservation trade-offs.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass


@dataclass(frozen=True)
class ActResult:
    """An agent's decision for one turn.

    Immutable: a turn's decision record must not be mutated after the agent
    returns it (it is written verbatim into the session trace).

    Attributes:
        action: The chosen action (one of the available actions).
        reasoning: The agent's stated/extracted rationale (CoT / thinking).
            Empty string if none was produced.
        raw_text: The full unprocessed model output.
        input_tokens: Token usage for the act call — prompt/input side.
        output_tokens: Token usage for the act call — completion/output side.
        thinking_tokens: Reasoning-token count (provider-reported or inline
            ``<think>`` whitespace-split count, whichever is available).
    """

    action: str
    reasoning: str
    raw_text: str
    input_tokens: int = 0
    output_tokens: int = 0
    thinking_tokens: int = 0


@dataclass(frozen=True)
class ProbeResult:
    """An agent's answer to a side-channel probe question for one turn.

    Immutable: written verbatim into the session trace as a 1st-class
    measurement target (the probe is scored offline later).

    Attributes:
        answer: The think-stripped probe answer.
        reasoning: The probe's stated/extracted rationale (CoT / thinking).
        raw_text: The full unprocessed model output for the probe.
        input_tokens: Prompt token usage for the probe call.
        output_tokens: Completion token usage for the probe call.
        thinking_tokens: Reasoning-token count for the probe call.
    """

    answer: str
    reasoning: str = ""
    raw_text: str = ""
    input_tokens: int = 0
    output_tokens: int = 0
    thinking_tokens: int = 0


class Agent(ABC):
    """Abstract LLM agent that plays a motive_grid scenario."""

    @property
    @abstractmethod
    def name(self) -> str:
        """Short identifier for this agent variant (e.g. ``"vanilla"``)."""

    @abstractmethod
    def act(
        self,
        observation: str,
        available_actions: list[str],
        system_prompt: str,
    ) -> ActResult:
        """Choose an action for the current turn.

        Args:
            observation: Text rendering of the current world (+ Cut history
                on the first turn).
            available_actions: Valid action strings to choose from.
            system_prompt: Rules + handover framing for the session.

        Returns:
            An :class:`ActResult`.
        """

    @abstractmethod
    def probe(
        self,
        observation: str,
        question: str,
        system_prompt: str,
    ) -> ProbeResult:
        """Answer a side-channel comprehension probe (does not change state).

        Args:
            observation: Text rendering of the current world.
            question: The probe question.
            system_prompt: Rules + handover framing for the session.

        Returns:
            A :class:`ProbeResult` containing the think-stripped answer,
            reasoning, raw model output, and token accounting.
        """

    @abstractmethod
    def reset(self) -> None:
        """Clear any per-session internal state."""