create_tool.py
1 import json 2 import re 3 4 5 _NAME_RE = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") 6 _MAX_CODE_SIZE = 10240 # 10KB 7 8 9 def create_tool(name: str, description: str, parameters: str, code: str, **kwargs) -> str: 10 """Create a reusable tool for this project that runs in a sandboxed Docker container. 11 The tool will be available in this and all future conversations for this project. 12 13 Args: 14 name (str): Tool name (letters, digits, underscore only, e.g. "fetch_page_title"). 15 description (str): What the tool does (shown to the LLM when deciding which tool to use). 16 parameters (str): JSON schema string describing the tool's input parameters (e.g. '{"type":"object","properties":{"url":{"type":"string"}},"required":["url"]}'). 17 code (str): Python code that receives an `args` dict with the parameters and should print the result to stdout. 18 """ 19 brain = kwargs.get("_brain") 20 chat_id = kwargs.get("_chat_id") 21 project_id = kwargs.get("_project_id") 22 23 if not brain or not getattr(brain, "docker_manager", None): 24 return "ERROR: Docker is not configured. An admin must configure Docker in Settings to create tools." 25 26 if not project_id: 27 return "ERROR: No project context available." 28 29 # Validate name 30 if not name or not _NAME_RE.match(name): 31 return f"ERROR: Invalid tool name '{name}'. Use only letters, digits, and underscores. Must start with a letter or underscore." 32 33 if not description or not description.strip(): 34 return "ERROR: Description is required." 35 36 if not code or not code.strip(): 37 return "ERROR: Code is required." 38 39 if len(code) > _MAX_CODE_SIZE: 40 return f"ERROR: Code is too large ({len(code)} bytes). Maximum is {_MAX_CODE_SIZE} bytes." 41 42 # Validate parameters JSON 43 try: 44 params_dict = json.loads(parameters) if isinstance(parameters, str) else parameters 45 except json.JSONDecodeError as e: 46 return f"ERROR: Invalid parameters JSON: {e}" 47 48 # Test the code by running it in Docker with empty args 49 script = f"import json, sys\nargs = json.loads(sys.stdin.readline() or '{{}}')\n{code}" 50 test_result = brain.docker_manager.run_script( 51 chat_id or "ephemeral", 52 script, 53 stdin_data="{}", 54 ) 55 if test_result.startswith("ERROR:"): 56 return f"ERROR: Code validation failed — {test_result}" 57 58 # Save to database 59 from restai.database import DBWrapper 60 db = DBWrapper() 61 try: 62 db.upsert_project_tool( 63 project_id=project_id, 64 name=name, 65 description=description, 66 parameters=json.dumps(params_dict) if isinstance(params_dict, dict) else parameters, 67 code=code, 68 ) 69 finally: 70 db.db.close() 71 72 return f"Tool '{name}' created successfully. You can now call it by name."