utils.py 2.19 KB
Newer Older
Timothy J. Baek's avatar
Timothy J. Baek committed
1
2
from importlib import util
import os
3
import re
Timothy J. Baek's avatar
Timothy J. Baek committed
4

Timothy J. Baek's avatar
Timothy J. Baek committed
5
from config import TOOLS_DIR, FUNCTIONS_DIR
Timothy J. Baek's avatar
Timothy J. Baek committed
6
7


8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
def extract_frontmatter(file_path):
    """
    Extract frontmatter as a dictionary from the specified file path.
    """
    frontmatter = {}
    frontmatter_pattern = re.compile(r"^([a-z_]+):\s*(.*)\s*$", re.IGNORECASE)

    with open(file_path, "r", encoding="utf-8") as file:
        for line in file:
            if line.strip() == '"""':
                # End of frontmatter section
                break
            match = frontmatter_pattern.match(line)
            if match:
                key, value = match.groups()
                frontmatter[key] = value

    return frontmatter


Timothy J. Baek's avatar
Timothy J. Baek committed
28
29
30
31
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)
32
    frontmatter = extract_frontmatter(toolkit_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
33
34
35
36
37

    try:
        spec.loader.exec_module(module)
        print(f"Loaded module: {module.__name__}")
        if hasattr(module, "Tools"):
38
            return module.Tools(), frontmatter
Timothy J. Baek's avatar
Timothy J. Baek committed
39
40
41
42
43
44
45
        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
Timothy J. Baek's avatar
Timothy J. Baek committed
46
47
48
49
50
51
52


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)
53
    frontmatter = extract_frontmatter(function_path)
Timothy J. Baek's avatar
Timothy J. Baek committed
54
55
56
57
58

    try:
        spec.loader.exec_module(module)
        print(f"Loaded module: {module.__name__}")
        if hasattr(module, "Pipe"):
59
            return module.Pipe(), "pipe", frontmatter
Timothy J. Baek's avatar
Timothy J. Baek committed
60
        elif hasattr(module, "Filter"):
61
            return module.Filter(), "filter", frontmatter
Timothy J. Baek's avatar
Timothy J. Baek committed
62
63
64
65
66
67
68
        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