logger.py

import time
import argparse
import logging
import shelve
import affirmative_interrogative
import statements
import extract_predicates
import consult_query
import run_prolog
from benchmark.utils import get_prompt
import csv
from datetime import datetime


def inform(msg: str):
    print(f"\n\033[1;32;40m{msg}\033[0m")


def main():
    
    parser = argparse.ArgumentParser()

    parser.add_argument("--id", type=str)
    args = parser.parse_args()

    if args.id is None:
        with shelve.open("extracted") as db:
            args.id = db["id"]
    else:
        with shelve.open("extracted") as db:
            db["id"] = args.id

    prompt_id = args.id

    with shelve.open("extracted") as db:
        prompt_dict = get_prompt(prompt_id)

        if prompt_dict is None:
            logging.error(f"Prompt with id {prompt_id} not found")
            return
        
        prompt = prompt_dict["prompt"]
        expected_result = prompt_dict["expected"]
        db["prompt"] = prompt


    # logging.basicConfig(level=logging.INFO)
    inform("Process prompt")
    print(prompt)
    start_time = time.time()

    inform("Separate affirmative and interrogative sentences (gpt-4o-mini)")
    affirmative_interrogative.main()

    inform("Extract statements (gpt-4o-mini)")
    statements.main()

    inform("Extract predicates (gpt-4o)")
    extract_predicates.main()

    inform("Generate query (gpt-4o)")
    consult_query.main()

    # execute prolog
    inform("Execute prolog")
    result = run_prolog.main()

    # append to historical_tests.csv
    inform("Registrating test")
    test_date = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    with open('historical_tests.csv', mode='a', newline='') as file:
        writer = csv.writer(file)
        writer.writerow([test_date, prompt_id, result, expected_result])

    end_time = time.time()
    inform(f"Execution time: {end_time - start_time} seconds")
    # logging.info(f"Execution time: {end_time - start_time} seconds")


if __name__ == "__main__":
    main()