Something went wrong. Try again.
Parametric learning in your code harness.
Something went wrong. Try again.
16 kB · 453 lines
Python
at letta-code
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454"""Pi-shaped file and shell tools for bounded multilingual post-training."""
from __future__ import annotations
import jsonimport reimport shlexfrom collections.abc import Sequencefrom dataclasses import dataclassfrom pathlib import PurePosixPathfrom typing import Annotated, Any
import chzfrom euler_tool_rl import close_sandbox_clients, sandbox_clientfrom euler_variants import LANGUAGES, SANDBOX_LANGUAGES, ensure_loopback_sandbox_urlfrom pi_rl_tasks import FAMILIES, PiRLTask, build_pi_rl_tasksfrom 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.tool_use import ( ToolResult, build_agent_tool_env, simple_tool_result, tool,)
RESULT_SCHEMA = "machine.pi-workspace-tool.v1"RUNTIME_COMMANDS = { "python": ("python", "python3"), "javascript": ("node",), "julia": ("julia",),}DEFAULT_FILENAMES = { "python": "solution.py", "javascript": "solution.js", "julia": "solution.jl",}
def _safe_path(value: str) -> str | None: path = PurePosixPath(value) if not value or path.is_absolute() or ".." in path.parts: return None normalized = str(path) return normalized if normalized not in {"", "."} else None
def _last_nonempty_line(value: str) -> str: lines = [line.strip() for line in value.splitlines() if line.strip()] return lines[-1] if lines else ""
class PiWorkspaceTools: def __init__( self, task: PiRLTask, language: str, sandbox_url: str, timeout: int, ) -> None: self.task = task self.language = language self.files = dict(task.files) self.mutated: set[str] = set() self.sandbox_url = ensure_loopback_sandbox_url(sandbox_url) self.timeout = timeout
def _result(self, tool_name: str, **payload: Any) -> ToolResult: return simple_tool_result( json.dumps( {"schema": RESULT_SCHEMA, "tool": tool_name, **payload}, sort_keys=True, ) )
@tool async def read( self, path: Annotated[str, "Workspace-relative file path."], ) -> ToolResult: """Read a UTF-8 text file from the assigned workspace.""" normalized = _safe_path(path) if normalized is None or normalized not in self.files: return self._result("read", status="error", error="file not found") return self._result("read", status="ok", path=normalized, content=self.files[normalized])
@tool async def write( self, path: Annotated[str, "Workspace-relative destination path."], content: Annotated[str, "Complete UTF-8 file content."], ) -> ToolResult: """Create or overwrite a UTF-8 text file in the assigned workspace.""" normalized = _safe_path(path) if normalized is None: return self._result("write", status="error", error="path denied") self.files[normalized] = content self.mutated.add(normalized) return self._result("write", status="ok", path=normalized, bytes=len(content.encode()))
@tool async def edit( self, path: Annotated[str, "Workspace-relative file path."], oldText: Annotated[str, "Exact text to replace."], newText: Annotated[str, "Replacement text."], ) -> ToolResult: """Replace one exact text occurrence in a workspace file.""" normalized = _safe_path(path) if normalized is None or normalized not in self.files: return self._result("edit", status="error", error="file not found") if not oldText or self.files[normalized].count(oldText) != 1: return self._result("edit", status="error", error="oldText must match once") self.files[normalized] = self.files[normalized].replace(oldText, newText, 1) self.mutated.add(normalized) return self._result("edit", status="ok", path=normalized)
@tool async def bash( self, command: Annotated[str, "One assigned-runtime command using a relative file path."], ) -> ToolResult: """Execute one program file with the assigned runtime in an isolated sandbox.""" try: argv = shlex.split(command) except ValueError as error: return self._result("bash", status="error", error=str(error)) if len(argv) != 2 or argv[0] not in RUNTIME_COMMANDS[self.language]: return self._result( "bash", status="error", error=f"use {RUNTIME_COMMANDS[self.language][0]} <relative-file>", ) path = _safe_path(argv[1]) if path is None or path not in self.files: return self._result("bash", status="error", error="program file not found") if path not in self.mutated: return self._result("bash", status="error", error="program file was not created or edited") client = sandbox_client(self.sandbox_url) _transport_success, response = await client.run( code=self.files[path], files={name: content for name, content in self.files.items() if name != path}, timeout=self.timeout, language=SANDBOX_LANGUAGES[self.language], ) run_result = response.get("run_result") if not isinstance(run_result, dict): return self._result( "bash", status="infrastructure_error", error=str(response.get("error") or response.get("message") or response)[:1000], ) stdout = str(run_result.get("stdout", "")) return_code = int(run_result.get("return_code", -1)) output_correct = return_code == 0 and _last_nonempty_line(stdout) == self.task.answer return self._result( "bash", status="ok" if return_code == 0 else "submission_error", return_code=return_code, stdout=stdout[:512], stderr=str(run_result.get("stderr", ""))[:1200], output_correct=output_correct, )
def _tool_payloads(history: list[Message]) -> list[dict[str, Any]]: payloads: list[dict[str, Any]] = [] for message in history: if message.get("role") != "tool": continue try: payload = json.loads(renderers.get_text_content(message)) except (TypeError, json.JSONDecodeError): continue if isinstance(payload, dict) and payload.get("schema") == RESULT_SCHEMA: payloads.append(payload) return payloads
def _final_answer_correct(history: list[Message], expected: str) -> bool: for message in reversed(history): if message.get("role") != "assistant": continue values = re.findall(r"(?<!\w)-?\d+(?!\w)", renderers.get_text_content(message)) return bool(values and values[-1] == expected) return False
@dataclass(frozen=True)class PiWorkspaceReward: task: PiRLTask binary: bool = False
async def __call__(self, history: list[Message]) -> tuple[float, dict[str, float]]: payloads = _tool_payloads(history) mutation = any( payload.get("tool") in {"write", "edit"} and payload.get("status") == "ok" for payload in payloads ) bash_results = [payload for payload in payloads if payload.get("tool") == "bash"] executed = any(payload.get("status") == "ok" for payload in bash_results) output_correct = any(bool(payload.get("output_correct")) for payload in bash_results) final_correct = _final_answer_correct(history, self.task.answer) infrastructure_error = any( payload.get("status") == "infrastructure_error" for payload in payloads ) fully_correct = mutation and executed and output_correct and final_correct if infrastructure_error: reward = -0.1 elif self.binary: reward = float(fully_correct) elif fully_correct: reward = 1.0 else: reward = ( 0.15 * float(mutation) + 0.15 * float(executed) + 0.6 * float(output_correct) + 0.1 * float(final_correct) ) return reward, { "tool_called": float(bool(payloads)), "mutation_succeeded": float(mutation), "execution_succeeded": float(executed), "execution_output_correct": float(output_correct), "response_answer_correct": float(final_correct), "functional_correct": float(fully_correct), "infrastructure_error": float(infrastructure_error), }
def _prompt(task: PiRLTask, language: str) -> str: filename = DEFAULT_FILENAMES[language] command = f"{RUNTIME_COMMANDS[language][0]} {filename}" files = "\n".join(f"- `{name}`" for name in sorted(task.files)) return f"""Use {language} to complete this workspace task.
{task.instruction}
Available files:{files}
Use the file tools to inspect inputs and write `{filename}`. Execute it with exactly `{command}`. Use the execution output to verify the result, then state the final integer in your response. Work only with relative paths in this workspace."""
def build_pi_workspace_env( *, task: PiRLTask, 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: bool,) -> Env: workspace = PiWorkspaceTools(task, language, sandbox_url, sandbox_timeout) tools = [workspace.read, workspace.write, workspace.edit, workspace.bash] prefix = renderer.create_conversation_prefix_with_tools( tools=[workspace_tool.to_spec() for workspace_tool in tools] ) environment = build_agent_tool_env( renderer=renderer, tools=tools, initial_messages=prefix + [{"role": "user", "content": _prompt(task, language)}], reward_fn=PiWorkspaceReward(task, binary=evaluation), 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, ) environment.example_id = f"{task.task_id}:{language}" return environment
@dataclass(frozen=True)class PiWorkspaceGroupBuilder(EnvGroupBuilder): task: PiRLTask 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 evaluation: bool = False
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), ) return [ build_pi_workspace_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, evaluation=self.evaluation, ) for _ in range(self.group_size) ]
def logging_tags(self) -> list[str]: return ["pi_workspace", self.task.family, self.task.split, self.language]
class RepeatingPiDataset(RLDataset): def __init__( self, builders: list[PiWorkspaceGroupBuilder], 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 PiWorkspaceDatasetBuilder(RLDatasetBuilder): model_name_for_tokenizer: str renderer_name: str families: str train_language_schedule: str eval_languages: str train_tasks_per_family: int eval_tasks_per_family: 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[PiRLTask], languages: list[str], group_size: int, *, evaluation: bool, ) -> list[PiWorkspaceGroupBuilder]: return [ PiWorkspaceGroupBuilder( 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, evaluation=evaluation, ) for task in tasks for language in languages ]
async def __call__(self) -> tuple[RLDataset, RLDataset | None]: ensure_loopback_sandbox_url(self.sandbox_url) families = tuple(value.strip() for value in self.families.split(",") if value.strip()) train_languages = [ value.strip() for value in self.train_language_schedule.split(",") if value.strip() ] eval_languages = [ value.strip() for value in self.eval_languages.split(",") if value.strip() ] unknown_families = sorted(set(families) - set(FAMILIES)) unknown_languages = sorted(set(train_languages + eval_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_pi_rl_tasks( split="train", tasks_per_family=self.train_tasks_per_family, seed=self.seed, families=families, ) eval_tasks = build_pi_rl_tasks( split="eval", tasks_per_family=self.eval_tasks_per_family, seed=self.seed, families=families, ) if {task.task_digest for task in train_tasks} & { task.task_digest for task in eval_tasks }: raise RuntimeError("Pi-RL train/eval task overlap") train_builders = self._builders( train_tasks, train_languages, self.group_size, evaluation=False, ) eval_builders = self._builders( eval_tasks, eval_languages, 1, evaluation=True, ) return ( RepeatingPiDataset(train_builders, self.batch_size, self.n_batches), RepeatingPiDataset(eval_builders, len(eval_builders), None), )
__all__ = [ "PiWorkspaceDatasetBuilder", "PiWorkspaceReward", "PiWorkspaceTools", "close_sandbox_clients",]