run.py

#!/usr/bin/env python3

"""

This is a work in progress

"""


import re
import os
from copy import copy
import json
import datetime
import argparse
from subprocess import check_call

from openai_helper import (
    ask_chatgpt,
    NO_CODE_GEN,
    NO_TEMPLATE_REGEN,
    SKELETON_REGEN,
    SKELETON_REGEN2,
    SKELETON_REGEN3,
    BASE_HTML,
)

file_dir = os.path.dirname(__file__)

parser = argparse.ArgumentParser(description="Generate a program using TSL")
parser.add_argument(
    "-d", "--dir", required=True, help="Directory of the prompts to generate from"
)
parser.add_argument(
    "-m",
    "--method",
    default="nl",
    choices=["nl", "nls"],
    help='"nl" means generate using only the natural language promps, "nls" means also use the TSL spec, and "nl+spec+synth" means also use the synthesized code.',
)
parser.add_argument(
    "-l",
    "--lang",
    default="JavaScript",
    help="Target generated implementation programming language (default: JavaScript)",
)
parser.add_argument(
    "--no-openai",
    action="store_true",
    help="Do not query the openai api and instead output prompts for the user to paste into AI chats.",
)
parser.add_argument(
    "--spec-prompt-file",
    default="Spec_template.prompt",
    help="filename to use for the spec prompt",
)
parser.add_argument(
    "--log-results",
    action="store_true",
    help='Copy the affected "{Experiment}/computed" dir into "results" with the timestamp"',
)
parser.add_argument(
    "--num-iter", default=1, help="repeat the experiment [num-iter] times"
)
parser.add_argument(
    "--regen-html", default=False, action="store_true", help="regenerate the template?"
)
parser.add_argument(  # defaults to 3.5-t for debug and testing
    "--model", default="gpt-3.5-turbo", help="which model to use from openai api"
)
parser.add_argument("--trusted", default=False, action="store_true")

parser.add_argument("--llm-only", default=False, action="store_true")
parser.add_argument("--llmtsl", default="tsl")

args = parser.parse_args()

>>>>>>> big

>>>>>>> 5035c3b09a82170526a4fd85bf615b9e7f9269d6
def run_with_args(args):
    print(datetime.datetime.now(), "Running with args", args)

    # handle repeated iterations
    if 1 < int(args.num_iter):
        for _ in range(int(args.num_iter)):
            new_args = copy(args)
            new_args.num_iter = 1
            run_with_args(new_args)
        return

    # handle "run all benchmarks"
    if args.dir.lower() == "all":
        for benchmark in json.load(open(os.path.join(file_dir, "benchmarks.json"))):
            new_args = copy(args)
            new_args.dir = benchmark
            run_with_args(new_args)
        return

    # handle LOPSTR experiment
    if args.spec_prompt_file.lower() == "lopstr":
        for prompt_file in [
            "Spec_template.prompt",
            "Spec_onlyNlAndSummary_template.prompt",
            "Spec_onlySummaryAndHeaders_template.prompt",
        ]:
            new_args = copy(args)
            new_args.spec_prompt_file = prompt_file
            run_with_args(new_args)
        return

    to_interpolate_tsl = {
        # "few_shot_header": "shotPrompt.txt",
        "functions_and_predicates": "Headers.txt",
        "natural_language_summary": "NL.summary.txt",
        "natural_language_description": "NL.txt",
        "wrapper_api": "Wrapper_api.js",
    }

    to_interpolate_llm = {
        "functions_and_predicates": "Headers.txt",
        "natural_language_summary": "NL.summary.txt",
        "natural_language_description": "NL.txt",
        "html_temp": "wrapper_template.html",
    }

    computed_dir = os.path.join(args.dir, "computed")
    if not os.path.exists(computed_dir):
        os.makedirs(computed_dir)

    with open(os.path.join(computed_dir, "args.json"), "w") as file:
        json.dump(args.__dict__, file, indent=2)

    def check_content_between_markers(filename, start_marker, end_marker):
        with open(os.path.join(args.dir, filename), "r") as file:
            content = file.read()
        start_index = content.find(start_marker) + len(start_marker)
        end_index = content.find(end_marker)
        if start_index == -1 or end_index == -1 or end_index <= start_index:
            return False  # Markers not found or in the wrong order
        return content[start_index:end_index].strip() == ""

    use_no_assumptions = check_content_between_markers(
        "NL.txt", "Assumptions:", "Guarantees:"
    )
    use_no_functions = check_content_between_markers(
        "Headers.txt", "Functions:", "Predicates:"
    )

    if use_no_assumptions and use_no_functions:
        spec_template = "Spec_withoutFA_template.prompt"
        impl_template = "Impl_withoutFunctions_template.prompt"
    elif use_no_assumptions:
        spec_template = "Spec_withoutAssumptions_template.prompt"
    elif use_no_functions:
        spec_template = "Spec_withoutFunctions_template.prompt"
        impl_template = "Impl_withoutFunctions_template.prompt"
    else:
        spec_template = "Spec_template.prompt"
        impl_template = "Impl_template.prompt"

    # override with filename
    if args.spec_prompt_file != "Spec_template.prompt":
        spec_template = args.spec_prompt_file

    def load_file_and_interpolate(filename, llm=False):

        to_interpolate = to_interpolate_llm if llm else to_interpolate_tsl

        with open(filename) as file:
            result = file.read()
            for name, filename in to_interpolate.items():
                path = os.path.join(args.dir, filename)
                with open(path) as file:
                    replacement = file.read()
                for to_replace in [
                    "[[" + name.upper() + "_GO_HERE]]",
                    "[[" + name.upper() + "_GOES_HERE]]",
                ]:
                    result = result.replace(to_replace, replacement)
            result = result.replace("[[LANGUAGE_GOES_HERE]]", args.lang)
        return result

    def output_error(message):
        print(message)
        with open(os.path.join(computed_dir, "err.log"), "a") as file:
            file.writelines([message])
        log_results()

