File size: 7,100 Bytes
576227b 6900003 576227b 2f85c93 0e58feb 2f85c93 576227b 0e58feb 576227b bf722a2 576227b 98d31e7 576227b bf722a2 6900003 bf722a2 576227b bf722a2 576227b 6900003 231c049 576227b bf722a2 231c049 bf722a2 231c049 576227b a3a158e bf722a2 576227b 231c049 bf722a2 576227b bf722a2 576227b bf722a2 576227b 967c695 576227b 967c695 576227b 6900003 bf722a2 576227b 6900003 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 |
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
} |