rnow 0.2.4__py3-none-any.whl

rnow/cli/token_count.py

@@ -0,0 +1,295 @@
+"""
+Token counting utilities for accurate context window validation.
+
+Supports:
+- gpt-oss: Uses openai-harmony render_conversation_for_completion (canonical API)
+- Qwen/other HF models: Uses HuggingFace tokenizers with chat template approximation
+- Fallback: Conservative char-based estimate
+"""
+
+from __future__ import annotations
+
+import heapq
+import json
+from pathlib import Path
+from typing import Any
+
+# Tokenizer cache to avoid reloading
+_tokenizer_cache: dict[str, Any] = {}
+
+
+def get_tokenizer_for_model(model_path: str) -> tuple[str, Any] | None:
+    """
+    Get the appropriate tokenizer for a model.
+
+    Returns:
+        Tuple of (tokenizer_type, tokenizer) or None if unavailable.
+        tokenizer_type is "harmony" for gpt-oss, "hf" for HuggingFace models.
+    """
+    if model_path in _tokenizer_cache:
+        return _tokenizer_cache[model_path]
+
+    # gpt-oss models use openai-harmony
+    if "gpt-oss" in model_path.lower():
+        try:
+            from openai_harmony import HarmonyEncodingName, load_harmony_encoding
+
+            tokenizer = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
+            _tokenizer_cache[model_path] = ("harmony", tokenizer)
+            return _tokenizer_cache[model_path]
+        except ImportError:
+            pass
+    else:
+        # Try HuggingFace tokenizers (requires tokenizer.json in repo)
+        try:
+            from tokenizers import Tokenizer
+
+            tokenizer = Tokenizer.from_pretrained(model_path)
+            _tokenizer_cache[model_path] = ("hf", tokenizer)
+            return _tokenizer_cache[model_path]
+        except Exception:
+            # tokenizer.json not available or repo requires auth
+            pass
+
+    _tokenizer_cache[model_path] = None
+    return None
+
+
+def count_tokens_for_gpt_oss(
+    messages: list[dict], tools: list[dict], instructions: str = ""
+) -> int:
+    """
+    Count tokens for gpt-oss using the canonical Harmony API.
+
+    This uses render_conversation_for_completion which is the correct way
+    to count tokens exactly as gpt-oss runtime sees them.
+
+    System messages from the dataset are merged into DeveloperContent.instructions,
+    since Harmony's SystemContent is for model metadata, not user system prompts.
+    """
+    try:
+        from openai_harmony import (
+            Conversation,
+            DeveloperContent,
+            HarmonyEncodingName,
+            Message,
+            Role,
+            SystemContent,
+            ToolDescription,
+            load_harmony_encoding,
+        )
+    except ImportError:
+        # Conservative fallback: count JSON chars as tokens
+        return sum(len(json.dumps(m, ensure_ascii=False)) for m in messages) + sum(
+            len(json.dumps(t, ensure_ascii=False)) for t in tools
+        )
+
+    encoding = load_harmony_encoding(HarmonyEncodingName.HARMONY_GPT_OSS)
+
+    # Extract system messages from dataset - these become developer instructions
+    # (Harmony's SystemContent is model metadata, not user system prompts)
+    sys_texts = [
+        m.get("content", "")
+        for m in messages
+        if m.get("role") == "system" and isinstance(m.get("content"), str)
+    ]
+    sys_text = "\n\n".join(t.strip() for t in sys_texts if t.strip())
+
+    # Combine with any additional instructions
+    combined_instructions = "\n\n".join(s for s in [instructions.strip(), sys_text] if s).strip()
+
+    # Build tool descriptions using canonical API
+    tool_descs = []
+    for t in tools:
+        if not isinstance(t, dict):
+            continue
+        try:
+            td = ToolDescription.new(
+                t.get("name", "unknown"),
+                t.get("description", "") or "",
+                parameters=t.get("schema") or {},
+            )
+            tool_descs.append(td)
+        except Exception:
+            # Skip malformed tools
+            pass
+
+    # Build DeveloperContent with instructions and tools
+    dev = DeveloperContent.new()
+    if combined_instructions:
+        dev = dev.with_instructions(combined_instructions)
+    if tool_descs:
+        dev = dev.with_function_tools(tool_descs)
+
+    # Build conversation using canonical from_role_and_content
+    convo_msgs = [
+        Message.from_role_and_content(Role.SYSTEM, SystemContent.new()),
+        Message.from_role_and_content(Role.DEVELOPER, dev),
+    ]
+
+    # Add non-system messages
+    for m in messages:
+        role = (m.get("role") or "").lower()
+        if role == "system":
+            continue  # Already handled in DeveloperContent
+
+        content = m.get("content", "") or ""
+
+        if role == "user":
+            convo_msgs.append(Message.from_role_and_content(Role.USER, content))
+        elif role == "assistant":
+            convo_msgs.append(Message.from_role_and_content(Role.ASSISTANT, content))
+
+    convo = Conversation.from_messages(convo_msgs)
+    tokens = encoding.render_conversation_for_completion(convo, Role.ASSISTANT)
+    return len(tokens)
+
+
+def count_tokens_for_hf_model(messages: list[dict], tools: list[dict], model_path: str) -> int:
+    """
+    Count tokens for HuggingFace models using tokenizers library.
+
+    Uses ChatML-style template as approximation since we don't have
+    the exact chat template. This may slightly over-count which is safer
+    than under-counting.
+    """
+    tokenizer_info = get_tokenizer_for_model(model_path)
+
+    # Build prompt with ChatML-style template (common for Qwen, etc.)
+    # This is an approximation but errs on the side of over-counting
+    prompt_parts = []
+
+    # Add tools as system content
+    if tools:
+        tool_lines = ["You have access to the following tools:\n"]
+        for tool in tools:
+            name = tool.get("name", "unknown")
+            desc = tool.get("description", "")
+            schema = tool.get("schema", {})
+            tool_lines.append(f"### {name}")
+            if desc:
+                tool_lines.append(f"{desc}")
+            if schema:
+                tool_lines.append(f"Parameters: {json.dumps(schema)}")
+            tool_lines.append("")
+        prompt_parts.append(f"<|im_start|>system\n{''.join(tool_lines)}<|im_end|>")
+
+    # Add messages
+    for msg in messages:
+        role = msg.get("role", "user")
+        content = msg.get("content", "")
+        prompt_parts.append(f"<|im_start|>{role}\n{content}<|im_end|>")
+
+    full_prompt = "\n".join(prompt_parts)
+
+    if tokenizer_info is None:
+        # Conservative estimate: 1 char = 1 token (over-estimates)
+        return len(full_prompt)
+
+    tokenizer_type, tokenizer = tokenizer_info
+    try:
+        if tokenizer_type == "hf":
+            encoded = tokenizer.encode(full_prompt)
+            return len(encoded.ids)
+        else:
+            # Shouldn't happen, but handle gracefully
+            return len(full_prompt)
+    except Exception:
+        return len(full_prompt)
+
+
+def count_tokens_for_prompt(messages: list[dict], tools: list[dict], model_path: str) -> int:
+    """
+    Count tokens for a prompt with the correct format for the model.
+
+    For gpt-oss: Uses Harmony render_conversation_for_completion (exact)
+    For HF models: Uses tokenizer with chat template approximation (conservative)
+    Fallback: Char-based estimate (very conservative)
+    """
+    is_gpt_oss = "gpt-oss" in model_path.lower()
+
+    if is_gpt_oss:
+        return count_tokens_for_gpt_oss(messages, tools)
+
+    return count_tokens_for_hf_model(messages, tools, model_path)
+
+
+def get_max_prompt_tokens(
+    path: Path,
+    tools: list[dict],
+    model_path: str = "",
+    sample_size: int = 0,
+    top_k: int = 20,
+) -> int:
+    """
+    Scan train.jsonl and return the maximum prompt token count.
+
+    Uses top-K strategy: keeps top K candidates by char count, then
+    token-counts all of them to find the true maximum. This handles
+    cases where char count doesn't correlate with token count
+    (e.g., CJK text, code, JSON).
+
+    Args:
+        path: Path to train.jsonl
+        tools: List of tool definitions to include in token count
+        model_path: Model path for tokenizer selection
+        sample_size: Max lines to scan (0 = all)
+        top_k: Number of candidates to keep for token counting
+    """
+    # Min-heap of (char_count, index, messages) - keeps smallest at top
+    # Index is needed to make tuples comparable when char_counts are equal
+    heap: list[tuple[int, int, list[dict]]] = []
+
+    try:
+        with open(path, encoding="utf-8") as f:
+            for i, line in enumerate(f):
+                if sample_size and i >= sample_size:
+                    break
+
+                line = line.strip()
+                if not line:
+                    continue
+
+                try:
+                    rec = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+
+                msgs = rec.get("messages")
+                if not isinstance(msgs, list):
+                    continue
+
+                # Calculate char count for this example
+                char_count = sum(
+                    len(m.get("content", ""))
+                    for m in msgs
+                    if isinstance(m, dict) and isinstance(m.get("content", ""), str)
+                )
+
+                # Keep top_k largest by char count
+                if len(heap) < top_k:
+                    heapq.heappush(heap, (char_count, i, msgs))
+                elif char_count > heap[0][0]:
+                    heapq.heappushpop(heap, (char_count, i, msgs))
+
+    except Exception:
+        return 0
+
+    if not heap:
+        return 0
+
+    # Token-count all candidates and return the maximum
+    max_tokens = 0
+    for _chars, _idx, msgs in heap:
+        tokens = count_tokens_for_prompt(msgs, tools, model_path)
+        max_tokens = max(max_tokens, tokens)
+
+    return max_tokens
+
+
+def estimate_tokens_from_chars(char_count: int) -> int:
+    """
+    Conservative token estimate from character count.
+    Uses 2 chars per token which typically over-estimates.
+    """
+    return char_count // 2

