uw_programmatic/uw_machine.py

from __future__ import annotations

import ast
import csv
import json
from pathlib import Path
import random
from typing import TYPE_CHECKING
from griptape.structures import Agent
import pandas as pd

from uw_programmatic.base_machine import UWBaseMachine

if TYPE_CHECKING:
    from griptape.tools import BaseTool


class UWMachine(UWBaseMachine):
    """State machine with GOAP"""

    @property
    def tools(self) -> dict[str, BaseTool]:
        return {}

    def start_machine(self) -> None:
        """Starts the machine."""
        # Clear input history.
        # Clear csv file
        with Path(Path.cwd().joinpath("outputs/similarity_step.csv")).open("w") as file:
            file.write("")
        self.send("enter_first_state")

    def on_event_gather_parameters(self, event_: dict) -> None:
        event_source = event_["type"]
        event_value = event_["value"]
        match event_source:
            case "user_input":
                parameters = event_value
                self.page_range = parameters["page_range"]
                self.question_number = parameters["question_number"]
                self.taxonomy = parameters["taxonomy"]
                self.send("next_state")
            case _:
                err_msg = f"Unexpected Transition Event ID: {event_value}."
                raise ValueError(err_msg)

    def on_enter_evaluate_q_count(self) -> None:
        if len(self.question_list) >= self.question_number:
            self.send("finish_state")  # go to output questions
        else:
            self.send("next_state")  # go to need more questions

    def on_event_evaluate_q_count(self, event_: dict) -> None:
        pass

    def on_enter_need_more_q(self) -> None:
        # Create the entire workflow to create another question.
        self.get_questions_workflow().run()

    def on_event_need_more_q(self, event_: dict) -> None:
        event_source = event_["type"]
        event_value = event_["value"]
        match event_source:
            case "griptape_event":
                event_type = event_value["type"]
                match event_type:
                    case "FinishStructureRunEvent":
                        structure_id = event_value["structure_id"]
                        match structure_id:
                            case "create_question_workflow":
                                values = event_value["output_task_output"]["value"]
                                questions = [
                                    ast.literal_eval(question["value"])
                                    for question in values
                                ]
                                self.most_recent_questions = (
                                    questions  # This is a ListArtifact I'm pretty sure
                                )
                                self.send("next_state")
                    case _:
                        print(f"Error:{event_} ")
            case _:
                print(f"Unexpected: {event_}")

    def on_enter_assess_generated_q(self) -> None:
        # TODO: Should it append it to the list already and remove duplicates? or not?
        # TODO: Merge incoming lists
        with Path(Path.cwd().joinpath("outputs/similarity_step.csv")).open(
            "a", newline=""
        ) as file:
            writer = csv.DictWriter(
                file,
                fieldnames=[
                    "Question",
                    "Answer",
                    "Wrong Answers",
                    "Page",
                    "Taxonomy",
                ],
            )
            writer.writerow({"Question": "LIST OF QUESTIONS GENERATED THIS ROUND"})
            writer.writerows(self.most_recent_questions)
        merged_list = [*self.question_list, *self.most_recent_questions]
        prompt = f"{merged_list}"
        self.get_structure("similarity_auditor").run(prompt)

    def on_event_assess_generated_q(self, event_: dict) -> None:
        event_source = event_["type"]
        event_value = event_["value"]
        match event_source:
            case "griptape_event":
                event_type = event_value["type"]
                match event_type:
                    case "FinishStructureRunEvent":
                        structure_id = event_value["structure_id"]
                        match structure_id:
                            case "similarity_auditor":
                                new_question_list = event_value["output_task_output"][
                                    "value"
                                ]
                                try:
                                    new_question_list = json.loads(
                                        new_question_list
                                    )  # This must be in that JSON format
                                except:
                                    new_question_list = self.question_list
                                merged_list = [
                                    *self.question_list,
                                    *self.most_recent_questions,
                                ]
                                deleted_q = [
                                    question1
                                    for question1 in merged_list
                                    if not any(
                                        question2["Question"] == question1["Question"]
                                        for question2 in new_question_list
                                    )
                                ]
                                with Path(
                                    Path.cwd().joinpath("outputs/similarity_step.csv")
                                ).open("a", newline="") as file:
                                    writer = csv.DictWriter(
                                        file,
                                        fieldnames=[
                                            "Question",
                                            "Answer",
                                            "Wrong Answers",
                                            "Page",
                                            "Taxonomy",
                                        ],
                                    )
                                    writer.writerow(
                                        {"Question": "QUESTIONS REMOVED THIS ROUND!"}
                                    )
                                    if len(deleted_q):
                                        writer.writerows(deleted_q)
                                    else:
                                        writer.writerow({"Question": "No q removed"})
                                self.question_list = new_question_list
                                self.send("next_state")  # move on

    def on_enter_output_q(self) -> None:
        columns = pd.MultiIndex.from_tuples(
            [
                ("Professor", "Page Range"),
                ("Professor", "Taxonomy"),
                ("Professor", "Question"),
                ("Professor", "Answer"),
                ("Professor", "Wrong Answers"),
                ("Student", "Question"),
                ("Student", "Answers"),
            ]
        )
        data = pd.DataFrame(columns=columns)
        for question in range(len(self.question_list)):
            shuffled_answers = [
                self.question_list[question]["Answer"],
                *self.question_list[question]["Wrong Answers"],
            ]
            random.shuffle(shuffled_answers)
            shuffled_answers = "\n".join(shuffled_answers)
            new_row = [
                self.question_list[question]["Page"],
                self.question_list[question]["Taxonomy"],
                self.question_list[question]["Question"],
                self.question_list[question]["Answer"],
                self.question_list[question]["Wrong Answers"],
                self.question_list[question]["Question"],
                shuffled_answers,
            ]
            data.loc[question] = new_row
        data.columns = ["_".join(col).strip() for col in data.columns.values]
        writer = pd.ExcelWriter("outputs/professor_guide.xlsx", engine="xlsxwriter")
        data.to_excel(writer, sheet_name="Quiz Questions", index=False)
        writer.close()
        self.send("next_state")

    def on_event_output_q(self, event_: dict) -> None:
        pass

    def on_exit_output_q(self) -> None:
        # Reset the state machine values
        self.question_list = []
        self.most_recent_questions = []

    if __name__ == "__main__":

        question_list = [
            {
                "Page": "1-2",
                "Taxonomy": "Knowledge",
                "Question": "What is Python?",
                "Answer": "A programming language",
                "Wrong Answers": ["A snake", "A car brand", "A fruit"],
            },
            {
                "Page": "3-4",
                "Taxonomy": "Comprehension",
                "Question": "What does HTML stand for?",
                "Answer": "HyperText Markup Language",
                "Wrong Answers": [
                    "High Text Machine Language",
                    "Hyperlink Text Mode Language",
                    "None of the above",
                ],
            },
        ]

        columns = pd.MultiIndex.from_tuples(
            [
                ("Professor", "Page Range"),
                ("Professor", "Taxonomy"),
                ("Professor", "Question"),
                ("Professor", "Answer"),
                ("Professor", "Wrong Answers"),
                ("Student", "Question"),
                ("Student", "Answers"),
            ]
        )
        data = pd.DataFrame(columns=columns)
        for question in range(len(question_list)):
            shuffled_answers = [
                question_list[question]["Answer"],
                *question_list[question]["Wrong Answers"],
            ]
            random.shuffle(shuffled_answers)
            shuffled_answers = "\n".join(shuffled_answers)
            new_row = [
                question_list[question]["Page"],
                question_list[question]["Taxonomy"],
                question_list[question]["Question"],
                question_list[question]["Answer"],
                question_list[question]["Wrong Answers"],
                question_list[question]["Question"],
                shuffled_answers,
            ]
            data.loc[question] = new_row
        data.columns = ["_".join(col).strip() for col in data.columns.values]
        writer = pd.ExcelWriter("outputs/professor_guide.xlsx", engine="xlsxwriter")
        data.to_excel(writer, sheet_name="Quiz Questions", index=False)
        writer.close()