toolkit.py

from collections import OrderedDict
from langchain_core.tools import tool
from typing import Any,Annotated
from typing import Annotated, Optional, List, Literal
from langchain.tools.base import BaseTool
from utils import generate_function_docstring
from virtual_machine import VM
import os
import traceback

TOOL_BUNDLE_REGISTRY = OrderedDict()

class Toolkit:
   def __init__(self, tool_bundles=[], **kwargs):
      self.registry = OrderedDict()
      self.env_registry = OrderedDict()
      for k in tool_bundles:
          bundle, env = TOOL_BUNDLE_REGISTRY[k](**kwargs) 
          self.env_registry[k] = env 
          for t in bundle: 
            self.add_tool(t)

   def add_tool(self, tool):
      self.registry[tool.name] = tool

   def get_prompt_description(self):
      tool_prompt = ""
      for k in self.registry:
         tool = self.registry[k]
         docstring = generate_function_docstring(tool.name,
                                                 tool.description,
                                                 tool.args)
         tool_prompt += f'\n{docstring}\n'
      return tool_prompt
   
   def execute_action(self, action_name, action_params):
      if(action_name not in self.registry):
          return "Action not registered."
      try:
          result = self.registry[action_name].invoke(action_params)
          return result
      except Exception as e:
          print(traceback.format_exc())
          return f"Action execution error: {e}"
      

import asyncio
def get_bash_tool(vm: VM, **kwargs):
    @tool(description=(
        "Executes a bash command in a sandboxed local environment. "
        "\nDo NOT use interactive commands including nano, vim, vi, and others. "
        "\nSet timeout=0 to run a command in the background."
        "\nNOTE: NEVER USE 'cat' command to read a whole file into the context."
        "\nNOTE: For background processes, you should check the progress by reading the log file but not use 'ps' and 'grep'."
        "\n!!! NEVER USE 'tail -f' because it is interactive. Simply use 'tail' instead."
        "\n!!! NEVER RUN 'tail', 'head', 'cat', 'ls', 'cd' or similar commands as a background process. Not necessary."
    ))
    def bash(
        command: Annotated[str, "The bash command to run."],
        timeout: Annotated[int, "Max seconds to wait; 0 runs in background."] = 25
    ) -> str:
        return vm.run_command(command)
    return [bash], vm


def get_done_tool(**kwargs):
    @tool(description="Claim the task is finished.")
    def done(text: Annotated[str, "Information provided to the user when the task is finished."]):
        return text
    return [done], None
    
### register the tools      
TOOL_BUNDLE_REGISTRY['virtual_bash'] = get_bash_tool
TOOL_BUNDLE_REGISTRY['done'] = get_done_tool