Something went wrong. Try again.
Parametric learning in your code harness.
Something went wrong. Try again.
14 kB · 421 lines
Python
at letta-code
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422"""Tinker agent-tool environment for hidden Project Euler variants."""
from __future__ import annotations
import jsonimport secretsfrom collections.abc import Sequencefrom dataclasses import dataclassfrom typing import Annotated, Any
import chzfrom 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_utilsfrom tinker_cookbook.renderers.base import Messagefrom tinker_cookbook.rl.types import Env, EnvGroupBuilder, RLDataset, RLDatasetBuilderfrom tinker_cookbook.sandbox import SandboxFusionClientfrom 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.chzclass 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), )