Spaces:
Running
Running
File size: 6,217 Bytes
576227b 6900003 576227b 2f85c93 0e58feb 2f85c93 576227b 0e58feb 576227b bf722a2 576227b 98d31e7 576227b bf722a2 6900003 bf722a2 576227b bf722a2 576227b 6900003 576227b bf722a2 576227b a3a158e bf722a2 576227b 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 |
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()
def __init__(self):
self.load_tools()
self._output_budgets()
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_resource_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):
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
} |