generatingDnds

main.py

import os
import glob
import argparse
import json

from supabase import create_client, Client

url: str = os.environ.get("NEXT_PUBLIC_SUPABASE_URL")
key: str = os.environ.get("SUPABASE_SERVICE_ROLE_KEY")
supabase: Client = create_client(url, key)


def read_file(file_path):
    with open(file_path, "r") as f:
        return f.read()


def parse_answer_key(answer_key_text):
    lines = answer_key_text.split("\n")
    optimal_answer_line = [line for line in lines if line.startswith("Optimal answer:")]
    impossibilities_line = [line for line in lines if "but would require" in line][0]

    optimal_answer = [
        x.strip()
        for x in optimal_answer_line[0].split("Optimal answer: ")[1].split(",")
    ]
    impossibilities = [x.strip() for x in impossibilities_line[1:-1].split(", ")]

    return {
        "full_text": answer_key_text,
        "optimal_answer": optimal_answer,
        "impossibilities": impossibilities,
    }


def parse_proposal(proposal_file):
    return [
        [line.split(",", 1)[0][-1], line.strip()]
        for line in open(proposal_file, "r").readlines()
    ]


def create_extra_data(history_data, answer_key, proposals):
    return {
        "history_data": history_data,
        "answer_key": answer_key,
        "proposals": proposals,
    }


def create_dnd_task_variant(task_spec_id, extra_data, slug):
    response = (
        supabase.table("TaskVariant")
        .insert(
            [
                {
                    "task_type": task_spec_id,
                    "extra_data": extra_data,
                    "specification": slug,
                }
            ]
        )
        .execute()
    )
    return response.data


def upload_task_variants(task_type_id, input_folder):
    data_files = glob.glob(os.path.join(input_folder, "data_*.csv"))
    answer_key_files = [
        file.replace("data_", "answerkey_").replace(".csv", ".txt")
        for file in data_files
    ]
    proposals_files = [
        file.replace("data_", "proposals_").replace(".csv", ".txt")
        for file in data_files
    ]

    for data_file, answer_key_file, proposals_file in zip(
        data_files, answer_key_files, proposals_files
    ):
        try:
            history_data = read_file(data_file)
            answer_key_text = read_file(answer_key_file)
            parsed_answer_key = parse_answer_key(answer_key_text)
            proposals = parse_proposal(proposals_file)

            extra_data = create_extra_data(history_data, parsed_answer_key, proposals)
            create_dnd_task_variant(task_type_id, extra_data, data_file)

            print(f"Task variant created for {data_file} and {answer_key_file}")
        except Exception as error:
            print(
                f"Error creating task variant for {data_file} and {answer_key_file}:",
                error,
            )
            raise error


def create_task_type(input_folder, task_type_id):
    long_description = read_file(
        os.path.join(input_folder, "spec", "long_description.html")
    )

    if task_type_id:
        response = (
            supabase.table("TaskType")
            .update(
                {
                    "long_description": long_description,
                    "name": "On the Construction of Inconstructible Structures",
                    "slug": "dnd",
                    "short_description": "Help Duke Arado expand his fantastical estate by choosing the top four plans most likely to defy the laws of nature. Will you uncover the secrets and bring more extraordinary structures to life?",
                    "image_name": "dnd.png",
                }
            )
            .eq("id", task_type_id)
            .execute()
        )
    else:
        response = (
            supabase.table("TaskType")
            .insert(
                {
                    "long_description": long_description,
                    "name": "On the Construction of Inconstructible Structures",
                    "slug": "dnd",
                    "short_description": "Help Duke Arado expand his fantastical estate by choosing the top four plans most likely to defy the laws of nature. Will you uncover the secrets and bring more extraordinary structures to life?",
                    "image_name": "dnd.png",
                }
            )
            .execute()
        )

    return response.data[0]["id"]


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--task-type-id",
        type=int,
        help="Task type ID of DND if exists to update. If not, create a new one.",
    )
    parser.add_argument(
        "--input-folder",
        type=str,
        required=True,
        help="Input folder with data and answer key files",
    )
    args = parser.parse_args()

    task_type = create_task_type(args.input_folder, args.task_type_id)
    upload_task_variants(task_type, args.input_folder)