<<<<<<< HEAD
    def log_results():
        if args.method == "no_tsl":
            template_filename = "No_tsl.prompt"
        else:
            template_filename = spec_template            
        log_outer_dir = os.path.join(
            file_dir, "results", "by_benchmark", args.dir, template_filename
=======
    def log_results(exp=args.llmtsl, trusted=args.trusted):
        log_outer_dir = os.path.join(
            file_dir,
            "results",
            f"{exp}_by_benchmark_{'regen_with_import' if trusted else 'regen_no_import'}",
            args.dir,
            spec_template,
            args.model,
>>>>>>> big
        )
        if not os.path.exists(log_outer_dir):
            os.makedirs(log_outer_dir)
        log_dir = os.path.join(log_outer_dir, datetime.datetime.now().isoformat())
        if args.verbose:
            print("Moving computed dir to", log_dir)
        check_call(["mv", computed_dir, log_dir])
        print(f"{log_dir}")

    def extract_first_code_block(text):
        match = re.search(r"```.*?\n((?:.|\n)+?)```", text)
        if match == None:
            return None
        else:
            return match[1]

    if args.method == "nl":
        num_retries = 10
        tries = 0
        realized = False

        while tries < num_retries and not realized:
            # TODO add retry loop for ill formed specs here. Add some prompt that passes the erorr message with a tesmplate back to llm
            spec_template_path = os.path.join(
                file_dir, spec_template
            )  # use dynamically chosen spec template
            spec_filename = os.path.join(computed_dir, "Spec.tsl")
            synth_filename = os.path.join(computed_dir, "Synth.js")
            spec_response_filename = os.path.join(computed_dir, "Spec_response.txt")

            spec_prompt = load_file_and_interpolate(spec_template_path)
            spec_prompt_filename = os.path.join(computed_dir, "Spec.prompt")
            with open(spec_prompt_filename, "w") as file:
                file.write(spec_prompt)

            if args.no_openai:
                print(
                    f"Please paste the contents of this file into the TSL GPT:\n\n    {spec_prompt_filename}\n\nThen, paste the spec in the response into:\n\n    {spec_filename}\n\nWhen you have done this, press Enter"
                )
                input()
            else:
                response = ask_chatgpt(spec_prompt, args.model)
                code_block = extract_first_code_block(
                    response.choices[0].message.content
                )
                with open(spec_response_filename, "w") as file:
                    file.write(response.choices[0].message.content)
                if code_block == None:
                    return output_error(
                        f"No valid code block in response. See {spec_response_filename}"
                    )
                else:
                    with open(spec_filename, "w") as file:
                        file.write(code_block)

            try:
                check_call(
                    [
                        "tsl",
                        "synthesize",
                        "-i",
                        spec_filename,
                        "--js",
                        "-o",
                        synth_filename,
                    ]
                )
                realized = True
            except BaseException as e:
                print(f"Error synthesizing: {str(e)}")
                tries += 1
                if tries == num_retries:
                    return output_error(str(e))

    elif args.method == "nls":
        # prompt_templates = (
        #     SKELETON_REGEN.format(spec_prompt, wrapper_contents, code_block),
        # )
        spec_template_path = os.path.join(file_dir, spec_template)
        spec_filename = os.path.join(computed_dir, "Spec.tsl")
        synth_filename = os.path.join(computed_dir, "Synth.js")
        spec_response_filename = os.path.join(computed_dir, "Spec_response.txt")

        with open(os.path.join(file_dir, args.dir, "Headers.txt"), "r") as headers_file:
            headers = headers_file.read()

        with open(
            os.path.join(file_dir, args.dir, "NL.summary.txt"), "r"
        ) as nl_summ_file:
            nl_summ = nl_summ_file.read()

        with open(os.path.join(file_dir, args.dir, "NL.txt"), "r") as nl_desc_file:
            nl_desc = nl_desc_file.read()

        with open(
            os.path.join(file_dir, args.dir, "wrapper_template.html"), "r"
        ) as html_temp_file:
            html_wrapp = html_temp_file.read()

        spec_prompt = SKELETON_REGEN.format(nl_summ, nl_desc, headers, " ")

        response = ask_chatgpt(
            f"{spec_prompt}\nGenerate using the above as inspiration a working html file which implements the described program",
            args.model,
        )
        code_block = extract_first_code_block(response.choices[0].message.content)
        with open(spec_response_filename, "w") as file:
            file.write(response.choices[0].message.content)
        if code_block == None:
            return output_error(
                f"No valid code block in response. See {spec_response_filename}"
            )
        else:
            dir_name = os.path.basename(os.path.normpath(args.dir))
            output_html_filename = f"{dir_name}.html"
            output_html_path = os.path.join(computed_dir, output_html_filename)

            with open(output_html_path, "w") as file:
                file.write(code_block)
        log_results()
        return 0

<<<<<<< HEAD
        try:
            check_call(
                ["timeout", str(args.tsl_timeout), "tsl", "synthesize", "-i", spec_filename, "--js", "-o", synth_filename]
            )
        except BaseException as e:
            return output_error(str(e))
    elif args.method == "no_tsl":
        impl_template_path = os.path.join(
            file_dir, "No_tsl.prompt"
        )
        impl_response_filename = os.path.join(computed_dir, "Impl_response.txt")
        impl_filename = os.path.join(computed_dir, "Impl.js")

        impl_prompt = load_file_and_interpolate(impl_template_path)
        impl_prompt_filename = os.path.join(computed_dir, "No_tsl.prompt")
        with open(impl_prompt_filename, "w") as file:
            file.write(impl_prompt)

        if args.no_openai:
            print(
                f"Please paste the contents of this file into the TSL GPT:\n\n    {impl_prompt_filename}\n\nThen, paste the first code block in the response into:\n\n    {impl_filename}\n\nWhen you have done this, press Enter"
            )
            input()
        else:
            response = ask_chatgpt(impl_prompt, args.model)
            code_block = extract_first_code_block(response.choices[0].message.content)
            with open(impl_response_filename, "w") as file:
                file.write(response.choices[0].message.content)
            if code_block == None:
                return output_error(
                    f"No valid code block in response. See {impl_response_filename}"
                )
            else:
                with open(impl_filename, "w") as file:
                    file.write(code_block)
=======
>>>>>>> big
    else:
        spec_filename = os.path.join(args.dir, spec_template)
        raise Exception('the only method implemented is "nl". the others are TODO.')

    if args.method != 'no_tsl':
        impl_prompt_filename = os.path.join(computed_dir, "Impl.prompt")
        impl_filename = os.path.join(computed_dir, "Impl.js")
        impl_template_path = os.path.join(file_dir, impl_template)
        impl_prompt = load_file_and_interpolate(impl_template_path)
        impl_response_filename = os.path.join(computed_dir, "Impl_response.txt")
        with open(impl_prompt_filename, "w") as file:
            file.write(impl_prompt)

        if args.no_openai:
            print(
                f"Please paste the contents of this file into the TSL GPT:\n\n    {impl_prompt_filename}\n\nThen, paste the spec in the response into:\n\n    {impl_filename}\n\nWhen you have done this, press Enter"
            )
            input()
        else:
            response = ask_chatgpt(impl_prompt)
            code_block = extract_first_code_block(response.choices[0].message.content)
            with open(impl_response_filename, "w") as file:
                file.write(response.choices[0].message.content)
            if code_block == None:
                return output_error(
                    f"No valid code block in response. See {impl_response_filename}"
                )
            else:
                with open(impl_filename, "w") as file:
                    file.write(code_block)

    dir_name = os.path.basename(os.path.normpath(args.dir))
    output_html_filename = f"{dir_name}.html"

    wrapper_template_path = os.path.join(args.dir, "wrapper_template.html")
    output_html_path = os.path.join(
        computed_dir, output_html_filename
    )  # Output file path in the computed directory with the directory name

    if os.path.exists(wrapper_template_path):
        with open(wrapper_template_path, "r") as file:
            wrapper_contents = file.read()

        # Nik - added new flag, uses code block to seed gpt generation of code
        if args.regen_html:
            print("Using synthesized code to regenerate html")

            # synth filename and impl filename
            spec_template_path = os.path.join(file_dir, spec_template)
            spec_filename = os.path.join(computed_dir, "Spec.tsl")
            synth_filename = os.path.join(computed_dir, "Synth.js")
            spec_response_filename = os.path.join(computed_dir, "Spec_response.txt")

            with open(
                os.path.join(file_dir, args.dir, "Headers.txt"), "r"
            ) as headers_file:
                headers = headers_file.read()

            with open(
                os.path.join(file_dir, args.dir, "NL.summary.txt"), "r"
            ) as nl_summ_file:
                nl_summ = nl_summ_file.read()

            with open(os.path.join(file_dir, args.dir, "NL.txt"), "r") as nl_desc_file:
                nl_desc = nl_desc_file.read()

            with open(
                os.path.join(file_dir, args.dir, "computed/Synth.js"), "r"
            ) as html_temp_file:
                synth_code = html_temp_file.read()

            with open(
                os.path.join(file_dir, args.dir, "computed/Impl.js"), "r"
            ) as html_temp_file:
                impl_code = html_temp_file.read()

            with open(
                os.path.join(file_dir, args.dir, "computed/Spec.tsl"), "r"
            ) as html_temp_file:
                spectsl = html_temp_file.read()

            if args.trusted:
                regen_prompt = SKELETON_REGEN2.format(
                    spec_prompt, impl_prompt, "", synth_code, spectsl
                )
            else:
                regen_prompt = SKELETON_REGEN3.format(
                    synth_code, impl_code, spectsl, nl_summ, nl_desc
                )

            response = ask_chatgpt(f"{regen_prompt}", args.model)

            code_block = extract_first_code_block(response.choices[0].message.content)

            with open(spec_response_filename, "w") as file:
                file.write(response.choices[0].message.content)
            if code_block == None:
                return output_error(
                    f"No valid code block in response. See {spec_response_filename}"
                )
            else:
                dir_name = os.path.basename(os.path.normpath(args.dir))
                output_html_filename = f"{dir_name}.html"
                output_html_path = os.path.join(computed_dir, output_html_filename)

                with open(output_html_path, "w") as file:
                    file.write(code_block)
            log_results()
            return 0

        print("Inserting generated code into base template")
        with open(impl_filename, "r") as file:
            impl_contents = file.read()
        placeholder_comment = "//[[GENERATED_FUNCTIONS_AND_PREDICATES_GO_HERE]]"
        if placeholder_comment in wrapper_contents:
            modified_wrapper = wrapper_contents.replace(
                placeholder_comment, impl_contents
            )
            with open(output_html_path, "w") as file:
                file.write(modified_wrapper)
            print(
                f"The Impl.js contents have been successfully inserted into {output_html_path}."
            )
        else:
            return output_error(
                "The placeholder comment does not match in the wrapper_template.html. Please ensure it is '//[[GENERATED_FUNCTIONS_AND_PREDICATES_GO_HERE]]'."
            )
    else:
        return output_error(
            "Please include a wrapper_template.html file with a line of comment '//[[GENERATED_FUNCTIONS_AND_PREDICATES_GO_HERE]]' inside the file to get the generated wrapper HTML."
        )

    log_results()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Generate a program using TSL")
    parser.add_argument(
        "-d", "--dir", required=True, help="Directory of the prompts to generate from"
    )
    parser.add_argument(
        "-m",
        "--method",
        default="nl",
        choices=["nl", "nl+spec", "nl+spec+synth", "no_tsl"],
        help='"nl" means generate using only the natural language promps, "nl+spec" means also use the TSL spec, and "nl+spec+synth" means also use the synthesized code.',
    )
    parser.add_argument(
        "-l",
        "--lang",
        default="JavaScript",
        help="Target generated implementation programming language (default: JavaScript)",
    )
    parser.add_argument(
        "--no-openai",
        action="store_true",
        help="Do not query the openai api and instead output prompts for the user to paste into AI chats.",
    )
    parser.add_argument(
        "--model",
        default='gpt-4-turbo',
        help="The openai model (e.g. \"gpt-4-turbo\")",
    )
    parser.add_argument(
        "--spec-prompt-file",
        default="Spec_template.prompt",
        help="filename to use for the spec prompt",
    )
    parser.add_argument(
        "--log-results",
        action="store_true",
        help='Copy the affected "{Experiment}/computed" dir into "results" with the timestamp"',
    )
    parser.add_argument(
        "--num-iter", default=1, help="repeat the experiment [num-iter] times"
    )
    parser.add_argument(
        "--regen-html",
        default=False,
        action="store_true",
        help="regenerate the template?",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="print messages other than errors.",
    )
    parser.add_argument(
        "--tsl-timeout",
        type=int,
        default=120,
        help="TSL synthesis timeout",
    )

    args = parser.parse_args()

    run_with_args(args)