open-webui/backend/apps/webui/utils.py

46 lines
1.5 KiB
Python
Raw Normal View History

2024-06-11 06:40:27 +00:00
from importlib import util
import os
2024-06-20 07:37:02 +00:00
from config import TOOLS_DIR, FUNCTIONS_DIR
2024-06-11 06:40:27 +00:00
def load_toolkit_module_by_id(toolkit_id):
toolkit_path = os.path.join(TOOLS_DIR, f"{toolkit_id}.py")
spec = util.spec_from_file_location(toolkit_id, toolkit_path)
module = util.module_from_spec(spec)
try:
spec.loader.exec_module(module)
print(f"Loaded module: {module.__name__}")
if hasattr(module, "Tools"):
return module.Tools()
else:
raise Exception("No Tools class found")
except Exception as e:
print(f"Error loading module: {toolkit_id}")
# Move the file to the error folder
os.rename(toolkit_path, f"{toolkit_path}.error")
raise e
2024-06-20 07:37:02 +00:00
def load_function_module_by_id(function_id):
function_path = os.path.join(FUNCTIONS_DIR, f"{function_id}.py")
spec = util.spec_from_file_location(function_id, function_path)
module = util.module_from_spec(spec)
try:
spec.loader.exec_module(module)
print(f"Loaded module: {module.__name__}")
if hasattr(module, "Pipe"):
2024-06-20 07:54:58 +00:00
return module.Pipe(), "pipe"
2024-06-20 07:37:02 +00:00
elif hasattr(module, "Filter"):
2024-06-20 07:54:58 +00:00
return module.Filter(), "filter"
2024-06-20 07:37:02 +00:00
else:
raise Exception("No Function class found")
except Exception as e:
print(f"Error loading module: {function_id}")
# Move the file to the error folder
os.rename(function_path, f"{function_path}.error")
raise e