rnow/core/__init__.py

@@ -0,0 +1,33 @@
+"""
+ReinforceNow Core - Entry points for reward and tool decorators.
+
+Users only need:
+- @reward decorator for defining reward functions
+- @tool decorator for defining tool functions
+- RewardArgs for type hints in reward functions
+"""
+
+from rnow.models import RewardArgs
+
+from .reward import (
+    REWARD_REGISTRY,
+    clear_reward_registry,
+    compute_total_reward,
+    is_precondition,
+    reward,
+)
+from .tool import TOOL_REGISTRY, clear_tool_registry, tool
+
+__all__ = [
+    # User-facing API
+    "reward",
+    "tool",
+    "RewardArgs",
+    # Registries (used by CLI and trainer)
+    "REWARD_REGISTRY",
+    "TOOL_REGISTRY",
+    "clear_reward_registry",
+    "clear_tool_registry",
+    "is_precondition",
+    "compute_total_reward",
+]

rnow/core/reward.py

@@ -0,0 +1,333 @@
+"""
+Reward entry point for ReinforceNow with validation.
+
+Validates at decorator-time:
+- Function has correct signature: (args: RewardArgs, messages: list) -> float
+- Function has docstring or description
+
+Both sync and async functions are supported. Execution strategy is
+determined automatically at runtime.
+"""
+
+import inspect
+from collections.abc import Callable
+from typing import get_type_hints
+
+from rnow.models import RewardArgs
+
+# Global registry for reward functions
+REWARD_REGISTRY: dict[str, Callable] = {}
+
+
+def clear_reward_registry() -> None:
+    """Clear the reward registry (useful for testing multiple projects)."""
+    REWARD_REGISTRY.clear()
+
+
+def is_precondition(name: str) -> bool:
+    """Check if a reward function is marked as a precondition."""
+    fn = REWARD_REGISTRY.get(name)
+    if fn is None:
+        return False
+    return getattr(fn, "_is_precondition", False)
+
+
+def compute_total_reward(reward_results: dict[str, float]) -> float:
+    """
+    Compute total reward with precondition logic.
+
+    Uses the _is_precondition attribute on registered reward functions.
+    If any precondition reward is 0, total reward is 0.
+    Otherwise, total reward is sum of all rewards.
+
+    Args:
+        reward_results: Dict mapping reward name to its value
+
+    Returns:
+        Total reward value
+    """
+    precondition_sum = 0.0
+    other_sum = 0.0
+
+    for name, value in reward_results.items():
+        if is_precondition(name):
+            # If any precondition fails (returns 0), total reward is 0
+            if value == 0.0:
+                return 0.0
+            precondition_sum += value
+        else:
+            other_sum += value
+
+    return precondition_sum + other_sum
+
+
+def _validate_reward_signature(func: Callable) -> None:
+    """
+    Validate that a reward function has the correct signature.
+
+    Expected signature:
+        def reward_fn(args: RewardArgs, messages: list) -> float
+
+    Both sync and async functions are supported.
+
+    Raises:
+        TypeError: If signature doesn't match expected format
+    """
+
+    sig = inspect.signature(func)
+    params = list(sig.parameters.values())
+
+    # Check parameter count (should be exactly 2: args and messages)
+    if len(params) != 2:
+        raise TypeError(
+            f"Reward '{func.__name__}' must have exactly 2 parameters: "
+            "(args: RewardArgs, messages: list). Got {len(params)} parameters."
+        )
+
+    # Get type hints
+    hints = get_type_hints(func)
+
+    # Check first parameter (args: RewardArgs)
+    first_param = params[0]
+    if first_param.name not in hints:
+        raise TypeError(
+            f"Reward '{func.__name__}': parameter '{first_param.name}' must have "
+            "type hint 'RewardArgs'."
+        )
+    first_type = hints[first_param.name]
+    if first_type is not RewardArgs:
+        raise TypeError(
+            f"Reward '{func.__name__}': first parameter must be typed as 'RewardArgs', "
+            f"got '{first_type}'."
+        )
+
+    # Check second parameter (messages: list)
+    second_param = params[1]
+    if second_param.name not in hints:
+        raise TypeError(
+            f"Reward '{func.__name__}': parameter '{second_param.name}' must have "
+            "type hint 'list'."
+        )
+    second_type = hints[second_param.name]
+    # Allow list or List (from typing)
+    from typing import get_origin
+
+    second_origin = get_origin(second_type) or second_type
+    if second_origin not in (list, list):
+        raise TypeError(
+            f"Reward '{func.__name__}': second parameter must be typed as 'list', "
+            f"got '{second_type}'."
+        )
+
+    # Check return type
+    if "return" not in hints:
+        raise TypeError(f"Reward '{func.__name__}' must declare return type '-> float'.")
+    return_type = hints["return"]
+    if return_type is not float:
+        raise TypeError(f"Reward '{func.__name__}' must return 'float', got '{return_type}'.")
+
+
+def reward(fn: Callable = None, *, precondition: bool = False) -> Callable:
+    """
+    Decorator to register reward functions with validation.
+
+    Validates at decorator-time:
+    - Signature is (args: RewardArgs, messages: list) -> float
+
+    Both sync and async functions are supported. Execution strategy
+    is determined automatically at runtime.
+
+    Usage:
+        @reward
+        def accuracy(args: RewardArgs, messages: list) -> float:
+            ground_truth = args.metadata.get("ground_truth")
+            response = messages[-1].get("content", "")
+            return 1.0 if ground_truth in response else 0.0
+
+        @reward(precondition=True)  # Mark as precondition reward
+        def format_check(args: RewardArgs, messages: list) -> float:
+            # If this returns 0, total reward is 0
+            # If this returns 1, total reward is 1 + sum(other rewards)
+            return 1.0 if valid_format else 0.0
+
+    Args:
+        precondition: If True, this reward acts as a gate:
+            - If precondition reward is 0, total reward is 0
+            - If precondition reward is 1, total reward is 1 + sum(other rewards)
+    """
+
+    def decorator(func):
+        # Validate signature (async, correct params, return type)
+        try:
+            _validate_reward_signature(func)
+        except TypeError as e:
+            raise TypeError(f"Reward registration failed: {e}") from e
+
+        # Warn if overwriting existing reward
+        if func.__name__ in REWARD_REGISTRY:
+            import warnings
+
+            warnings.warn(
+                f"Reward '{func.__name__}' is being overwritten in the registry.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        # Store metadata
+        func._is_reward = True
+        func._reward_name = func.__name__
+        func._is_precondition = precondition
+
+        # Register the function
+        REWARD_REGISTRY[func.__name__] = func
+        return func
+
+    # Support both @reward and @reward(precondition=False)
+    return decorator(fn) if fn else decorator
+
+
+def validate_rewards_file(filepath) -> list:
+    """
+    Validate a rewards.py file without executing it.
+
+    Parses the AST to find @reward decorated functions and checks:
+    - Has correct number of parameters
+    - Return values are between 0.0 and 1.0
+
+    Both sync and async functions are supported.
+
+    Returns a list of error messages (empty if valid).
+    """
+    import ast
+    from pathlib import Path
+
+    errors = []
+    filepath = Path(filepath)
+
+    try:
+        source = filepath.read_text()
+        tree = ast.parse(source, filename=str(filepath))
+    except SyntaxError as e:
+        return [f"Syntax error in {filepath.name}: {e}"]
+
+    def get_numeric_value(node) -> float | None:
+        """Extract numeric value from AST node if it's a constant."""
+        # Python 3.8+: ast.Constant
+        if isinstance(node, ast.Constant) and isinstance(node.value, int | float):
+            return float(node.value)
+        # Python 3.7: ast.Num
+        if hasattr(ast, "Num") and isinstance(node, ast.Num):
+            return float(node.n)
+        # Handle negative numbers: ast.UnaryOp with ast.USub
+        if isinstance(node, ast.UnaryOp) and isinstance(node.op, ast.USub):
+            inner = get_numeric_value(node.operand)
+            if inner is not None:
+                return -inner
+        return None
+
+    def check_return_values(func_node, func_name: str):
+        """Check all return statements in a function for valid reward values."""
+        for child in ast.walk(func_node):
+            if isinstance(child, ast.Return) and child.value is not None:
+                # Check direct numeric returns
+                value = get_numeric_value(child.value)
+                if value is not None and (value < 0.0 or value > 1.0):
+                    errors.append(
+                        f"Reward '{func_name}' returns {value}, must be between 0.0 and 1.0 "
+                        f"(line {child.lineno})"
+                    )
+
+                # Check ternary expressions: x if cond else y
+                if isinstance(child.value, ast.IfExp):
+                    for branch in [child.value.body, child.value.orelse]:
+                        value = get_numeric_value(branch)
+                        if value is not None and (value < 0.0 or value > 1.0):
+                            errors.append(
+                                f"Reward '{func_name}' returns {value}, must be between 0.0 and 1.0 "
+                                f"(line {child.lineno})"
+                            )
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
+            # Check if function has @reward decorator
+            is_reward = False
+            for decorator in node.decorator_list:
+                if (
+                    isinstance(decorator, ast.Name)
+                    and decorator.id == "reward"
+                    or (
+                        isinstance(decorator, ast.Call)
+                        and isinstance(decorator.func, ast.Name)
+                        and decorator.func.id == "reward"
+                    )
+                ):
+                    is_reward = True
+
+            if is_reward:
+                # Both async and sync functions are allowed
+
+                # Check parameter count
+                params = node.args.args
+                if len(params) != 2:
+                    errors.append(
+                        f"Reward '{node.name}' must have exactly 2 parameters: "
+                        f"(args: RewardArgs, messages: list). Got {len(params)} parameters."
+                    )
+
+                # Check return type annotation
+                if node.returns is None:
+                    errors.append(
+                        f"Reward '{node.name}' must have a return type annotation '-> float'."
+                    )
+
+                # Check parameter type annotations
+                for i, arg in enumerate(params):
+                    if arg.annotation is None:
+                        param_hint = "RewardArgs" if i == 0 else "list"
+                        errors.append(
+                            f"Reward '{node.name}': parameter '{arg.arg}' must have type hint '{param_hint}'."
+                        )
+
+                # Check return values are between 0.0 and 1.0
+                check_return_values(node, node.name)
+
+    return errors
+
+
+def get_reward_names_from_file(filepath) -> set[str]:
+    """
+    Extract reward function names from a rewards.py file without executing it.
+
+    Parses the AST to find @reward decorated functions and returns their names.
+
+    Returns:
+        Set of reward function names defined in the file.
+    """
+    import ast
+    from pathlib import Path
+
+    names = set()
+    filepath = Path(filepath)
+
+    try:
+        source = filepath.read_text()
+        tree = ast.parse(source, filename=str(filepath))
+    except SyntaxError:
+        return names
+
+    for node in ast.walk(tree):
+        if isinstance(node, ast.FunctionDef | ast.AsyncFunctionDef):
+            # Check if function has @reward decorator
+            for decorator in node.decorator_list:
+                if (
+                    isinstance(decorator, ast.Name)
+                    and decorator.id == "reward"
+                    or (
+                        isinstance(decorator, ast.Call)
+                        and isinstance(decorator.func, ast.Name)
+                        and decorator.func.id == "reward"
+                    )
+                ):
+                    names.add(node.name)
+
+    return names