enh: repopulate tools & functions from db if non existent

This commit is contained in:
Timothy J. Baek 2024-08-26 12:18:12 +02:00
parent e98a20fce9
commit efd4b03f78

View File

@ -4,6 +4,9 @@ import re
import sys import sys
import subprocess import subprocess
from apps.webui.models.tools import Tools
from apps.webui.models.functions import Functions
from config import TOOLS_DIR, FUNCTIONS_DIR from config import TOOLS_DIR, FUNCTIONS_DIR
@ -49,6 +52,15 @@ def extract_frontmatter(file_path):
def load_toolkit_module_by_id(toolkit_id): def load_toolkit_module_by_id(toolkit_id):
toolkit_path = os.path.join(TOOLS_DIR, f"{toolkit_id}.py") 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) spec = util.spec_from_file_location(toolkit_id, toolkit_path)
module = util.module_from_spec(spec) module = util.module_from_spec(spec)
frontmatter = extract_frontmatter(toolkit_path) frontmatter = extract_frontmatter(toolkit_path)
@ -71,6 +83,14 @@ def load_toolkit_module_by_id(toolkit_id):
def load_function_module_by_id(function_id): def load_function_module_by_id(function_id):
function_path = os.path.join(FUNCTIONS_DIR, f"{function_id}.py") 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) spec = util.spec_from_file_location(function_id, function_path)
module = util.module_from_spec(spec) module = util.module_from_spec(spec)
frontmatter = extract_frontmatter(function_path) frontmatter = extract_frontmatter(function_path)