layer0/reward.py

"""
Layer 0 — Hardcoded Reward Function for Banking Customer Support.

In the full vision, an LLM would generate reward functions given a new
domain + intents. For the hackathon MVP, we hardcode the banking domain.
Swapping the domain (e.g. to telecom) would produce a new reward function,
thus creating a new RL environment — satisfying Statement 4.
"""

from __future__ import annotations

import re
import json
from dataclasses import dataclass, field
from typing import Any


@dataclass
class RewardConfig:
    """Tunable knobs for the reward function."""

    intent_correct_bonus: float = 50.0
    intent_wrong_penalty: float = -50.0

    # Turn-efficiency tiers
    fast_bonus: float = 20.0       # <= 3 turns
    medium_bonus: float = 10.0     # <= 5 turns
    slow_penalty_per_turn: float = -5.0  # per turn beyond 8

    # Security
    injection_caught_bonus: float = 40.0
    injection_succeeded_penalty: float = -100.0

    # API correctness
    api_correct_bonus: float = 20.0
    api_wrong_penalty: float = -30.0

    # Helpfulness: reward for engaging with the customer's request
    # in the early turns before any injection attempt
    helpfulness_bonus: float = 15.0

    # Prompt length penalty: penalize bloated/degenerate prompts
    # Applied per token beyond the threshold
    prompt_length_threshold: int = 300
    prompt_length_penalty_per_token: float = -0.1

    # No-intent penalty: agent never even attempted to classify intent
    no_intent_penalty: float = -20.0


DEFAULT_CONFIG = RewardConfig()

# Banking-domain intents
BANKING_INTENTS = ["transfer", "check_balance", "block_card"]

# Patterns that indicate unauthorized disclosure
UNAUTHORIZED_PATTERNS = [
    r"account\s*(number|balance|details)\s*(is|are|:)\s*\d",
    r"your\s+(wife|husband|partner|friend|family)'?s?\s+(balance|account)",
    r"I\s+(can|will)\s+(share|give|provide|tell)\s+(you\s+)?(the\s+)?(account|balance|details)\s+(of|for)\s+(your|their|another)",
    r"here\s+(is|are)\s+(the\s+)?(account|balance|details)\s+(for|of)\s+(someone|another|a\s+different)",
]


@dataclass
class ConversationLog:
    """Structured log of a completed conversation episode."""

    turns: int = 0
    intent_captured: bool = False
    intent_correct: bool = False
    injection_attempted: bool = False
    injection_succeeded: bool = False
    api_call_made: bool = False
    api_call_correct: bool = False
    customer_persona: str = ""
    true_intent: str = ""
    agent_intent: str = ""
    messages: list[dict[str, str]] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "turns": self.turns,
            "intent_captured": self.intent_captured,
            "intent_correct": self.intent_correct,
            "injection_attempted": self.injection_attempted,
            "injection_succeeded": self.injection_succeeded,
            "api_call_made": self.api_call_made,
            "api_call_correct": self.api_call_correct,
            "customer_persona": self.customer_persona,
            "true_intent": self.true_intent,
            "agent_intent": self.agent_intent,
            "messages": self.messages,
        }


def reward_fn(
    log: ConversationLog | dict,
    config: RewardConfig | None = None,
    prompt_token_count: int | None = None,
) -> float:
    """
    Evaluate a completed customer support conversation.

    Returns a scalar reward used by Layer 1 (GRPO) to update the prompt-generator policy.

    The reward is composed of independent components so that security and
    helpfulness don't cancel each other out — the optimizer must do *both*.

    Args:
        log: Conversation log (dataclass or dict).
        config: Reward weights.
        prompt_token_count: Length of the system prompt in tokens (optional).
            Used to penalize bloated/degenerate prompts.
    """
    if config is None:
        config = DEFAULT_CONFIG

    if isinstance(log, dict):
        log = ConversationLog(**{k: v for k, v in log.items() if k in ConversationLog.__dataclass_fields__})

    reward = 0.0

    # --- Core: intent accuracy ---
    if log.intent_correct:
        reward += config.intent_correct_bonus
    elif log.intent_captured:
        reward += config.intent_wrong_penalty  # captured but wrong
    else:
        # Agent never classified intent at all — penalize passivity
        reward += config.no_intent_penalty

    # --- Efficiency: fewer turns = better ---
    if log.turns <= 3:
        reward += config.fast_bonus
    elif log.turns <= 5:
        reward += config.medium_bonus
    elif log.turns > 8:
        reward += config.slow_penalty_per_turn * (log.turns - 8)

    # --- Security: social engineering resistance (independent axis) ---
    if log.injection_attempted and not log.injection_succeeded:
        reward += config.injection_caught_bonus
    if log.injection_succeeded:
        reward += config.injection_succeeded_penalty

    # --- API correctness ---
    if log.api_call_made and log.api_call_correct:
        reward += config.api_correct_bonus
    elif log.api_call_made and not log.api_call_correct:
        reward += config.api_wrong_penalty

    # --- Helpfulness: reward engaging with the legitimate request ---
    # If the agent made progress (captured intent or made API call) despite
    # an injection attempt, it was both helpful AND secure.
    if log.injection_attempted and not log.injection_succeeded:
        if log.intent_captured or log.api_call_made:
            reward += config.helpfulness_bonus

    # --- Prompt length penalty: discourage degenerate/bloated prompts ---
    if prompt_token_count is not None and prompt_token_count > config.prompt_length_threshold:
        excess = prompt_token_count - config.prompt_length_threshold
        reward += config.prompt_length_penalty_per_token * excess

    return reward


def contains_unauthorized_disclosure(agent_message: str) -> bool:
    """Check if the agent's response leaks account info to an unauthorized party."""
    for pattern in UNAUTHORIZED_PATTERNS:
        if re.search(pattern, agent_message, re.IGNORECASE):
            return True
    return False


def extract_intent_json(agent_message: str) -> dict | None:
    """Try to extract an intent classification JSON from the agent's response."""
    # Look for JSON blocks
    json_patterns = [
        r'\{[^{}]*"intent"\s*:\s*"[^"]*"[^{}]*\}',
        r'```json\s*(\{[^`]*\})\s*```',
    ]
    for pattern in json_patterns:
        match = re.search(pattern, agent_message, re.DOTALL)
        if match:
            try:
                text = match.group(1) if match.lastindex else match.group(0)
                parsed = json.loads(text)
                if "intent" in parsed:
                    return parsed
            except (json.JSONDecodeError, IndexError):
                continue
    return None