rewards.py

"""
Genesis-2.0 RLHF — Reward Functions Module

Rule-based reward functions for scoring tool-use trajectories.
Used by both the DPO preference pair generator and the GRPO trainer.
"""

import json
import re
from typing import Any, Optional


def extract_tool_calls(text: str) -> list[dict[str, Any]]:
    """Extract <tool_call> JSON blocks from assistant messages."""
    tool_calls = []
    pattern = re.compile(r'<tool_call>\s*(\{.*?\})\s*</tool_call>', re.DOTALL)
    for match in pattern.finditer(text):
        try:
            tc = json.loads(match.group(1))
            if isinstance(tc, dict) and "name" in tc:
                tool_calls.append(tc)
        except json.JSONDecodeError:
            continue
    return tool_calls


def extract_assistant_messages(text: str) -> list[str]:
    """Split a conversation into individual assistant response blocks."""
    parts = text.split("<|im_start|>assistant\n")
    if len(parts) <= 1:
        return []
    messages = []
    for p in parts[1:]:
        end_idx = p.find("<|im_end|>")
        if end_idx >= 0:
            messages.append(p[:end_idx].strip())
    return messages


def reward_tool_call_validity(trajectory: str) -> float:
    """
    +1.0 if all <tool_call> blocks contain valid JSON with valid tool names.
    0.0 if ANY tool call has invalid JSON or missing 'name' field.
    """
    tool_calls = extract_tool_calls(trajectory)
    if not tool_calls:
        return 0.0
    for tc in tool_calls:
        if not isinstance(tc.get("name"), str) or not tc["name"].strip():
            return 0.0
        if "arguments" not in tc:
            return 0.5  # partial: has name but no args
    return 1.0


def reward_task_completion(trajectory: str, expected_keywords: Optional[list[str]] = None) -> float:
    """
    +1.0 if the final assistant message provides a substantive answer.
    0.0 if the trajectory ends without a meaningful response.
    -1.0 if it ends with an error or hallucination.
    """
    msgs = extract_assistant_messages(trajectory)
    if not msgs:
        return -1.0

    last = msgs[-1]

    # If the last message has tool calls but no text answer, task isn't complete
    text_after_tools = re.sub(r'<tool_call>.*?</tool_call>', '', last, flags=re.DOTALL).strip()
    text_after_tools = re.sub(r'<tool_result>.*?</tool_result>', '', text_after_tools, flags=re.DOTALL).strip()

    tool_calls = extract_tool_calls(last)
    if not text_after_tools and tool_calls:
        return 0.0  # only tool calls, no answer

    if text_after_tools and len(text_after_tools) > 20:
        # Has substantive text — likely gave an answer
        if expected_keywords:
            hits = sum(1 for kw in expected_keywords if kw.lower() in text_after_tools.lower())
            return min(1.0, hits / max(len(expected_keywords), 1) * 1.5)
        return 1.0

    if text_after_tools and len(text_after_tools) > 5:
        return 0.5  # short answer, partial completion

    return 0.0


def reward_tool_result_usage(trajectory: str) -> float:
    """
    +1.0 if tool calls are followed by reasoning that references them.
    0.0 if tool calls are made but no subsequent reasoning.
    """
    segments = trajectory.split("<|im_start|>assistant\n")
    tool_call_segments = 0
    reasoning_after_call = 0

    for seg in segments[1:]:
        has_tool_call = bool(re.search(r'<tool_call>', seg))
        if has_tool_call:
            tool_call_segments += 1
            # Remove tool_call blocks and check remaining text
            text_after = re.sub(r'<tool_call>.*?</tool_call>', '', seg, flags=re.DOTALL)
            text_after = re.sub(r'<think>.*?</think>', '', text_after, flags=re.DOTALL)
            text_after = text_after.strip()
            if len(text_after) > 50:
                reasoning_after_call += 1

    if tool_call_segments == 0:
        return 0.5  # no tools used — neutral
    if reasoning_after_call >= tool_call_segments * 0.5:
        return 1.0
    if reasoning_after_call > 0:
        return 0.5
    return 0.0


def reward_trajectory_efficiency(trajectory: str) -> float:
    """
    Score based on how efficiently tools are used.
    Ideal: 1-3 tool calls for a typical query.
    Penalty for excessive tool calls (>5) or zero tool calls when tools would help.
    """
    tool_calls = extract_tool_calls(trajectory)
    n = len(tool_calls)

    if n == 0:
        return 0.3  # no tools used — might be fine for simple queries
    elif n <= 3:
        return 1.0  # efficient
    elif n <= 5:
        return 0.7  # acceptable
    elif n <= 8:
        return 0.4  # somewhat excessive
    else:
        return 0.0  # too many tool calls


def reward_format_validity(trajectory: str) -> float:
    """
    +1.0 if tool calls use the correct <tool_call> JSON format.
    0.0 if there are malformed tool calls.
    """
    # Check for malformed tool_call markers
    malformed = re.findall(r'<tool_call>[^<]*?$', trajectory)
    if malformed:
        return 0.0

    # Check for unclosed tool_call tags
    open_tags = len(re.findall(r'<tool_call>', trajectory))
    close_tags = len(re.findall(r'</tool_call>', trajectory))
    if open_tags != close_tags:
        return 0.0

    # All parsed tool calls should have valid JSON
    tool_calls = extract_tool_calls(trajectory)
    expected = len(re.findall(r'<tool_call>', trajectory))
    if expected > 0 and len(tool_calls) < expected:
        return 0.5  # some valid, some not

    return 1.0


def reward_recovery_from_error(trajectory: str) -> float:
    """
    +1.0 if the model retries after a failed tool call (multi-turn recovery).
    Detects consecutive assistant segments making tool calls.
    0.0 otherwise.
    """
    segments = trajectory.split("<|im_start|>assistant\n")[1:]
    consecutive_tool_calls = 0
    for i, seg in enumerate(segments):
        has_tc = bool(re.search(r'<tool_call>', seg))
        if has_tc and i > 0:
            prev_has_tc = bool(re.search(r'<tool_call>', segments[i-1]))
            if prev_has_tc:
                consecutive_tool_calls += 1

    if consecutive_tool_calls >= 1:
        return 1.0  # did a multi-step tool sequence
    return 0.0


def combined_reward(trajectory: str, expected_keywords: Optional[list[str]] = None) -> float:
    """
    Weighted combination of all reward signals.
    Used by both DPO pair generation and GRPO training.
    """
    weights = {
        "validity": 0.25,
        "completion": 0.40,
        "usage": 0.15,
        "efficiency": 0.10,
        "format": 0.05,
        "recovery": 0.05,
    }

    scores = {
        "validity": reward_tool_call_validity(trajectory),
        "completion": reward_task_completion(trajectory, expected_keywords),
        "usage": reward_tool_result_usage(trajectory),
        "efficiency": reward_trajectory_efficiency(trajectory),
        "format": reward_format_validity(trajectory),
        "recovery": reward_recovery_from_error(trajectory),
    }

    total = sum(scores[k] * weights[k] for k in weights)
    return max(-1.0, min(1.0, total))  # clamp to [-1, 1]


def reward_debug(trajectory: str) -> dict:
    """Return breakdown of all reward components for debugging."""
    return {
        "validity": reward_tool_call_validity(trajectory),
        "completion": reward_task_completion(trajectory),
        "usage": reward_tool_result_usage(trajectory),
        "efficiency": reward_trajectory_efficiency(trajectory),
        "format": reward_format_validity(trajectory),
        "recovery": reward_recovery_from_error(trajectory),
        "combined": combined_reward(trajectory),
    }