allow arbitrary python as tools

Author: lolololol

core/hatchery.py

@@ -37,10 +37,12 @@ def run(self,ctx):
     return self.tool_func(**fixed_args)
 
 class Drone(Agent):
-  def __init__(self,node_name,sys_prompt,usr_prompt,_tools=[],next=None,model=None,base_url=None,parent_hatchery=None,mcps=[]):
+  def __init__(self,node_name,sys_prompt,usr_prompt,_tools=[],next=None,model=None,base_url=None,parent_hatchery=None,mcps=[],pytools=[]):
     print("drone: initializing drone '%s'" % node_name)
     self.mcp_loader = core.mcp.MCPLoader()
     self.toolbox = tools.ToolLoader(Hatchery)
+    for p in pytools:
+      self.toolbox.load_pytool(p)
     self.parent_hatchery = parent_hatchery # allows cross-node calls
     self.name = node_name
     self.usr_prompt = usr_prompt
@@ -109,8 +111,9 @@ def __init__(self, fn):
         sys_prompt = node.get("sys_prompt","You are a helpful assistant.")
         usr_prompt = node["usr_prompt"]
         tools  = node.get("tools",[])
+        pytools = node.get("pytools",[])
         mcps  = node.get("mcp",[])
-        self.nodes[node_name] = Drone(node_name,sys_prompt,usr_prompt,tools,next=node.get("next",None),model=node_model,base_url=node_base_url,parent_hatchery=self,mcps = mcps)
+        self.nodes[node_name] = Drone(node_name,sys_prompt,usr_prompt,tools,next=node.get("next",None),model=node_model,base_url=node_base_url,parent_hatchery=self,mcps = mcps,pytools=pytools)
         self.nodes[node_name].save_output = node.get("save_output",None)
         self.nodes[node_name].write_output = node.get("write_output",None)
       elif node_type == "tool":

harness.py

@@ -61,7 +61,7 @@ def main():
   CFG_REASONING = None
   CFG_HATCHERY = None
   DENIED_ALREADY = False
-  args,extra = getopt.getopt(sys.argv[1:],"ip:s:t:m:r:a:",["interactive","prompt=","system=","tool=","model=","reasoning=","persona=","toolbox=","agentic=","mcp=","mcp-deny="])
+  args,extra = getopt.getopt(sys.argv[1:],"ip:s:t:m:r:a:",["interactive","prompt=","system=","tool=","model=","reasoning=","persona=","toolbox=","agentic=","mcp=","mcp-deny=","pytool="])
   for arg,val in args:
     if arg in ["-p","--prompt"]:
       if CFG_USR_PROMPT is not None:
@@ -91,6 +91,8 @@ def main():
       CFG_TOOLS.append(val)
     elif arg == "--toolbox":
       CFG_TOOLS += Toolbox.fetch_toolbox(val)
+    elif arg == "--pytool":
+      CFG_TOOLS += Toolbox.load_pytool(val)
     elif arg in ["-m","--model"]:
       CFG_MODEL = val
     elif arg == "--persona":

hatchery/pytool.json

@@ -0,0 +1,13 @@
+{
+  "name":"pytool",
+  "desc":"test the pytool system in the context of hatchery",
+  "start":"FruitStoryBananas",
+  "nodes":[
+  {
+    "name":"FruitStoryBananas",
+    "usr_prompt":"Use test_func1 to get the flag",
+    "pytools":["/var/tmp/pytool.py"],
+    "tools":["test_func1"]
+  }
+  ]
+}

tools/__init__.py

@@ -7,13 +7,17 @@
 import tools.chrome
 import tools.debug
 import tools.r2tool
+import inspect
+
+class Namespace:
+  pass
 
 class ToolLoader:
   def registerHatchery(self,hatch_name):
     temp_h = self.HatcheryClass(hatch_name)
     self.hatch_store[temp_h.name] = temp_h
     h_name = temp_h.name # hatch_name is the file, h_name is the internal name
-    print("info: registering hatch '%s'" % h_name)
+    print("tools: registering hatch '%s'" % h_name)
     self.hatch_names[hatch_name] = h_name
     if h_name in self.tools.keys():
       print("warning: ToolLoader trying to register '%s' as hatchery, already loaded" % name)
@@ -30,7 +34,7 @@ def registerFunction(self,name,function,func_desc):
       self.tools[name].__doc__ = func_desc
 
   def run_hatch(self,hatch_name,input_str):
-    print("info: ToolLoader calling run_hatch(%s,%s)" % (hatch_name,input_str))
+    print("tools: ToolLoader calling run_hatch(%s,%s)" % (hatch_name,input_str))
     return self.hatch_store[hatch_name].run(ctx={"input":input_str})
     # return "ok"
 
@@ -39,6 +43,7 @@ def __init__(self,HatcheryClass=None):
       raise ValueError("fatal: i didn't get hatcheryclass")
     else:
       self.HatcheryClass = HatcheryClass
+    self.exec_ns_array = []
     self.hatch_store = {}
     self.hatch_names = {}
     self.short_name_store = {}
@@ -65,23 +70,39 @@ def __init__(self,HatcheryClass=None):
     self.registerFunction("r2_cmd",r2tool.r2_cmd, "Run an r2 command.")
     self.registerFunction("r2_close",r2tool.r2_close, "Close the r2 session.")
 
+  def load_pytool(self,name):
+    print("tools: loading pytool '%s'. caveat emptor..." % name)
+    namespace = {}
+    with open(name,"r") as f:
+      code = f.read()
+    ns = Namespace()
+    exec(code,ns.__dict__)
+    out = []
+    self.exec_ns_array.append(ns)
+    for i in ns.__dict__.keys():
+      if callable(ns.__dict__[i]):
+        print("tools: got a callable '%s'" % i)
+        self.registerFunction(i,ns.__dict__[i],ns.__dict__[i].__doc__ or "")
+        out.append(i)
+    return out
+
   def fetch_toolbox(self,name):
     out = []
     for i in self.tools.keys():
       if i.startswith(name):
         out.append(i)
-    print("info: fetch_toolbox(%s) returned %d results" % (name,len(out)))
+    print("tools: fetch_toolbox(%s) returned %d results" % (name,len(out)))
     return out
 
   def fetch(self,name):
-    print("info: attempting to grab tool '%s'" % name)
+    print("tools: attempting to grab tool '%s'" % name)
     if name.startswith("hatch:"):
       shortname = name[6:]
       if shortname in self.hatch_names.keys():
         # translate the name first.
         return self.tools[self.hatch_names[shortname]]
       else:
-        print("info: detected hatchery-as-tool pattern, passing")
+        print("tools: detected hatchery-as-tool pattern, passing")
         self.registerHatchery(shortname)
         return self.tools[self.hatch_names[shortname]]
     else: