benchmax 0.1.2.dev16__tar.gz → 0.1.2.dev17__tar.gz

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev16
+Version: 0.1.2.dev17
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: cgft.io
 Classifier: Programming Language :: Python :: 3

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "benchmax"
-version = "0.1.2.dev16"
+version = "0.1.2.dev17"
 description = "Framework-Agnostic RL Environments for LLM Fine-Tuning"
 readme = "README.md"
 authors = [{ name = "cgft.io" }]

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/src/benchmax/adapters/benchmax_wrapper.py

@@ -5,6 +5,7 @@ from ray.actor import ActorClass, ActorProxy
 from typing import Dict, List, Any, Optional, Type, Union
 
 from benchmax.envs.base_env import BaseEnv
+from benchmax.envs.types import Completion
 
 # 5 minutes timeout in seconds
 RAY_GET_TIMEOUT = 300
@@ -67,7 +68,7 @@ class BenchmaxEnv:
 
     @ray.method
     async def compute_reward(
-        self, rollout_id: str, completion: str, ground_truth: Any, **kwargs: Any
+        self, rollout_id: str, completion: Completion, ground_truth: Any, **kwargs: Any
     ) -> Dict[str, float]:
         return await self._env.compute_reward(
             rollout_id=rollout_id,
@@ -258,7 +259,7 @@ class BenchmaxEnvWrapper:
     async def compute_reward(
         self,
         rollout_id: str,
-        completion: str,
+        completion: Completion,
         ground_truth: Any,
         **kwargs: Any,
     ) -> Dict[str, float]:
@@ -271,7 +272,7 @@ class BenchmaxEnvWrapper:
     def compute_reward_sync(
         self,
         rollout_id: str,
-        completion: str,
+        completion: Completion,
         ground_truth: Any,
         **kwargs: Any,
     ) -> Dict[str, float]:

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/src/benchmax/envs/base_env.py

@@ -1,9 +1,10 @@
 from abc import ABC, abstractmethod
+from functools import wraps
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
 
-from benchmax.envs.tracking import TrackingConfig, log_env, with_tracking
-from benchmax.envs.types import StandardizedExample, ToolDefinition
+from benchmax.envs.tracking import TrackingConfig, log_env, tracking_context
+from benchmax.envs.types import Completion, StandardizedExample, ToolDefinition
 from benchmax.prompts.tools import render_tools_prompt
 
 if TYPE_CHECKING:
@@ -16,32 +17,28 @@ class BaseEnv(ABC):
     system_prompt: str = ""
     _tracking_config: TrackingConfig | None = None
 
-    def __init_subclass__(cls, **kwargs):
-        super().__init_subclass__(**kwargs)
+    def __init__(self, **kwargs):
+        self._tracking_config: Optional[TrackingConfig] = None
 
-        compute_reward = cls.__dict__.get("compute_reward")
-        if compute_reward is None:
-            return
-        if getattr(compute_reward, "__benchmax_tracking_wrapped__", False):
-            return
-
-        wrapped = with_tracking(lambda self, *a, **kw: self.get_tracking_config())(
-            compute_reward
-        )
-        setattr(wrapped, "__benchmax_tracking_wrapped__", True)
-        setattr(cls, "compute_reward", wrapped)
-
-    def __init__(
+    def enable_tracking(
         self,
         experiment_id: Optional[str] = None,
         api_key: Optional[str] = None,
-        **kwargs,
-    ):
+    ) -> None:
+        """Enable experiment tracking. Wraps compute_reward on this instance with a tracking context."""
         self._tracking_config = TrackingConfig(
             experiment_id=experiment_id, api_key=api_key
         )
+        cls_compute_reward = type(self).compute_reward
 
-    def get_tracking_config(self) -> TrackingConfig | None:
+        @wraps(cls_compute_reward)
+        async def _tracked(*args, **kwargs):
+            with tracking_context(self._tracking_config):
+                return await cls_compute_reward(self, *args, **kwargs)
+
+        self.compute_reward = _tracked
+
+    def get_tracking_config(self) -> Optional[TrackingConfig]:
         return self._tracking_config
 
     def log_env(self, rollout_id: str, message: str) -> None:
@@ -106,7 +103,7 @@ class BaseEnv(ABC):
 
     @abstractmethod
     async def compute_reward(
-        self, rollout_id: str, completion: str, ground_truth: Any, **kwargs: Any
+        self, rollout_id: str, completion: Completion, ground_truth: Any, **kwargs: Any
     ) -> Dict[str, float]:
         """Compute rewards using registered functions
 
@@ -114,6 +111,34 @@ class BaseEnv(ABC):
         """
         pass
 
+    async def compute_group_reward(
+        self,
+        rollout_ids: List[str],
+        completions: List[str | List[Dict[str, str]]],
+        ground_truths: List[Any],
+        **kwargs: Any,
+    ) -> List[Dict[str, float]]:
+        """Compute rewards across a group of rollouts jointly.
+
+        Override this when reward computation requires cross-rollout context (e.g.,
+        relative scoring, group normalization, or deduplication). Can be used alongside
+        ``compute_reward`` — the two are not mutually exclusive. The default implementation
+        returns empty reward dicts, deferring entirely to per-rollout ``compute_reward`` calls.
+
+        Args:
+            rollout_ids: Identifiers for each rollout in the group.
+            completions: Model outputs, one per rollout. Each entry is either a
+                plain string or a list of message dicts.
+            ground_truths: Reference answers, one per rollout.
+            **kwargs: Additional environment-specific arguments.
+
+        Returns:
+            A list of reward dicts (one per rollout), each mapping reward function
+            names to their computed scores. An empty dict signals that no group
+            reward was computed for that rollout.
+        """
+        return [{} for _ in rollout_ids]
+
     async def get_system_prompt(self, add_tool_defs: bool = False) -> str:
         """Get system prompt. To add tool definitions, set add_tool_defs to True."""
         if add_tool_defs:

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/src/benchmax/envs/crm/workdir/reward_fn.py

@@ -95,7 +95,7 @@ def get_all_metrics(proposed_answer: str, ground_truth: str) -> float:
 
 
 def crm_matching_reward_function(
-    completion: str,
+    completion: List[Dict[str, Any]],
     ground_truth: List[str],
     mcp_client: Client,
     workspace: Path,
@@ -119,7 +119,8 @@ def crm_matching_reward_function(
     if not reward_metric:
         raise ValueError("kwargs must contain reward metric")
 
-    proposed_answer = completion.strip() if completion else ""
+    completion_text = completion[-1].get("content", "") if completion else ""
+    proposed_answer = completion_text.strip() if completion_text else ""
     proposed_answer = parse_answers(proposed_answer)
 
     if reward_metric == "exact_match":

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/src/benchmax/envs/excel/excel_env.py

@@ -7,7 +7,7 @@ from benchmax.envs.mcp.parallel_mcp_env import ParallelMcpEnv
 from benchmax.envs.mcp.provisioners.base_provisioner import BaseProvisioner
 from benchmax.envs.mcp.provisioners.local_provisioner import LocalProvisioner
 from benchmax.envs.mcp.provisioners.skypilot_provisioner import SkypilotProvisioner
-from benchmax.envs.types import StandardizedExample
+from benchmax.envs.types import Completion, StandardizedExample
 from .data_utils import download_and_extract
 
 # Using library shared with mcp workdir
@@ -162,7 +162,7 @@ Output Path: {output_filename}"""
         await self.copy_to_workspace(rollout_id, Path(input_src_path))
 
     async def compute_reward(
-        self, rollout_id: str, completion: str, ground_truth: Any, **kwargs: Any
+        self, rollout_id: str, completion: Completion, ground_truth: Any, **kwargs: Any
     ) -> Dict[str, float]:
         answer_position: Optional[str] = kwargs.get("answer_position")
         output_filename: Optional[str] = kwargs.get("output_filename")

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/src/benchmax/envs/excel/workdir/reward_fn.py

@@ -1,5 +1,5 @@
 from pathlib import Path
-from typing import Any, Awaitable, Callable, Dict, Optional, Union
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Union
 
 from fastmcp import Client
 
@@ -13,7 +13,7 @@ RewardFunction = Callable[..., Union[float, Awaitable[float]]]
 
 
 def spreadsheet_comparison_reward(
-    completion: str,
+    completion: List[Dict[str, Any]],
     ground_truth: dict,
     mcp_client: Client,
     workspace: Path,
@@ -34,23 +34,17 @@ def spreadsheet_comparison_reward(
 
     output_path = workspace / output_filename
     ground_truth_path = workspace / ground_truth_filename
-    rollout_id = kwargs.get("rollout_id", "unknown_rollout")
-    log_env(
-        rollout_id, f"excel_reward:compare_files={ground_truth_filename}:{output_filename}:{answer_position}"
-    )
 
     # Return 1.0 score if the output completely matches the ground truth
     try:
         match, _ = compare_excel_cells(
             str(ground_truth_path), str(output_path), answer_position
         )
-        log_env(rollout_id, f"excel_reward:spreadsheet_match={float(match)}")
         return 1.0 if match else 0.0
     except Exception as e:
         print(
             f"Error comparing spreadsheets {ground_truth_path} and {output_path}: {e}"
         )
-        log_env(rollout_id, f"excel_reward:error={str(e)}")
         return 0.0
 
 

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/src/benchmax/envs/math/workdir/reward_fn.py

@@ -1,6 +1,6 @@
 from pathlib import Path
 import re
-from typing import Any, Callable, Dict, Union, Awaitable
+from typing import Any, Callable, Dict, List, Union, Awaitable
 from fastmcp import Client
 from html import unescape
 
@@ -8,7 +8,7 @@ RewardFunction = Callable[..., Union[float, Awaitable[float]]]
 
 
 async def text_match_reward(
-    completion: str,
+    completion: List[Dict[str, Any]],
     ground_truth: str,
     mcp_client: Client,
     workspace: Path,
@@ -21,9 +21,11 @@ async def text_match_reward(
     Falls back to 0 if the tag is missing or empty.
     """
 
+    completion_text = completion[-1].get("content", "") if completion else ""
+
     # Grab only the text inside the first <answer> … </answer> pair (case-insensitive).
     m = re.search(
-        r"<answer>(.*?)</answer>", completion, flags=re.IGNORECASE | re.DOTALL
+        r"<answer>(.*?)</answer>", completion_text, flags=re.IGNORECASE | re.DOTALL
     )
     if m is None:
         return 0.0

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/src/benchmax/envs/mcp/example_workdir/reward_fn.py

@@ -23,7 +23,7 @@ and reaading from the workspace that the MCP is operating in.
 """
 
 from pathlib import Path
-from typing import Any, Callable, Dict, Union, Awaitable
+from typing import Any, Callable, Dict, List, Union, Awaitable
 from mcp.types import TextContent
 from fastmcp import Client
 from fastmcp.exceptions import ToolError
@@ -35,7 +35,7 @@ RewardFunction = Callable[..., Union[float, Awaitable[float]]]
 # Reward 0: Stateless completion check
 # -------------------------------
 async def completion_match_reward(
-    completion: str,
+    completion: List[Dict[str, Any]],
     ground_truth: dict,
     mcp_client: Client,
     workspace: Path,
@@ -47,6 +47,9 @@ async def completion_match_reward(
     Uses: ground_truth['completion'] (str)
     """
     expected = ground_truth.get("completion", "")
+    if isinstance(completion, list):
+        completion = completion[-1].get("content", "") if completion else ""
+    completion = str(completion)
     return 1.0 if completion.strip() == expected.strip() else 0.0
 
 
@@ -54,7 +57,7 @@ async def completion_match_reward(
 # Reward 1: Tool call variable in memory check
 # -------------------------------
 async def variable_in_memory_reward(
-    completion: str, ground_truth: dict, mcp_client: Client, workspace: Path, **kwargs
+    completion: str | list[dict], ground_truth: dict, mcp_client: Client, workspace: Path, **kwargs
 ) -> float:
     """
     Reward uses tool call to match in-memory variable value.
@@ -96,7 +99,7 @@ async def variable_in_memory_reward(
 # Reward 2: Workspace log check
 # -------------------------------
 async def log_in_workspace_reward(
-    completion: str,
+    completion: str | list[dict],
     ground_truth: dict,
     mcp_client: Client,
     workspace: Path,

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/src/benchmax/envs/mcp/parallel_mcp_env.py

@@ -20,7 +20,7 @@ except ModuleNotFoundError as e:
 
 from benchmax.envs.base_env import BaseEnv
 from benchmax.envs.tracking import to_tracking_payload
-from benchmax.envs.types import ToolDefinition
+from benchmax.envs.types import Completion, ToolDefinition
 from .server_pool import ServerPool
 from .provisioners.base_provisioner import BaseProvisioner
 from .utils import (
@@ -344,7 +344,7 @@ class ParallelMcpEnv(BaseEnv):
             return str(e)
 
     async def compute_reward(
-        self, rollout_id: str, completion: str, ground_truth: Any, **kwargs: Any
+        self, rollout_id: str, completion: Completion, ground_truth: Any, **kwargs: Any
     ) -> Dict[str, float]:
         """
         Compute reward and cleanup rollout.

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/src/benchmax/envs/types.py

@@ -1,9 +1,11 @@
 from dataclasses import dataclass
-from typing import Any, Dict, Optional, TypedDict
+from typing import Any, Dict, List, Optional, TypedDict
+
+Completion = List[Dict[str, Any]]
 
 
 class StandardizedExample(TypedDict):
-    prompt: str
+    prompt: str | List[Dict[str, Any]]
     ground_truth: Any
     init_rollout_args: Optional[Dict[str, Any]]
 

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/src/benchmax/envs/wikipedia/wiki_env.py

@@ -5,7 +5,7 @@ from typing import Any, Callable, Dict, List, Optional, Tuple
 
 from benchmax.envs.base_env import BaseEnv
 from benchmax.envs.tracking import log_env
-from benchmax.envs.types import ToolDefinition, StandardizedExample
+from benchmax.envs.types import Completion, ToolDefinition, StandardizedExample
 from benchmax.envs.wikipedia.utils import APIKeyRotator, clean_html, safe_request
 
 SYSTEM_PROMPT = """Please use the tools provided to get accurate, up-to-date information.
@@ -14,12 +14,12 @@ Write your complete answer on the final line only as a concise entity, within th
 """
 
 
-def text_match_reward_function(completion: str, ground_truth: str, rollout_id: str, **kwargs) -> float:
+def text_match_reward_function(completion: Completion, ground_truth: str, rollout_id: str, **kwargs) -> float:
     """
     Score 1.0 if ground truth appears in <answer> tags, else 0.0.
 
     Args:
-        completion: The model's generated text
+        completion: The model's generated text (str or list of message dicts)
         ground_truth: Expected answer (case-insensitive)
         **kwargs: Catch-all for BaseEnv compatibility
 
@@ -27,17 +27,22 @@ def text_match_reward_function(completion: str, ground_truth: str, rollout_id: s
         1.0 if ground_truth matches the answer text, else 0.0
     """
     assert ground_truth is not None
+    completion_str = ""
+    if isinstance(completion, list):
+        completion_str = completion[-1].get("content", "") if completion else ""
+    elif isinstance(completion, str):
+        completion_str = completion
+    else:
+        completion_str = ""
 
     m = re.search(
-        r"<answer>(.*?)</answer>", completion, flags=re.IGNORECASE | re.DOTALL
+        r"<answer>(.*?)</answer>", completion_str, flags=re.IGNORECASE | re.DOTALL
     )
     if not m:
-        log_env(rollout_id, "wikipedia_reward:no_answer_tag")
         return 0.0
 
     answer_text = unescape(m.group(1)).strip().lower()
     score = float(ground_truth.lower() == answer_text)
-    log_env(rollout_id, f"wikipedia_reward:text_match={score}")
     return score
 
 
@@ -264,7 +269,7 @@ class WikipediaEnv(BaseEnv):
         pass
 
     async def compute_reward(
-        self, rollout_id: str, completion: str, ground_truth: Any, **kwargs: Any
+        self, rollout_id: str, completion: Completion, ground_truth: Any, **kwargs: Any
     ) -> Dict[str, float]:
         """Compute rewards using the text match reward function."""
         return {

{benchmax-0.1.2.dev16 → benchmax-0.1.2.dev17}/src/benchmax.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: benchmax
-Version: 0.1.2.dev16
+Version: 0.1.2.dev17
 Summary: Framework-Agnostic RL Environments for LLM Fine-Tuning
 Author: cgft.io
 Classifier: Programming Language :: Python :: 3