refac: tools

This commit is contained in:
Timothy Jaeryang Baek 2024-11-16 18:31:13 -08:00
parent 393ab5d457
commit b562067e91

View File

@ -32,10 +32,11 @@ def apply_extra_params_to_tool_function(
def get_tools( def get_tools(
webui_app, tool_ids: list[str], user: UserModel, extra_params: dict webui_app, tool_ids: list[str], user: UserModel, extra_params: dict
) -> dict[str, dict]: ) -> dict[str, dict]:
tools = {} tools_dict = {}
for tool_id in tool_ids: for tool_id in tool_ids:
toolkit = Tools.get_tool_by_id(tool_id) tools = Tools.get_tool_by_id(tool_id)
if toolkit is None: if tools is None:
continue continue
module = webui_app.state.TOOLS.get(tool_id, None) module = webui_app.state.TOOLS.get(tool_id, None)
@ -53,11 +54,19 @@ def get_tools(
**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)
) )
for spec in toolkit.specs: for spec in tools.specs:
# TODO: Fix hack for OpenAI API # TODO: Fix hack for OpenAI API
for val in spec.get("parameters", {}).get("properties", {}).values(): for val in spec.get("parameters", {}).get("properties", {}).values():
if val["type"] == "str": if val["type"] == "str":
val["type"] = "string" val["type"] = "string"
# Remove internal parameters
spec["parameters"]["properties"] = {
key: val
for key, val in spec["parameters"]["properties"].items()
if not key.startswith("__")
}
function_name = spec["name"] 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
@ -77,13 +86,14 @@ def get_tools(
} }
# TODO: if collision, prepend toolkit name # TODO: if collision, prepend toolkit name
if function_name in tools: if function_name in tools_dict:
log.warning(f"Tool {function_name} already exists in another toolkit!") log.warning(f"Tool {function_name} already exists in another tools!")
log.warning(f"Collision between {toolkit} and {tool_id}.") log.warning(f"Collision between {tools} and {tool_id}.")
log.warning(f"Discarding {toolkit}.{function_name}") log.warning(f"Discarding {tools}.{function_name}")
else: else:
tools[function_name] = tool_dict tools_dict[function_name] = tool_dict
return tools
return tools_dict
def doc_to_dict(docstring): def doc_to_dict(docstring):