"""Tinker agent-tool environment for hidden Project Euler variants.""" from __future__ import annotations import json import secrets from collections.abc import Sequence from dataclasses import dataclass from typing import Annotated, Any import chz from euler_variants import ( FAMILIES, LANGUAGES, RESULT_SCHEMA, SANDBOX_LANGUAGES, VariantTask, assert_split_isolation, assess_sandbox_response, build_sandbox_program, build_variant_prompt, build_variant_tasks, ensure_loopback_sandbox_url, ) from tinker_cookbook import model_info, renderers, tokenizer_utils from tinker_cookbook.renderers.base import Message from tinker_cookbook.rl.types import Env, EnvGroupBuilder, RLDataset, RLDatasetBuilder from tinker_cookbook.sandbox import SandboxFusionClient from tinker_cookbook.tool_use import ( ToolResult, build_agent_tool_env, simple_tool_result, tool, ) _SANDBOX_CLIENTS: dict[str, SandboxFusionClient] = {} def sandbox_client(url: str) -> SandboxFusionClient: if url not in _SANDBOX_CLIENTS: _SANDBOX_CLIENTS[url] = SandboxFusionClient(url=url) return _SANDBOX_CLIENTS[url] async def close_sandbox_clients() -> None: clients = list(_SANDBOX_CLIENTS.values()) _SANDBOX_CLIENTS.clear() for client in clients: await client.close() async def verify_sandbox_languages(sandbox_url: str, timeout: int) -> None: """Fail before Tinker setup if the local sandbox lacks a required runtime.""" url = ensure_loopback_sandbox_url(sandbox_url) client = SandboxFusionClient(url=url) programs = { "python": "print('__MACHINE_SANDBOX_READY__')", "javascript": "console.log('__MACHINE_SANDBOX_READY__');", "julia": 'println("__MACHINE_SANDBOX_READY__")', } try: for language, code in programs.items(): success, response = await client.run( code=code, files={}, timeout=timeout, language=SANDBOX_LANGUAGES[language], ) run_result = response.get("run_result") ready = ( success and isinstance(run_result, dict) and run_result.get("return_code") == 0 and "__MACHINE_SANDBOX_READY__" in str(run_result.get("stdout", "")) ) if not ready: diagnostic = ( response.get("error") or response.get("message") or response ) raise RuntimeError( f"SandboxFusion {language} preflight failed: {diagnostic}" ) finally: await client.close() class EulerVariantTool: def __init__( self, task: VariantTask, language: str, sandbox_url: str, timeout: int, ) -> None: self.task = task self.language = language self.sandbox_url = ensure_loopback_sandbox_url(sandbox_url) self.timeout = timeout @tool async def check_solution( self, code: Annotated[str, "Complete source defining the requested solve function."], ) -> ToolResult: """Run the proposed solve function against hidden cases in the local sandbox.""" client = sandbox_client(self.sandbox_url) result_marker = f"{secrets.token_hex(16)}__" transport_success, response = await client.run( code=build_sandbox_program( self.task, self.language, code, result_marker=result_marker, ), files={}, timeout=self.timeout, language=SANDBOX_LANGUAGES[self.language], ) assessment = assess_sandbox_response( self.task, transport_success=transport_success, response=response, result_marker=result_marker, ) payload = { "schema": RESULT_SCHEMA, "status": assessment.status, "passed": assessment.passed, "passed_cases": assessment.passed_cases, "total_cases": assessment.total_cases, } if assessment.diagnostic: payload["diagnostic"] = assessment.diagnostic return simple_tool_result(json.dumps(payload, sort_keys=True)) def _tool_results(history: list[Message]) -> list[dict[str, Any]]: results: list[dict[str, Any]] = [] for message in history: if message.get("role") != "tool": continue content = renderers.get_text_content(message) try: value = json.loads(content) except (TypeError, json.JSONDecodeError): continue if isinstance(value, dict) and value.get("schema") == RESULT_SCHEMA: results.append(value) return results @dataclass(frozen=True) class EulerVariantReward: """Reward hidden-case execution while preserving aggregate-only feedback.""" async def __call__(self, history: list[Message]) -> tuple[float, dict[str, float]]: results = _tool_results(history) if not results: return -0.1, { "tool_called": 0.0, "infrastructure_error": 0.0, "submission_error": 0.0, "hidden_pass_rate": 0.0, "hidden_all_passed": 0.0, } valid = [ result for result in results if result.get("status") != "infrastructure_error" ] if not valid: return -0.1, { "tool_called": 1.0, "infrastructure_error": 1.0, "submission_error": 0.0, "hidden_pass_rate": 0.0, "hidden_all_passed": 0.0, } best = max( valid, key=lambda result: ( int(result.get("passed_cases", 0)), int(bool(result.get("passed", False))), ), ) total = max(1, int(best.get("total_cases", 0))) pass_rate = min(1.0, max(0.0, int(best.get("passed_cases", 0)) / total)) all_passed = float(bool(best.get("passed", False))) reward = 0.5 * pass_rate + 0.5 * all_passed return reward, { "tool_called": 1.0, "infrastructure_error": 0.0, "submission_error": float(best.get("status") == "submission_error"), "hidden_pass_rate": pass_rate, "hidden_all_passed": all_passed, } @dataclass(frozen=True) class EulerVariantEvaluationReward: """Use binary all-cases success for checkpoint comparison and pass@k.""" async def __call__(self, history: list[Message]) -> tuple[float, dict[str, float]]: _, metrics = await EulerVariantReward()(history) return metrics["hidden_all_passed"], metrics def _initial_messages( task: VariantTask, language: str, renderer: renderers.Renderer, variant_tool: EulerVariantTool, ) -> list[Message]: prefix = renderer.create_conversation_prefix_with_tools( tools=[variant_tool.check_solution.to_spec()] ) return prefix + [{"role": "user", "content": build_variant_prompt(task, language)}] def build_variant_env( *, task: VariantTask, language: str, model_name: str, renderer: renderers.Renderer, sandbox_url: str, sandbox_timeout: int, max_turns: int, max_tool_calls: int, max_trajectory_tokens: int, max_generation_tokens: int, evaluation_reward: bool = False, example_id: str | None = None, ) -> Env: variant_tool = EulerVariantTool( task, language, sandbox_url, sandbox_timeout, ) reward_fn = ( EulerVariantEvaluationReward() if evaluation_reward else EulerVariantReward() ) environment = build_agent_tool_env( renderer=renderer, tools=[variant_tool.check_solution], initial_messages=_initial_messages( task, language, renderer, variant_tool, ), reward_fn=reward_fn, model_name=model_name, max_turns=max_turns, max_tool_calls=max_tool_calls, max_trajectory_tokens=max_trajectory_tokens, max_generation_tokens=max_generation_tokens, context_overflow_reward=-0.1, ) if example_id is not None: environment.example_id = example_id return environment @dataclass(frozen=True) class EulerVariantGroupBuilder(EnvGroupBuilder): task: VariantTask language: str model_name: str renderer_name: str group_size: int sandbox_url: str sandbox_timeout: int max_turns: int max_tool_calls: int max_trajectory_tokens: int max_generation_tokens: int async def make_envs(self) -> Sequence[Env]: renderer = renderers.get_renderer( self.renderer_name or model_info.get_recommended_renderer_name(self.model_name), tokenizer_utils.get_tokenizer(self.model_name), ) environments: list[Env] = [] for _ in range(self.group_size): environments.append( build_variant_env( task=self.task, language=self.language, model_name=self.model_name, renderer=renderer, sandbox_url=self.sandbox_url, sandbox_timeout=self.sandbox_timeout, max_turns=self.max_turns, max_tool_calls=self.max_tool_calls, max_trajectory_tokens=self.max_trajectory_tokens, max_generation_tokens=self.max_generation_tokens, ) ) return environments def logging_tags(self) -> list[str]: return [ "project_euler_variant", self.task.family, self.task.split, self.language, self.task.task_id, ] class RepeatingVariantDataset(RLDataset): def __init__( self, builders: list[EulerVariantGroupBuilder], batch_size: int, n_batches: int | None, ) -> None: self.builders = builders self.batch_size = batch_size self.n_batches = n_batches self.offset = 0 def get_batch(self, index: int) -> Sequence[EnvGroupBuilder]: if not self.builders: return [] start = (index * self.batch_size + self.offset) % len(self.builders) return [ self.builders[(start + position) % len(self.builders)] for position in range(min(self.batch_size, len(self.builders))) ] def set_epoch(self, seed: int = 0) -> None: self.offset = seed % max(1, len(self.builders)) def __len__(self) -> int: natural = (len(self.builders) + self.batch_size - 1) // self.batch_size return max(natural, self.n_batches or 0) @chz.chz class EulerVariantDatasetBuilder(RLDatasetBuilder): model_name_for_tokenizer: str renderer_name: str families: str languages: str train_tasks_per_family: int eval_tasks_per_family: int cases_per_task: int seed: int batch_size: int group_size: int n_batches: int sandbox_url: str sandbox_timeout: int max_turns: int max_tool_calls: int max_trajectory_tokens: int max_generation_tokens: int def _builders( self, tasks: list[VariantTask], languages: list[str], group_size: int, ) -> list[EulerVariantGroupBuilder]: return [ EulerVariantGroupBuilder( task=task, language=language, model_name=self.model_name_for_tokenizer, renderer_name=self.renderer_name, group_size=group_size, sandbox_url=self.sandbox_url, sandbox_timeout=self.sandbox_timeout, max_turns=self.max_turns, max_tool_calls=self.max_tool_calls, max_trajectory_tokens=self.max_trajectory_tokens, max_generation_tokens=self.max_generation_tokens, ) for task in tasks for language in languages ] async def __call__(self) -> tuple[RLDataset, RLDataset | None]: ensure_loopback_sandbox_url(self.sandbox_url) families = [ value.strip() for value in self.families.split(",") if value.strip() ] languages = [ value.strip() for value in self.languages.split(",") if value.strip() ] unknown_families = sorted(set(families) - set(FAMILIES)) unknown_languages = sorted(set(languages) - set(LANGUAGES)) if unknown_families: raise ValueError(f"unknown families: {', '.join(unknown_families)}") if unknown_languages: raise ValueError(f"unknown languages: {', '.join(unknown_languages)}") train_tasks = build_variant_tasks( families, split="train", tasks_per_family=self.train_tasks_per_family, cases_per_task=self.cases_per_task, seed=self.seed, ) eval_tasks = build_variant_tasks( families, split="eval", tasks_per_family=self.eval_tasks_per_family, cases_per_task=self.cases_per_task, seed=self.seed, ) assert_split_isolation(train_tasks, eval_tasks) train_builders = self._builders(train_tasks, languages, self.group_size) eval_builders = self._builders(eval_tasks, languages, 1) return ( RepeatingVariantDataset(train_builders, self.batch_size, self.n_batches), RepeatingVariantDataset(eval_builders, len(eval_builders), None), )