/ restai / llms / tools / create_tool.py
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."