diff --git a/examples/README.md b/examples/README.md index e1b822c2b..518eab493 100644 --- a/examples/README.md +++ b/examples/README.md @@ -4,10 +4,9 @@ These examples provide concrete examples to leverage vime in your own RL workflo ## Directory Structure -- **[eval_multi_task](./eval_multi_task)**: Example for supporting evaluation multiple tasks with different configs. +- **[coding_agent_rl](./coding_agent_rl)**: End-to-end SWE coding-agent RL: a real coding agent (claude-code) edits code in a per-sample sandbox, and the resulting `git diff` is graded against the dataset's test harness. - **[fully_async](./fully_async)**: Demonstrates fully asynchronous rollout generation for higher efficiency. - **[geo3k_vlm](./geo3k_vlm)**: Training VLMs on a single-turn reasoning task using GRPO on the GEO3K dataset. - **[geo3k_vlm_multi_turn](./geo3k_vlm_multi_turn)**: VLM multi-turn training on Geo3k dataset. - **[multi_agent](./multi_agent)**: Example of running multi-agent RL with `vime`. -- **[reproducibility](./reproducibility)**: Guides on achieving bitwise experiment reproduction using deterministic modes. - **[train_infer_mismatch_helper](./train_infer_mismatch_helper)**: Algorithmic methods for rollout correction (e.g., TIS, MIS). diff --git a/examples/retool/generate_with_retool.py b/examples/retool/generate_with_retool.py deleted file mode 100644 index f0506db66..000000000 --- a/examples/retool/generate_with_retool.py +++ /dev/null @@ -1,461 +0,0 @@ -# Adapted from https://github.com/volcengine/verl/blob/cb809d66e46dfd3342d008628891a14a054fa424/recipe/retool/retool.py -# -# vime/vLLM counterpart of slime's ``examples/retool/generate_with_retool.py``. -# The only structural change versus slime is the rollout engine: requests go to -# vLLM's ``/inference/v1/generate`` (token-only body, ``token_ids`` not -# ``input_ids``) and the response ``choices[0]`` is normalized into a -# sglang-shaped ``meta_info`` (``finish_reason.type`` + ``output_token_logprobs``) -# via the canonical helpers in :mod:`vime.rollout.vllm_rollout`, so the -# multi-turn tool-call control flow below is identical to slime's. -import re -from typing import Any - -try: - from jinja2 import Template -except ImportError as e: - raise ImportError("Jinja2 is required. Please install it with: pip install jinja2") from e - -from vime.rollout.vllm_rollout import ( - GenerateState, - _build_inference_sampling_params, - _coerce_flat_int_token_ids, - _inference_generate_tokens_and_logprobs, - _vllm_meta_from_generate_choice, -) -from vime.utils.http_utils import post -from vime.utils.types import Sample - -# Import reward models -try: - from vime.rollout.rm_hub.math_dapo_utils import compute_score as math_dapo_compute_score -except ImportError as e: - raise ImportError("MathDapo is not installed") from e - -# Import tool sandbox functionality -from tool_sandbox import SEMAPHORE, TOOL_CONFIGS, tool_registry - -# Jinja2 template for tool-enabled conversations -TOOL_TEMPLATE = """<|im_start|>system -{%- if messages[0]['role'] == 'system' %} -{{- messages[0]['content'] }} -{%- else %} -You are a helpful assistant. -{%- endif %} -{%- if tools %} -# Tools - -You may call one or more functions to assist with the user query. - -You are provided with function signatures within XML tags: - -{%- for tool in tools %} -{{- tool | tojson }} -{%- endfor %} - - -For each function call, return a json object with function name and arguments within XML tags: - -{"name": , "arguments": } - -{%- endif %} -<|im_end|> -{%- for message in messages %} -{%- if message['role'] == 'user' %} -<|im_start|>user -{{- message['content'] }}<|im_end|> -{%- elif message['role'] == 'assistant' %} -<|im_start|>assistant -{{- message['content'] }}<|im_end|> -{%- endif %} -{%- endfor %} -<|im_start|>assistant -""" - - -def format_conversation_with_tools( - prompt: str, tools: list[dict[str, Any]] = None, system_prompt: str = None, messages: list[dict[str, Any]] = None -) -> str: - """Format conversation using Jinja2 template with tool support""" - template = Template(TOOL_TEMPLATE) - - # Prepare messages - messages_to_render = [] - - # Always add system message - use provided one or default - if system_prompt: - system_content = system_prompt - else: - system_content = ( - "You are a helpful assistant that can use Python " - "tools to solve mathematical problems. When you need " - "to perform calculations, use the code_interpreter " - "tool to execute code and get results." - ) - - messages_to_render.append({"role": "system", "content": system_content}) - - # Add user message if provided - if prompt: - messages_to_render.append({"role": "user", "content": prompt}) - - # Add assistant responses from previous turns if provided - if messages: - messages_to_render.extend(messages) - - # Render template - formatted_text = template.render(messages=messages_to_render, tools=tools or []) - - return formatted_text - - -def postprocess_predictions(prediction: str): - """Extract action and content from prediction string""" - # Check for Answer: \boxed{...} format (only format we need for math_dapo) - # Use a more robust regex that handles nested braces - answer_pattern = r"Answer:\s*\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}" - answer_match = re.search(answer_pattern, prediction, re.DOTALL) - if answer_match: - content = answer_match.group(1).strip() - return "answer", content - - # Then check for tags (new format from Jinja2 template) - tool_call_pattern = r"\s*(\{.*?\})\s*" - tool_call_match = re.search(tool_call_pattern, prediction, re.DOTALL) - if tool_call_match: - try: - import json - - # Clean up the JSON string by removing newlines and extra - # whitespace - json_str = tool_call_match.group(1) - # Replace newlines in string values with \n - json_str = json_str.replace("\n", "\\n") - tool_call_data = json.loads(json_str) - tool_name = tool_call_data.get("name") - arguments = tool_call_data.get("arguments", {}) - - if tool_name == "code_interpreter": - code = arguments.get("code", "") - if code.strip(): - return "code", code - except (json.JSONDecodeError, KeyError, AttributeError): - pass - - # Then check for tags - code_pattern = r"(.*?)" - code_match = re.search(code_pattern, prediction, re.DOTALL) - if code_match: - content = code_match.group(1).strip() - return "code", content - - # Finally check for ```python code blocks (lowest priority) - python_code_pattern = r"```python\s*(.*?)\s*```" - python_code_match = re.search(python_code_pattern, prediction, re.DOTALL) - if python_code_match: - content = python_code_match.group(1).strip() - return "code", content - - return None, "" - - -def postprocess_responses(resp: str) -> str: - """Post-process response to ensure tag completeness""" - # Handle tags (new format from Jinja2 template) - if "" in resp: - # Find the last occurrence of ... - tool_call_pattern = r"\s*\{.*?\}\s*" - matches = list(re.finditer(tool_call_pattern, resp, re.DOTALL)) - if matches: - last_match = matches[-1] - return resp[: last_match.end()] - - # Handle tags - if "" in resp: - return resp.split("")[0] + "" - - # Handle ```python code blocks - if "```python" in resp: - # Find the last occurrence of ```python...``` - python_pattern = r"```python\s*.*?```" - matches = list(re.finditer(python_pattern, resp, re.DOTALL)) - if matches: - last_match = matches[-1] - return resp[: last_match.end()] - - # Handle Answer: \boxed{...} format (only format we need for math_dapo) - if "Answer:" in resp and "\\boxed{" in resp: - # Find the last occurrence of Answer: \boxed{...} with nested braces support - answer_pattern = r"Answer:\s*\\boxed\{((?:[^{}]|\{[^{}]*\})*)\}" - matches = list(re.finditer(answer_pattern, resp, re.DOTALL)) - if matches: - last_match = matches[-1] - return resp[: last_match.end()] - - return resp - - -async def execute_predictions(prediction: str) -> str: - """Execute predictions and return results""" - action, content = postprocess_predictions(prediction) - - if action == "code": - # Content is already the Python code (extracted by - # postprocess_predictions) - code = content.strip() - if code: - async with SEMAPHORE: - result = await tool_registry.execute_tool("code_interpreter", {"code": code}) - next_obs = f"\n\n\n{result}\n\n\n" - done = False - else: - next_obs = "\n\n\nError: No Python code found" "\n\n\n" - done = False - elif action == "answer": - next_obs = "" - done = True - else: - next_obs = ( - "\nMy previous action is invalid. " - "If I want to execute code, I should put the code between " - " and . " - "If I want to give the final answer, I should use the format " - "'Answer: \\boxed{answer}'. Let me try again.\n" - ) - done = False - - return next_obs, done - - -async def generate(args, sample: Sample, sampling_params) -> Sample: - """Custom generation function supporting tool calls""" - assert not args.partial_rollout, "Partial rollout is not supported for " "this function at the moment." - - # Retried samples (previously aborted / partial) arrive here with stale - # rollout state from the first attempt. Clear it so this generation starts - # clean; otherwise the concatenation below appends new tokens to old ones - # and downstream `slice_log_prob_with_cp` sees a length mismatch. - sample.rollout_log_probs = None - sample.response = "" - sample.response_length = 0 - sample.loss_mask = None - - state = GenerateState(args) - url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}/inference/v1/generate" - - # Set up the initial prompt with system prompt and tools (outside the loop) - tool_specs = tool_registry.get_tool_specs() - prompt = format_conversation_with_tools(prompt=sample.prompt, tools=tool_specs) - - prompt_tokens_ids = state.tokenizer(prompt, add_special_tokens=False)["input_ids"] - response = "" - response_token_ids = [] - loss_masks = [] - tool_call_count = 0 # Track actual tool call rounds - - if args.rollout_max_context_len is not None: - max_context_length = args.rollout_max_context_len - else: - max_context_length = args.context_parallel_size * args.max_tokens_per_gpu - - # Track the finish reason of the most recent turn so the post-loop status - # mapping mirrors slime (which read it back off the last response). - finish_type = "stop" - - for turn in range(TOOL_CONFIGS["max_turns"]): - # Check if total length exceeds max context length - total_length = len(prompt_tokens_ids) + len(response_token_ids) - if total_length >= max_context_length: - sample.status = Sample.Status.TRUNCATED - break - - # Clamp per-turn max_new_tokens to the remaining context budget so a - # single turn cannot push total_length past max_context_length. Without - # this, a turn can append up to rollout_max_response_len tokens on top - # of a total that was just barely under the cap, producing samples - # that exceed the training-side max_tokens_per_gpu * cp_size budget - # and crash the partition/batch code (asserts or OOMs on an oversized - # partition). - remaining_budget = max_context_length - total_length - per_turn_sampling_params = dict(sampling_params) - per_turn_sampling_params["max_new_tokens"] = min( - sampling_params.get("max_new_tokens", remaining_budget), - remaining_budget, - ) - - # Use token IDs instead of text. vLLM's /inference/v1/generate is - # token-only; sampling params are nested under their own key and built - # by the canonical helper (logprobs enabled there). - current_token_ids = _coerce_flat_int_token_ids(prompt_tokens_ids + response_token_ids) - payload = { - "model": args.hf_checkpoint, - "token_ids": current_token_ids, - "sampling_params": _build_inference_sampling_params(per_turn_sampling_params), - } - - # Log payload to wandb for debugging - try: - import wandb - - if wandb.run is not None: - # Count available tools (from tool_specs) - available_tools = len(tool_specs) - # Count tools used in the current response - tools_used = response.count("") - - wandb.log( - { - "debug/payload_length": len(prompt + response), - "debug/available_tools": available_tools, - "debug/tools_used": tools_used, - "debug/turn": turn, - } - ) - except ImportError: - pass # wandb not available - - output = await post(url, payload) - - # Normalize the vLLM choice into a sglang-shaped meta_info: finish - # reason wrapped as {"type": str} and per-token (logprob, token_id) - # pairs under output_token_logprobs, exactly what the non-streaming - # vllm_rollout.generate produces. The retool control flow below then - # reads meta_info the same way slime did. - choice = output["choices"][0] - meta_info = _vllm_meta_from_generate_choice(args, choice, output.get("usage")) - cur_token_ids, cur_logprobs = _inference_generate_tokens_and_logprobs(choice) - if cur_token_ids: - meta_info["output_token_logprobs"] = [ - [float(lp), int(tid)] for lp, tid in zip(cur_logprobs, cur_token_ids, strict=False) - ] - finish_type = meta_info["finish_reason"]["type"] - - # Handle abort - if finish_type == "abort": - sample.status = Sample.Status.ABORTED - return sample - - if not meta_info.get("output_token_logprobs"): - # vLLM returned a choice but no per-token logprobs — we cannot - # recover per-token logprobs for this turn, which would desync - # rollout_log_probs from response_token_ids and blow up - # `slice_log_prob_with_cp` downstream. Abort the sample so the - # fully_async rollout manager returns the whole group to the - # buffer for retry instead of poisoning the trainer. - sample.status = Sample.Status.ABORTED - return sample - - cur_response_token_ids = [item[1] for item in meta_info["output_token_logprobs"]] - cur_response = state.tokenizer.decode(cur_response_token_ids) - cur_log_probs = [item[0] for item in meta_info["output_token_logprobs"]] - if sample.rollout_log_probs is None: - sample.rollout_log_probs = [] - sample.rollout_log_probs += cur_log_probs - - response += cur_response - response_token_ids += cur_response_token_ids - loss_masks += [1] * len(cur_response_token_ids) - - # Check length limit - if finish_type == "length": - break - - next_obs, done = await execute_predictions(cur_response) - if done: - break - - # Count tool calls (when we get interpreter output, it means a tool - # was called) - if "" in next_obs: - tool_call_count += 1 - - assert next_obs != "", "Next observation should not be empty." - obs_tokens_ids = state.tokenizer(next_obs, add_special_tokens=False)["input_ids"] - response += next_obs - response_token_ids += obs_tokens_ids - loss_masks += [0] * len(obs_tokens_ids) - - # Add dummy log probs for observation tokens (they won't be used due to loss_mask=0) - # Check if maximum tool call count reached - if sample.rollout_log_probs is not None: - sample.rollout_log_probs += [0.0] * len(obs_tokens_ids) - - assert len(response_token_ids) == len( - sample.rollout_log_probs - ), f"Token/logp length mismatch at turn {turn}: {len(response_token_ids)} tokens vs {len(sample.rollout_log_probs)} logps" - - # Tool output is appended verbatim and can push total_length past - # max_context_length (the per-turn generation was clamped to the - # remaining budget, but tool output is unconstrained). Trim tail - # tokens so the final sample fits the training budget exactly. - overflow = len(prompt_tokens_ids) + len(response_token_ids) - max_context_length - if overflow > 0: - response_token_ids = response_token_ids[:-overflow] - loss_masks = loss_masks[:-overflow] - if sample.rollout_log_probs is not None: - sample.rollout_log_probs = sample.rollout_log_probs[:-overflow] - # Resync the text field from the trimmed token list so - # reward_func's `sample.prompt + sample.response` matches what - # the model was actually trained on. decode(tokenize(text)) can - # be lossy on some tokenizers (whitespace / special-token - # collapse), but reward_func's regex is whitespace-robust and - # the trainer sees tokens, not text — so the drift is safe. - response = state.tokenizer.decode(response_token_ids) - sample.status = Sample.Status.TRUNCATED - break - - if tool_call_count >= TOOL_CONFIGS["max_tool_calls"]: - break - - # Set sample attributes - sample.tokens = prompt_tokens_ids + response_token_ids - sample.response_length = len(response_token_ids) - sample.response = response - sample.loss_mask = loss_masks - - # Store payload information for wandb logging - sample.payload_text = prompt + response - sample.payload_has_system = "<|im_start|>system" in prompt + response - sample.payload_has_tools = "# Tools" in prompt + response - - # Store tool call count for reward calculation - sample.tool_call_count = tool_call_count - - # Set status (skip if a context-budget branch above already set it) - if sample.status not in (Sample.Status.TRUNCATED, Sample.Status.ABORTED): - match finish_type: - case "length": - sample.status = Sample.Status.TRUNCATED - case "abort": - sample.status = Sample.Status.ABORTED - case "stop": - sample.status = Sample.Status.COMPLETED - - return sample - - -async def reward_func(args, sample, **kwargs): - """Tool call reward function using math_dapo as primary reward model""" - if not isinstance(sample, Sample): - raise TypeError("Sample must be an instance of Sample class.") - - # Build complete solution string - solution_str = sample.prompt + sample.response - - # Get ground truth answer - label is a string, not a dict - ground_truth = sample.label if sample.label is not None else "" - - # Get tool call count as num_turns - num_turns = getattr(sample, "tool_call_count", 0) - - # use \\boxed{...} answer - result = math_dapo_compute_score(solution_str, ground_truth, strict_box_verify=True) - - # encourage model to call tools - if result["score"] < 0: - tool_call_reward = (num_turns - 2) / 2 * 0.1 - result["score"] = min(-0.6, result["score"] + tool_call_reward) - - if result["pred"] is None: - result["pred"] = "" - - return result diff --git a/examples/retool/tool_sandbox.py b/examples/retool/tool_sandbox.py deleted file mode 100644 index cdf68aa02..000000000 --- a/examples/retool/tool_sandbox.py +++ /dev/null @@ -1,360 +0,0 @@ -""" -Tool sandbox module for safe code execution and tool management. - -This module provides: -- PythonSandbox: Safe Python code execution environment -- ToolRegistry: Tool registration and execution management -- Memory management utilities -""" - -import asyncio -import gc -import os -import re -import subprocess -import tempfile -from contextlib import contextmanager -from typing import Any - -import psutil - -# Configuration for tool execution -TOOL_CONFIGS = { - "max_turns": 16, - "max_tool_calls": 16, - "tool_concurrency": 32, # Aggressive: 32 concurrent processes - # Python interpreter settings - "python_timeout": 120, # 2 minutes for complex calculations - "python_memory_limit": "4GB", # 4GB per Python process - "python_cpu_limit": 1, - # Memory management settings - "max_memory_usage": 12288, # 12GB total (75% of 16GB) - "cleanup_threshold": 6144, # 6GB - "aggressive_cleanup_threshold": 3072, # 3GB - "force_cleanup_threshold": 9216, # 9GB -} - -# Global semaphore for controlling concurrent tool executions -SEMAPHORE = asyncio.Semaphore(TOOL_CONFIGS["tool_concurrency"]) - - -def get_memory_usage() -> float: - """Get current memory usage in MB""" - process = psutil.Process() - return process.memory_info().rss / 1024 / 1024 - - -def cleanup_memory(): - """Force garbage collection to free memory""" - gc.collect() - - -def aggressive_cleanup_memory(): - """More aggressive memory cleanup""" - # Force multiple garbage collection cycles - for _ in range(3): - gc.collect() - - # Clear Python's internal caches - import sys - - # Note: sys.intern doesn't have a clear method, so we skip this - # Clear module cache if possible - if hasattr(sys, "modules"): - # Don't clear all modules, but clear some common ones that might cache data - modules_to_clear = ["numpy", "pandas", "matplotlib", "scipy"] - for module_name in modules_to_clear: - if module_name in sys.modules: - module = sys.modules[module_name] - if hasattr(module, "clear_cache"): - module.clear_cache() - - -def check_and_cleanup_memory(): - """Check memory usage and perform appropriate cleanup""" - current_memory = get_memory_usage() - - if current_memory > TOOL_CONFIGS["force_cleanup_threshold"]: - # Force aggressive cleanup - aggressive_cleanup_memory() - return f"Warning: High memory usage ({current_memory:.1f}MB), performed aggressive cleanup" - elif current_memory > TOOL_CONFIGS["cleanup_threshold"]: - # Normal cleanup - cleanup_memory() - return f"Info: Memory usage ({current_memory:.1f}MB), performed cleanup" - elif current_memory > TOOL_CONFIGS["aggressive_cleanup_threshold"]: - # Light cleanup - gc.collect() - return f"Info: Memory usage ({current_memory:.1f}MB), performed light cleanup" - - return None - - -class PythonSandbox: - """Python code sandbox, provides safe code execution environment""" - - def __init__(self, timeout: int = 10, memory_limit: str = "100MB"): - self.timeout = timeout - self.memory_limit = memory_limit - self.allowed_modules = { - "math", - "random", - "datetime", - "collections", - "itertools", - "functools", - "operator", - "statistics", - "decimal", - "fractions", - } - - def _check_code_safety(self, code: str) -> tuple[bool, str]: - """Check code safety by scanning for dangerous patterns""" - # Check for dangerous operations - dangerous_patterns = [ - r"import\s+os", - r"import\s+sys", - r"import\s+subprocess", - r"import\s+shutil", - r"import\s+glob", - r"import\s+pathlib", - r"__import__", - r"eval\s*\(", - r"exec\s*\(", - r"open\s*\(", - r"file\s*\(", - r"input\s*\(", - r"raw_input\s*\(", - r"compile\s*\(", - r"execfile\s*\(", - r"getattr\s*\(", - r"setattr\s*\(", - r"delattr\s*\(", - r"hasattr\s*\(", - r"globals\s*\(", - r"locals\s*\(", - r"vars\s*\(", - r"dir\s*\(", - r"type\s*\(", - r"isinstance\s*\(", - r"issubclass\s*\(", - r"super\s*\(", - r"property\s*\(", - r"staticmethod\s*\(", - r"classmethod\s*\(", - r"__\w+__", # double underscore methods - ] - - for pattern in dangerous_patterns: - if re.search(pattern, code, re.IGNORECASE): - return False, f"Code contains dangerous pattern: {pattern}" - - # Check imported modules - import_pattern = r"import\s+(\w+)" - from_pattern = r"from\s+(\w+)" - - imports = re.findall(import_pattern, code) - froms = re.findall(from_pattern, code) - - all_imports = set(imports + froms) - for imp in all_imports: - if imp not in self.allowed_modules: - return False, f"Import of '{imp}' is not allowed" - - return True, "Code is safe" - - @contextmanager - def _create_safe_environment(self): - """Create safe execution environment with temporary directory""" - # Create temporary directory - temp_dir = tempfile.mkdtemp(prefix="python_sandbox_") - - try: - # Create safe Python script - script_path = os.path.join(temp_dir, "code.py") - - # Set environment variables - env = os.environ.copy() - env["PYTHONPATH"] = temp_dir - env["PYTHONUNBUFFERED"] = "1" - - yield script_path, env, temp_dir - - finally: - # Clean up temporary directory - try: - import shutil - - shutil.rmtree(temp_dir) - except Exception: - pass - - async def execute_code(self, code: str) -> str: - """Execute Python code in sandbox with safety checks""" - # Check memory usage before execution - current_memory = get_memory_usage() - if current_memory > TOOL_CONFIGS["max_memory_usage"]: - aggressive_cleanup_memory() - return "Error: Memory usage too high, please try again" - - # Check code safety - is_safe, message = self._check_code_safety(code) - if not is_safe: - return f"Error: {message}" - - # Add necessary wrapper code with memory limits - # Properly indent the user code within the try block - # Handle indentation properly by adding 4 spaces to each line - indented_code = "\n".join(" " + line for line in code.split("\n")) - - wrapped_code = f"""import sys -import traceback -from io import StringIO -import resource - -# Set memory limit (4GB) -try: - resource.setrlimit(resource.RLIMIT_AS, (4 * 1024 * 1024 * 1024, -1)) -except Exception: - pass - -# Redirect stdout and stderr -old_stdout = sys.stdout -old_stderr = sys.stderr -stdout_capture = StringIO() -stderr_capture = StringIO() -sys.stdout = stdout_capture -sys.stderr = stderr_capture - -try: - # User code -{indented_code} - - # Get output - stdout_output = stdout_capture.getvalue() - stderr_output = stderr_capture.getvalue() - - # Restore standard output - sys.stdout = old_stdout - sys.stderr = old_stderr - - # Return result - result = "" - if stdout_output: - result += f"Output:\\n{{stdout_output}}" - if stderr_output: - result += f"\\nErrors:\\n{{stderr_output}}" - - print(result) - -except Exception as e: - # Restore standard output - sys.stdout = old_stdout - sys.stderr = old_stderr - - # Return error information - error_msg = f"Error: {{str(e)}}\\nTraceback:\\n{{traceback.format_exc()}}" - print(error_msg)""" - - with self._create_safe_environment() as (script_path, env, temp_dir): - # Write code to file - with open(script_path, "w") as f: - f.write(wrapped_code) - - try: - # Use subprocess to run code - process = subprocess.Popen( - ["python3", script_path], - stdout=subprocess.PIPE, - stderr=subprocess.PIPE, - env=env, - cwd=temp_dir, - text=True, - ) - - # Set timeout - try: - stdout, stderr = process.communicate(timeout=self.timeout) - - if process.returncode == 0: - result = stdout.strip() - else: - result = f"Error: Process exited with code {process.returncode}\n{stderr}" - - except subprocess.TimeoutExpired: - process.kill() - result = f"Error: Code execution timed out after {self.timeout} seconds" - - except Exception as e: - result = f"Error: Failed to execute code: {str(e)}" - - # Check memory usage after execution and cleanup if needed - cleanup_message = check_and_cleanup_memory() - if cleanup_message: - print(f"Memory cleanup: {cleanup_message}") - - return result - - -class ToolRegistry: - """Tool registry, manages available tools and their execution""" - - def __init__(self): - self.tools = {} - self.python_sandbox = PythonSandbox( - timeout=TOOL_CONFIGS["python_timeout"], memory_limit=TOOL_CONFIGS["python_memory_limit"] - ) - self._register_default_tools() - - def _register_default_tools(self): - """Register default tools in the registry""" - # Python code interpreter - self.register_tool( - "code_interpreter", - { - "type": "function", - "function": { - "name": "code_interpreter", - "description": "A tool for executing Python code in a safe sandbox environment.", - "parameters": { - "type": "object", - "properties": {"code": {"type": "string", "description": "The Python code to execute"}}, - "required": ["code"], - }, - }, - }, - ) - - def register_tool(self, name: str, tool_spec: dict[str, Any]): - """Register a new tool in the registry""" - self.tools[name] = tool_spec - - def get_tool_specs(self) -> list[dict[str, Any]]: - """Get all tool specifications as a list""" - return list(self.tools.values()) - - async def execute_tool(self, tool_name: str, arguments: dict[str, Any]) -> str: - """Execute a tool call with the given arguments""" - if tool_name not in self.tools: - return f"Error: Tool '{tool_name}' not found" - - async with SEMAPHORE: - if tool_name == "code_interpreter": - return await self._execute_python(arguments) - else: - return f"Error: Tool '{tool_name}' not implemented" - - async def _execute_python(self, arguments: dict[str, Any]) -> str: - """Execute Python code using the sandbox""" - code = arguments.get("code", "") - if not code.strip(): - return "Error: No code provided" - - # Execute code in sandbox - result = await self.python_sandbox.execute_code(code) - return result - - -# Global tool registry instance -tool_registry = ToolRegistry()