This commit is contained in:
Timothy Jaeryang Baek 2025-04-05 05:03:15 -06:00
parent 66db2e1515
commit 93bb77ede3
3 changed files with 19 additions and 19 deletions

View File

@ -226,7 +226,7 @@ async def chat_completion_tools_handler(
if isinstance(tool_result, str): if isinstance(tool_result, str):
tool = tools[tool_function_name] tool = tools[tool_function_name]
tool_id = tool.get("toolkit_id", "") tool_id = tool.get("tool_id", "")
if tool.get("metadata", {}).get("citation", False) or tool.get( if tool.get("metadata", {}).get("citation", False) or tool.get(
"direct", False "direct", False
): ):

View File

@ -68,23 +68,23 @@ def replace_imports(content):
return content return content
def load_tools_module_by_id(toolkit_id, content=None): def load_tools_module_by_id(tool_id, content=None):
if content is None: if content is None:
tool = Tools.get_tool_by_id(toolkit_id) tool = Tools.get_tool_by_id(tool_id)
if not tool: if not tool:
raise Exception(f"Toolkit not found: {toolkit_id}") raise Exception(f"Toolkit not found: {tool_id}")
content = tool.content content = tool.content
content = replace_imports(content) content = replace_imports(content)
Tools.update_tool_by_id(toolkit_id, {"content": content}) Tools.update_tool_by_id(tool_id, {"content": content})
else: else:
frontmatter = extract_frontmatter(content) frontmatter = extract_frontmatter(content)
# Install required packages found within the frontmatter # Install required packages found within the frontmatter
install_frontmatter_requirements(frontmatter.get("requirements", "")) install_frontmatter_requirements(frontmatter.get("requirements", ""))
module_name = f"tool_{toolkit_id}" module_name = f"tool_{tool_id}"
module = types.ModuleType(module_name) module = types.ModuleType(module_name)
sys.modules[module_name] = module sys.modules[module_name] = module
@ -108,7 +108,7 @@ def load_tools_module_by_id(toolkit_id, content=None):
else: else:
raise Exception("No Tools class found in the module") raise Exception("No Tools class found in the module")
except Exception as e: except Exception as e:
log.error(f"Error loading module: {toolkit_id}: {e}") log.error(f"Error loading module: {tool_id}: {e}")
del sys.modules[module_name] # Clean up del sys.modules[module_name] # Clean up
raise e raise e
finally: finally:

View File

@ -40,15 +40,14 @@ def apply_extra_params_to_tool_function(
return new_function return new_function
# Mutation on extra_params
def get_tools( def get_tools(
request: Request, tool_ids: list[str], user: UserModel, extra_params: dict request: Request, tool_ids: list[str], user: UserModel, extra_params: dict
) -> dict[str, dict]: ) -> dict[str, dict]:
tools_dict = {} tools_dict = {}
for tool_id in tool_ids: for tool_id in tool_ids:
tools = Tools.get_tool_by_id(tool_id) tool = Tools.get_tool_by_id(tool_id)
if tools 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])
@ -57,7 +56,7 @@ def get_tools(
tool_dict = { tool_dict = {
"spec": spec, "spec": spec,
"callable": callable, "callable": callable,
"toolkit_id": tool_id, "tool_id": tool_id,
# Misc info # Misc info
"metadata": { "metadata": {
"file_handler": hasattr(module, "file_handler") "file_handler": hasattr(module, "file_handler")
@ -65,9 +64,9 @@ def get_tools(
"citation": hasattr(module, "citation") and module.citation, "citation": hasattr(module, "citation") and module.citation,
}, },
} }
else: else:
continue
else:
module = request.app.state.TOOLS.get(tool_id, None) module = request.app.state.TOOLS.get(tool_id, None)
if module is None: if module is None:
module, _ = load_tools_module_by_id(tool_id) module, _ = load_tools_module_by_id(tool_id)
@ -83,7 +82,7 @@ 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 tools.specs: for spec in tool.specs:
# TODO: Fix hack for OpenAI API # TODO: Fix hack for OpenAI API
# Some times breaks OpenAI but others don't. Leaving the comment # Some times breaks OpenAI but others don't. Leaving the comment
for val in spec.get("parameters", {}).get("properties", {}).values(): for val in spec.get("parameters", {}).get("properties", {}).values():
@ -114,7 +113,7 @@ def get_tools(
tool_dict = { tool_dict = {
"spec": spec, "spec": spec,
"callable": callable, "callable": callable,
"toolkit_id": tool_id, "tool_id": tool_id,
# Misc info # Misc info
"metadata": { "metadata": {
"file_handler": hasattr(module, "file_handler") "file_handler": hasattr(module, "file_handler")
@ -128,8 +127,8 @@ 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 {tools} and {tool_id}.") log.warning(f"Collision between {tool} and {tool_id}.")
log.warning(f"Discarding {tools}.{function_name}") log.warning(f"Discarding {tool}.{function_name}")
else: else:
tools_dict[function_name] = tool_dict tools_dict[function_name] = tool_dict
@ -239,8 +238,9 @@ def get_callable_attributes(tool: object) -> list[Callable]:
def get_tools_specs(tool_class: object) -> list[dict]: def get_tools_specs(tool_class: object) -> list[dict]:
function_list = get_callable_attributes(tool_class) function_model_list = map(
function_model_list = map(function_to_pydantic_model, function_list) function_to_pydantic_model, get_callable_attributes(tool_class)
)
return [ return [
convert_to_openai_function(function_model) convert_to_openai_function(function_model)
for function_model in function_model_list for function_model in function_model_list