File size: 5,369 Bytes
576227b 6900003 576227b 2f85c93 0e58feb 2f85c93 576227b 0e58feb 576227b 98d31e7 576227b 6900003 576227b 6900003 576227b 967c695 434b328 576227b a3a158e 6900003 576227b 434b328 576227b 6900003 576227b 434b328 576227b 967c695 576227b 967c695 576227b 6900003 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 |
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")]
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_cost = None
self.invoke_cost = None
if "create_cost" in self.tool.inputSchema:
self.create_cost = self.tool.inputSchema["create_cost"]
if "invoke_cost" in self.tool.inputSchema:
self.invoke_cost = self.tool.inputSchema["invoke_cost"]
if self.dependencies:
self.install_dependencies()
def install_dependencies(self):
for package in self.dependencies:
try:
__import__(package.split('==')[0])
except ImportError:
print(f"Installing {package}")
if '==' in package:
package = package.split('==')[0]
pip.main(['install', 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()
output_assistant_response(f"Budget Remaining: {self.budget_manager.get_current_remaining_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_cost is not None:
self.budget_manager.add_to_expense(toolObj.create_cost)
self.toolsImported = newToolsImported
def runTool(self, toolName, query):
output_assistant_response(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
for tool in self.toolsImported:
if tool.name == toolName:
if tool.invoke_cost is not None:
self.budget_manager.add_to_expense(tool.invoke_cost)
return tool.run(query)
output_assistant_response(f"Budget Remaining: {self.budget_manager.get_current_remaining_budget()}")
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_cost is not None:
self.budget_manager.remove_from_expense(tool.create_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
} |