File size: 4,221 Bytes
a733dcb | 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 | import os
import re
import subprocess
import sys
from importlib import util
from apps.webui.models.functions import Functions
from apps.webui.models.tools import Tools
from config import FUNCTIONS_DIR, TOOLS_DIR
def extract_frontmatter(file_path):
"""
Extract frontmatter as a dictionary from the specified file path.
"""
frontmatter = {}
frontmatter_started = False
frontmatter_ended = False
frontmatter_pattern = re.compile(r"^\s*([a-z_]+):\s*(.*)\s*$", re.IGNORECASE)
try:
with open(file_path, "r", encoding="utf-8") as file:
first_line = file.readline()
if first_line.strip() != '"""':
# The file doesn't start with triple quotes
return {}
frontmatter_started = True
for line in file:
if '"""' in line:
if frontmatter_started:
frontmatter_ended = True
break
if frontmatter_started and not frontmatter_ended:
match = frontmatter_pattern.match(line)
if match:
key, value = match.groups()
frontmatter[key.strip()] = value.strip()
except FileNotFoundError:
print(f"Error: The file {file_path} does not exist.")
return {}
except Exception as e:
print(f"An error occurred: {e}")
return {}
return frontmatter
def load_toolkit_module_by_id(toolkit_id):
toolkit_path = os.path.join(TOOLS_DIR, f"{toolkit_id}.py")
if not os.path.exists(toolkit_path):
tool = Tools.get_tool_by_id(toolkit_id)
if tool:
with open(toolkit_path, "w") as file:
file.write(tool.content)
else:
raise Exception(f"Toolkit not found: {toolkit_id}")
spec = util.spec_from_file_location(toolkit_id, toolkit_path)
module = util.module_from_spec(spec)
frontmatter = extract_frontmatter(toolkit_path)
try:
install_frontmatter_requirements(frontmatter.get("requirements", ""))
spec.loader.exec_module(module)
print(f"Loaded module: {module.__name__}")
if hasattr(module, "Tools"):
return module.Tools(), frontmatter
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
def load_function_module_by_id(function_id):
function_path = os.path.join(FUNCTIONS_DIR, f"{function_id}.py")
if not os.path.exists(function_path):
function = Functions.get_function_by_id(function_id)
if function:
with open(function_path, "w") as file:
file.write(function.content)
else:
raise Exception(f"Function not found: {function_id}")
spec = util.spec_from_file_location(function_id, function_path)
module = util.module_from_spec(spec)
frontmatter = extract_frontmatter(function_path)
try:
install_frontmatter_requirements(frontmatter.get("requirements", ""))
spec.loader.exec_module(module)
print(f"Loaded module: {module.__name__}")
if hasattr(module, "Pipe"):
return module.Pipe(), "pipe", frontmatter
elif hasattr(module, "Filter"):
return module.Filter(), "filter", frontmatter
elif hasattr(module, "Action"):
return module.Action(), "action", frontmatter
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
def install_frontmatter_requirements(requirements):
if requirements:
req_list = [req.strip() for req in requirements.split(",")]
for req in req_list:
print(f"Installing requirement: {req}")
subprocess.check_call([sys.executable, "-m", "pip", "install", req])
else:
print("No requirements found in frontmatter.")
|