src/manager/tool_manager.py

import importlib
import importlib.util
import os
import types
from typing import List
import pip
from google.genai import types

from src.manager.budget_manager import BudgetManager
from src.manager.utils.singleton import singleton
from src.manager.utils.suppress_outputs import suppress_output
from src.tools.default_tools.tool_deletor import ToolDeletor
from src.manager.utils.streamlit_interface import output_assistant_response

toolsImported = []

TOOLS_DIRECTORIES = [os.path.abspath("./src/tools/default_tools"), os.path.abspath("./src/tools/user_tools")]

installed_packages = set()

class Tool:
    def __init__(self, toolClass):
        suppress_output(self.load_tool)(toolClass)
        
    def load_tool(self, toolClass):
        self.tool = toolClass()
        self.inputSchema = self.tool.inputSchema
        self.name = self.inputSchema["name"]
        self.description = self.inputSchema["description"]
        self.dependencies = self.tool.dependencies
        self.create_resource_cost = self.inputSchema.get("create_resource_cost", 0)
        self.invoke_resource_cost = self.inputSchema.get("invoke_resource_cost", 0)
        self.create_expense_cost = self.inputSchema.get("create_expense_cost", 0)
        self.invoke_expense_cost = self.inputSchema.get("invoke_expense_cost", 0)
        if self.dependencies:
            self.install_dependencies()
    
    def install_dependencies(self):
        for package in self.dependencies:
            if package in installed_packages:
                continue
            try:
                __import__(package.split('==')[0])
            except ImportError:
                print(f"Installing {package}")
                if '==' in package:
                    package = package.split('==')[0]
                pip.main(['install', package])
            installed_packages.add(package)

    def run(self, query):
        return self.tool.run(**query)

@singleton
class ToolManager:
    toolsImported: List[Tool] = []
    budget_manager: BudgetManager = BudgetManager()
    is_creation_enabled: bool = True
    is_invocation_enabled: bool = True

    def __init__(self):
        self.load_tools()
        self._output_budgets()
    
    def set_creation_mode(self, status: bool):
        self.is_creation_enabled = status
        if status:
            output_assistant_response("Tool creation mode is enabled.")
        else:
            output_assistant_response("Tool creation mode is disabled.")
    
    def set_invocation_mode(self, status: bool):
        self.is_invocation_enabled = status
        if status:
            output_assistant_response("Tool invocation mode is enabled.")
        else:
            output_assistant_response("Tool invocation mode is disabled.")
    
    def _output_budgets(self):
        output_assistant_response(f"Resource budget Remaining: {self.budget_manager.get_current_remaining_resource_budget()}")
        output_assistant_response(f"Expense budget Remaining: {self.budget_manager.get_current_remaining_expense_budget()}")

    def load_tools(self):
        newToolsImported = []
        for directory in TOOLS_DIRECTORIES:
            for filename in os.listdir(directory):
                if filename.endswith(".py") and filename != "__init__.py":
                    module_name = filename[:-3]
                    spec = importlib.util.spec_from_file_location(module_name, f"{directory}/{filename}")
                    foo = importlib.util.module_from_spec(spec)
                    spec.loader.exec_module(foo)
                    class_name = foo.__all__[0]
                    toolClass = getattr(foo, class_name)
                    toolObj = Tool(toolClass)
                    newToolsImported.append(toolObj)
                    if toolObj.create_resource_cost is not None:
                        self.budget_manager.add_to_resource_budget(toolObj.create_resource_cost)
                    if toolObj.create_expense_cost is not None:
                        self.budget_manager.add_to_resource_budget(toolObj.create_expense_cost)
        self.toolsImported = newToolsImported

    def runTool(self, toolName, query):
        if not self.is_invocation_enabled:
            raise Exception("Tool invocation mode is disabled")
        if toolName == "ToolCreator":
            if not self.is_creation_enabled:
                raise Exception("Tool creation mode is disabled")
        self._output_budgets()
        for tool in self.toolsImported:
            if tool.name == toolName:
                if tool.invoke_resource_cost is not None:
                    if not self.budget_manager.can_spend_resource(tool.invoke_resource_cost):
                        raise Exception("No resource budget remaining")
                if tool.invoke_expense_cost is not None:
                    self.budget_manager.add_to_resource_budget(tool.invoke_expense_cost)
                return tool.run(query)
        self._output_budgets()
        return {
            "status": "error",
            "message": f"Tool {toolName} not found",
            "output": None
        }
    

    def getTools(self):
        toolsList = []
        for tool in self.toolsImported:
            parameters = types.Schema()
            parameters.type = tool.inputSchema["parameters"]["type"]
            properties = {}
            for prop, value in tool.inputSchema["parameters"]["properties"].items():
                properties[prop] = types.Schema(
                    type=value["type"],
                    description=value["description"]
                )
            parameters.properties = properties
            parameters.required = tool.inputSchema["parameters"].get("required", [])
            function = types.FunctionDeclaration(
                name=tool.inputSchema["name"],
                description=tool.inputSchema["description"],
                parameters=parameters,
            )
            toolType = types.Tool(function_declarations=[function])
            toolsList.append(toolType)
        return toolsList
    
    def delete_tool(self, toolName, toolFile):
        try:
            tool_deletor = ToolDeletor()
            tool_deletor.run(name=toolName, file_path=toolFile)
            for tool in self.toolsImported:
                if tool.name == toolName:
                    # remove budget for the tool
                    if tool.create_resource_cost is not None:
                        self.budget_manager.remove_from_resource_expense(tool.create_resource_cost)
                    if tool.create_expense_cost is not None:
                        self.budget_manager.remove_from_resource_expense(tool.create_expense_cost)
                    self.toolsImported.remove(tool)
                    return {
                        "status": "success",
                        "message": f"Tool {toolName} deleted",
                        "output": None
                    }
        except Exception as e:
            return {
                "status": "error",
                "message": f"Tool {toolName} not found",
                "output": None
            }