prompt_loader.py

import os
import re
import sys
import importlib.util
import argparse


def extract_variables(template_content):
    pattern = re.compile(r"<var\s+(\w+)\s*>")
    return sorted(set(pattern.findall(template_content)))


def generate_function_code(
    function_name, template_content, variables, template_path
):
    params = ", ".join([f"{var}: str = None" for var in variables])
    template_args = ", ".join([f"{var}={var}" for var in variables])
    return f"""
def {function_name}({params}) -> str:
    '''/{template_path}\n---\n{template_content}'''
    template = env.get_template('{os.path.basename(template_path)}')
    return template.render({template_args})
"""


def setup(prompts_dir):
    code_string = f"""from jinja2 import Environment, FileSystemLoader
    
env = Environment(
    loader=FileSystemLoader('{prompts_dir}'),
    variable_start_string='<var',
    variable_end_string='>',
    comment_start_string='<!--',
    comment_end_string='-->',
)
    """

    templates = [f for f in os.listdir(prompts_dir) if f.endswith(".qmd")]

    for filename in templates:
        template_path = os.path.join(prompts_dir, filename)
        function_name = filename[:-4]  # Remove .qmd extension
        with open(template_path, "r") as f:
            template_content = f.read()

        variables = extract_variables(template_content)
        code_string += generate_function_code(
            function_name, template_content, variables, template_path
        )

    functions_dir = os.path.join(prompts_dir, "functions")
    os.makedirs(functions_dir, exist_ok=True)

    temp_file = os.path.join(functions_dir, "pyprompts.py")

    with open(temp_file, "w") as f:
        f.write(code_string)

    pycache_path = os.path.join(functions_dir, "__pycache__")
    if not os.path.exists(pycache_path) or not any(
        f.endswith(".pyc") for f in os.listdir(pycache_path)
    ):
        spec = importlib.util.spec_from_file_location("pyprompts", temp_file)
        module = importlib.util.module_from_spec(spec)
        sys.modules[spec.name] = module
        spec.loader.exec_module(module)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Provide path to prompts directory"
    )

    parser.add_argument("--path", type=str, help="path to prompts directory")
    args = parser.parse_args()
    setup(args.path)