refac: substandard codebase overhauled

This commit is contained in:
Timothy Jaeryang Baek 2025-04-05 05:31:22 -06:00
parent 93bb77ede3
commit e570a98bf7

View File

@ -23,21 +23,23 @@ import copy
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
def apply_extra_params_to_tool_function( def get_async_tool_function_and_apply_extra_params(
function: Callable, extra_params: dict function: Callable, extra_params: dict
) -> Callable[..., Awaitable]: ) -> Callable[..., Awaitable]:
sig = inspect.signature(function) sig = inspect.signature(function)
extra_params = {k: v for k, v in extra_params.items() if k in sig.parameters} extra_params = {k: v for k, v in extra_params.items() if k in sig.parameters}
partial_func = partial(function, **extra_params) partial_func = partial(function, **extra_params)
if inspect.iscoroutinefunction(function): if inspect.iscoroutinefunction(function):
update_wrapper(partial_func, function) update_wrapper(partial_func, function)
return partial_func return partial_func
else:
# Make it a coroutine function
async def new_function(*args, **kwargs):
return partial_func(*args, **kwargs)
async def new_function(*args, **kwargs): update_wrapper(new_function, function)
return partial_func(*args, **kwargs) return new_function
update_wrapper(new_function, function)
return new_function
def get_tools( def get_tools(
@ -48,22 +50,49 @@ def get_tools(
for tool_id in tool_ids: for tool_id in tool_ids:
tool = Tools.get_tool_by_id(tool_id) tool = Tools.get_tool_by_id(tool_id)
if tool is None: if tool is None:
if tool_id.startswith("server:"): if tool_id.startswith("server:"):
server_idx = int(tool_id.split(":")[1]) server_idx = int(tool_id.split(":")[1])
tool_server_connection = (
request.app.state.config.TOOL_SERVER_CONNECTIONS[server_idx]
)
tool_server_data = request.app.state.TOOL_SERVERS[server_idx] tool_server_data = request.app.state.TOOL_SERVERS[server_idx]
specs = tool_server_data.get("specs", [])
tool_dict = { for spec in specs:
"spec": spec, function_name = spec["name"]
"callable": callable,
"tool_id": tool_id, auth_type = tool_server_connection.get("auth_type", "bearer")
# Misc info token = None
"metadata": {
"file_handler": hasattr(module, "file_handler") if auth_type == "bearer":
and module.file_handler, token = tool_server_connection.get("key", "")
"citation": hasattr(module, "citation") and module.citation, elif auth_type == "session":
}, token = request.state.token.credentials
}
callable = get_async_tool_function_and_apply_extra_params(
execute_tool_server,
{
"token": token,
"url": tool_server_data["url"],
"name": function_name,
"server_data": tool_server_data,
},
)
tool_dict = {
"tool_id": tool_id,
"callable": callable,
"spec": spec,
}
# TODO: if collision, prepend toolkit name
if function_name in tools_dict:
log.warning(
f"Tool {function_name} already exists in another tools!"
)
log.warning(f"Discarding {tool_id}.{function_name}")
else:
tools_dict[function_name] = tool_dict
else: else:
continue continue
else: else:
@ -73,10 +102,11 @@ def get_tools(
request.app.state.TOOLS[tool_id] = module request.app.state.TOOLS[tool_id] = module
extra_params["__id__"] = tool_id extra_params["__id__"] = tool_id
# Set valves for the tool
if hasattr(module, "valves") and hasattr(module, "Valves"): if hasattr(module, "valves") and hasattr(module, "Valves"):
valves = Tools.get_tool_valves_by_id(tool_id) or {} valves = Tools.get_tool_valves_by_id(tool_id) or {}
module.valves = module.Valves(**valves) module.valves = module.Valves(**valves)
if hasattr(module, "UserValves"): if hasattr(module, "UserValves"):
extra_params["__user__"]["valves"] = module.UserValves( # type: ignore extra_params["__user__"]["valves"] = module.UserValves( # type: ignore
**Tools.get_user_valves_by_id_and_user_id(tool_id, user.id) **Tools.get_user_valves_by_id_and_user_id(tool_id, user.id)
@ -89,21 +119,21 @@ def get_tools(
if val["type"] == "str": if val["type"] == "str":
val["type"] = "string" val["type"] = "string"
# Remove internal parameters # Remove internal reserved parameters (e.g. __id__, __user__)
spec["parameters"]["properties"] = { spec["parameters"]["properties"] = {
key: val key: val
for key, val in spec["parameters"]["properties"].items() for key, val in spec["parameters"]["properties"].items()
if not key.startswith("__") if not key.startswith("__")
} }
function_name = spec["name"]
# convert to function that takes only model params and inserts custom params # convert to function that takes only model params and inserts custom params
original_func = getattr(module, function_name) function_name = spec["name"]
callable = apply_extra_params_to_tool_function( tool_function = getattr(module, function_name)
original_func, extra_params callable = get_async_tool_function_and_apply_extra_params(
tool_function, extra_params
) )
# TODO: Support Pydantic models as parameters
if callable.__doc__ and callable.__doc__.strip() != "": if callable.__doc__ and callable.__doc__.strip() != "":
s = re.split(":(param|return)", callable.__doc__, 1) s = re.split(":(param|return)", callable.__doc__, 1)
spec["description"] = s[0] spec["description"] = s[0]
@ -111,9 +141,9 @@ def get_tools(
spec["description"] = function_name spec["description"] = function_name
tool_dict = { tool_dict = {
"spec": spec,
"callable": callable,
"tool_id": tool_id, "tool_id": tool_id,
"callable": callable,
"spec": spec,
# Misc info # Misc info
"metadata": { "metadata": {
"file_handler": hasattr(module, "file_handler") "file_handler": hasattr(module, "file_handler")
@ -127,8 +157,7 @@ def get_tools(
log.warning( log.warning(
f"Tool {function_name} already exists in another tools!" f"Tool {function_name} already exists in another tools!"
) )
log.warning(f"Collision between {tool} and {tool_id}.") log.warning(f"Discarding {tool_id}.{function_name}")
log.warning(f"Discarding {tool}.{function_name}")
else: else:
tools_dict[function_name] = tool_dict tools_dict[function_name] = tool_dict