uw_programmatic/single_question_machine.py

from __future__ import annotations

import os
import random
import schema
from typing import TYPE_CHECKING, cast

from griptape.configs import Defaults
from griptape.configs.drivers import (
    OpenAiDriversConfig,
)
from griptape.drivers import GriptapeCloudVectorStoreDriver, OpenAiChatPromptDriver
from griptape.engines.rag import RagEngine
from griptape.engines.rag.modules import (
    TextChunksResponseRagModule,
    VectorStoreRetrievalRagModule,
)
from griptape.engines.rag.stages import ResponseRagStage, RetrievalRagStage
from griptape.events import (
    BaseEvent,
    EventBus,
    EventListener,
    FinishStructureRunEvent,
)
from griptape.rules import Rule, Ruleset
from griptape.structures import Agent
from griptape.tasks import ToolTask
from griptape.tools import RagTool
from statemachine import State, StateMachine
from statemachine.factory import StateMachineMetaclass

if TYPE_CHECKING:
    from griptape.structures import Structure

Defaults.drivers_config = OpenAiDriversConfig(
    prompt_driver=OpenAiChatPromptDriver(model="gpt-4o", max_tokens=4096)
)

# States will be:
# random_selection (does dice roll and kb selection plus information task)
# Question generation (generates the question and answer properly)
# Wrong answer generation (generates a wrong answer?)
# Compile task (finishes all and compiles it into a neat thing)
# TODO: How to get it to return everything
STATES = [
    "start",
    "random_selection",
    "get_textbook",
    "question_generation",
    "wrong_answer_generation",
    "audit_question",
    "compile_task",
    "end",
]
START = "start"
END = "end"
TRANSITIONS = [
    {
        "event": "next_state",
        "transitions": [
            {"from": "random_selection", "to": "get_textbook"},
            {"from": "get_textbook", "to": "question_generation"},
            {"from": "question_generation", "to": "wrong_answer_generation"},
            {"from": "wrong_answer_generation", "to": "audit_question"},
            {"from": "audit_question", "to": "compile_task"},
            {"from": "compile_task", "to": "end"},
        ],
    },
    {
        "event": "redo",
        "transitions": [
            {"from": "audit_question", "to": "wrong_answer_generation"},
        ],
    },
    {
        "event": "start_up",
        "transitions": [
            {"from": "start", "to": "random_selection"},
        ],
    },
    {"event": "end_state", "transitions": [{"from": "random_selection", "to": "end"}]},
]
RULESETS = {
    "specific_question_creator": [
        """Question should be a multiple choice quiz style question that assesses a students knowledge of the information in the knowledge base (which should be referred
      to as 'the textbook'). Answer should be a correct answer to the question that
      uses information from the knowledge base. Do not return incorrect answers.""",
        """The length of the question should be 30 words at most.""",
        """Question should never reference or ask about an entire section, never reference
      or ask about a quote in the knowledge base, never ask for the page number of
      some information, and never ask for information about the file, document, or
      knowledge base.""",
        """The answer to the question should be short, but should not omit important
      information.""",
    ],
    "incorrect_answers_creator": [
        """All incorrect answers should be different, but plausible answers to the question.""",
        """Incorrect answers may reference material from the info provided as context, but must not be correct answers to the question""",
        """Incorrect answers should always have a similar structure to the correct answer.""",
        """The length of all incorrect answers should be as close to the correct answer as possible while remaining plausible.""",
    ],
    "question_auditor_ruleset": [
        # """If any of the rules are false, return false and why. If they are all true, return true.""",
        """If any of the rules are false, return True for the part of the question why they are false.""",
        """The reason why it is false is between 3-7 words""",
        """There is exactly one correct answer.""",
        """The correct answer has a clearly distinct meaning from all incorrect answers.""",
        """Incorrect answers are plausible to someone who does not know the correct answer""",
        """All answer choices are on the same topic as the question""",
        """All answer choices are relevant to the context of the question, with no unrelated concepts or entities.""",
        """All answers have semantically different meanings from one another, even if they are syntactically similar.""",
        """All answer choices are parallel to one another with respect to grammatical structure, length, and complexity""",
    ],
}
STRUCTURES = {
    "subject_matter_expert": {"ruleset_ids": ["specific_question_creator"]},
    "wrong_answers_generator": {"ruleset_ids": ["incorrect_answers_creator"]},
    "question_auditor": {"ruleset_ids": ["question_auditor_ruleset"]},
}


class SingleQuestion(StateMachine):
    "Base class for machine"

    def __init__(self, **kwargs):
        self._structures = {}
        self.kb_ids = kwargs["kb_ids"]
        self.page_range: tuple = kwargs["page_range"]
        self.taxonomy_choices: list = kwargs["taxonomy_choices"]
        self.question: str = ""
        self.answer: str = ""
        self.wrong_answers: list = []
        self.generated_question: dict = {}
        self.taxonomy: str = ""
        self.rejected: bool = False
        self.give_up: int = 0
        self.reject_reason: str = ""

        def on_event(event: BaseEvent) -> None:
            "Takes in griptape events and fixes them"
            try:
                self.send("griptape_event", event_=event.to_dict())
            except Exception as e:
                errormsg = f"Would not allow Griptape Event to be sent"
                raise ValueError(errormsg) from e

        EventBus.clear_event_listeners()
        EventBus.add_event_listener(
            EventListener(on_event, event_types=[FinishStructureRunEvent]),
        )
        super().__init__()

    @classmethod
    def create_statemachine(
        cls, taxonomy_choices: list, kb_ids: dict, page_range: tuple
    ) -> SingleQuestion:
        states_instances = {}
        events = {}
        for state in STATES:
            initial = state == START
            final = state == END
            # Creates the states
            states_instances[state] = State(value=state, initial=initial, final=final)
            if not (initial or final or state in ("random_selection", "compile_task")):
                # Creates the internal transition
                transition = states_instances[state].to(
                    states_instances[state],
                    event="griptape_event",
                    on=f"on_event_{state}",
                    internal=True,
                )
                if "griptape_event" in events:
                    events["griptape_event"] |= transition
                else:
                    events["griptape_event"] = transition
        for transition in TRANSITIONS:
            for transition_data in transition["transitions"]:
                transition_value = states_instances[transition_data["from"]].to(
                    states_instances[transition_data["to"]],
                    event=transition["event"],
                    internal=False,
                )
                if transition["event"] in events:
                    events[transition["event"]] |= transition_value
                else:
                    events[transition["event"]] = transition_value
        attrs_mapper = {
            **states_instances,
            **events,
        }
        kwargs = {
            "taxonomy_choices": taxonomy_choices,
            "kb_ids": kb_ids,
            "page_range": page_range,
        }
        return cast(
            SingleQuestion,
            StateMachineMetaclass(cls.__name__, (cls,), attrs_mapper)(**kwargs),
        )

    # BENEATH ARE THE NECESSARY METHODS

    def on_enter_random_selection(self) -> None:
        # Get the random taxonomy
        self.taxonomy = random.choice(self.taxonomy_choices)
        # I changed this so I didn't have to do an "eval". Not sure how it'll work.
        taxonomy_prompt = {
            "Knowledge": "Generate a quiz question based ONLY on the information. Then write the answer to the question. The interrogative verb for the question should be randomly chosen from: 'define', 'list', 'state', 'identify','label'. INFORMATION: ",
            "Comprehension": "Generate a quiz question based ONLY on the information. Then write the answer to the question. The interrogative verb for the question should be randomly chosen from: 'explain', 'predict', 'interpret', 'infer', 'summarize', 'convert','give an example of x'. INFORMATION: ",
            "Application": "Generate a quiz question based ONLY on the information. Then write the answer to the question. The structure of the question should be randomly chosen from: 'How could x be used to y?', 'How would you show/make use of/modify/demonstrate/solve/apply x to conditions y?' INFORMATION: ",
        }
        self.taxonomy_prompt = taxonomy_prompt[self.taxonomy]
        # get the random page range and GTCVectorStoreDriver
        pages, driver = self.get_vector_store_id_from_page()
        if driver is None:
            self.send("end_state")
            self.rejected = True
            self.reject_reason = "BAD KB PAGE RANGE"
            print(self.reject_reason)
            return
        self.pages = pages
        self.driver = driver
        self.send("next_state")

    def on_enter_get_textbook(self) -> None:
        # I am going to create the agent in this method
        if "get_information" not in self._structures:
            tool = self.build_rag_tool(self.build_rag_engine(self.driver))
            use_rag_task = ToolTask(tool=tool)
            information_retriever = Agent(id="get_information", tasks=[use_rag_task])
            self._structures["get_information"] = information_retriever
        self._structures["get_information"].run("What is the information in KB?")

    def on_event_get_textbook(self, event_: dict) -> None:
        event_type = event_["type"]
        match event_type:
            case "FinishStructureRunEvent":
                structure_id = event_["structure_id"]
                match structure_id:
                    case "get_information":
                        self.information = event_["output_task_output"]["value"]
                        self.send("next_state")

    def on_enter_question_generation(self) -> None:
        if "subject_matter_expert" not in self._structures:
            rulesets = self.get_rulesets("subject_matter_expert")
            subject_matter_expert = Agent(id="subject_matter_expert", rulesets=rulesets)
            subject_matter_expert.task.output_schema = schema.Schema(
                {"Question": str, "Answer": str}
            )
            self._structures["subject_matter_expert"] = subject_matter_expert
        self._structures["subject_matter_expert"].run(
            f"{self.taxonomy_prompt}{self.information}"
        )  # TODO: Will this work the same as before

    def on_event_question_generation(self, event_: dict) -> None:
        event_type = event_["type"]
        match event_type:
            case "FinishStructureRunEvent":
                structure_id = event_["structure_id"]
                match structure_id:
                    case "subject_matter_expert":
                        question = event_["output_task_output"]["value"]
                        # save question and answer separately
                        self.question = question["Question"]
                        self.answer = question["Answer"]
                        self.send("next_state")

    def on_enter_wrong_answer_generation(self) -> None:
        if "wrong_answers_generator" not in self._structures:
            rulesets = self.get_rulesets("wrong_answers_generator")
            wrong_answers_generator = Agent(
                id="wrong_answers_generator", rulesets=rulesets
            )
            wrong_answers_generator.task.output_schema = schema.Schema(
                {"1": str, "2": str, "3": str, "4": str}
            )
            self._structures["wrong_answers_generator"] = wrong_answers_generator
        if not self.rejected:
            prompt = f"""Write and return four incorrect answers for this question: {self.question}. The correct answer to the question is: {self.answer}, and incorrect answers should have similar sentence structure to the correct answer. Write the incorrect answers from this information: {self.information}"""
        else:
            prompt = f"""Write and return four incorrect answers for this question: {self.question}. The correct answer to the question is: {self.answer}, and incorrect answers should have similar sentence structure to the correct answer. Write the incorrect answers from this information: {self.information}. Answers should not be: {self.reject_reason}."""
            print(self.reject_reason)
        self._structures["wrong_answers_generator"].run(prompt)

    def on_event_wrong_answer_generation(self, event_: dict) -> None:
        event_type = event_["type"]
        match event_type:
            case "FinishStructureRunEvent":
                structure_id = event_["structure_id"]
                match structure_id:
                    case "wrong_answers_generator":
                        wrong_answers = event_["output_task_output"]["value"]
                        wrong_answers = [wrong_answers[x] for x in ["1", "2", "3", "4"]]
                        # save question and answer separately
                        self.wrong_answers = wrong_answers
                        self.send("next_state")

    def on_enter_audit_question(self) -> None:
        if "question_auditor" not in self._structures:
            rulesets = self.get_rulesets("question_auditor")
            question_auditor = Agent(id="question_auditor", rulesets=rulesets)
            # question_auditor.task.output_schema = schema.Schema(
            # {
            #     "keep": bool,
            #     "why": schema.Optional(
            #         {
            #             "Question": bool,
            #             "Answer": bool,
            #             "Wrong Answers": bool,
            #             "Reason": str,
            #         }
            #     ),
            # }
            # )
            question_auditor.task.output_schema = schema.Schema(
                {
                    "Bad_Question": bool,
                    "Bad_Answer": bool,
                    "Bad_Wrong_Answers": bool,
                    "Reason": schema.Optional(str),
                }
            )
            self._structures["question_auditor"] = question_auditor
        # prompt = f"This is the question: {self.question}. This is the answer: {self.answer}. These are the incorrect answers:{self.wrong_answers}. This is the information given:{self.information}. IF the question is not kept, return True for the reason why from 'Question', 'Answers', 'Wrong Answers'."
        prompt = f"This is the question: {self.question}. This is the answer: {self.answer}. These are the incorrect answers:{self.wrong_answers}. This is the information given:{self.information}. IF the question is should not be kept, return True for the reason why from 'Bad_Question', 'Bad_Answer', 'Bad_Wrong_Answers'."
        self._structures["question_auditor"].run(prompt)

    def on_event_audit_question(self, event_: dict) -> None:
        event_type = event_["type"]
        match event_type:
            case "FinishStructureRunEvent":
                structure_id = event_["structure_id"]
                match structure_id:
                    case "question_auditor":
                        if self.give_up >= 3:
                            self.rejected = True
                            self.reject_reason += " \n Too many tries"
                            self.send("next_state")
                            return
                        self.give_up += 1
                        audit = event_["output_task_output"]["value"]
                        # TODO: Go back to some other state that checks the quality bar
                        # if audit["keep"]:
                        #     self.send("next_state")
                        # else:
                        #     self.rejected = True
                        #     self.reject_reason = audit["why"]["Reason"]
                        #     if audit["why"]["Question"]:
                        #         self.send("next_state")
                        #         return
                        #     if audit["why"]["Answer"]:
                        #         self.send("next_state")
                        #         return
                        #     if audit["why"]["Wrong Answers"]:
                        #         self.send(
                        #             "redo"
                        #         )  # Goes back to generate more wrong answers
                        #         return
                        #     self.send("next_state")
                        print(audit)
                        if audit["Bad_Question"]:
                            self.rejected = True
                            self.reject_reason = audit["Reason"]
                            self.reject_classification = "Bad_Question"
                            self.send("next_state")
                            return
                        if audit["Bad_Answer"]:
                            self.rejected = True
                            self.reject_reason = audit["Reason"]
                            self.reject_classification = "Bad_Answer"
                            self.send("next_state")
                            return
                        if audit["Bad_Wrong_Answers"]:
                            self.rejected = True
                            self.reject_reason = audit["Reason"]
                            self.reject_classification = "Bad_Wrong_Answers"
                            self.send("redo")
                            return
                        self.rejected = False
                        self.send("next_state")

    def on_enter_compile_task(self) -> None:
        # TODO: Logic to determine if I should go back to wrong answers
        question = {
            "Question": self.question,
            "Answer": self.answer,
            "Wrong Answers": self.wrong_answers,
            "Page": self.pages,
            "Taxonomy": self.taxonomy,
        }
        if self.rejected:
            question["Reject Classification"] = self.reject_classification
            question["Reason"] = self.reject_reason
        self.generated_question = question
        self.send("next_state")

    # TODO : Does this return output
    def on_enter_end(self) -> dict:
        return self.generated_question

    # HELPER METHODS BELOW

    def get_rulesets(self, structure_id: str) -> list:
        final_ruleset_list = []
        ruleset_ids = STRUCTURES[structure_id]["ruleset_ids"]
        for ruleset_id in ruleset_ids:
            ruleset_rules = RULESETS[ruleset_id]
            rules = [Rule(rule) for rule in ruleset_rules]
            final_ruleset_list.append(Ruleset(ruleset_id, rules=rules))
        return final_ruleset_list

    def get_vector_store_id_from_page(
        self,
    ) -> tuple[str, GriptapeCloudVectorStoreDriver | None]:
        possible_kbs = {}
        for name, kb_id in self.kb_ids.items():
            page_nums = name.split("p")[1:]
            start_page = int(page_nums[0].split("-")[0])
            end_page = int(page_nums[1])
            if end_page <= self.page_range[1] and start_page >= self.page_range[0]:
                possible_kbs[kb_id] = f"{start_page}-{end_page}"
        if not len(list(possible_kbs.keys())):
            return ("No KBs in range", None)
        kb_id = random.choice(list(possible_kbs.keys()))
        page_value = possible_kbs[kb_id]
        return page_value, GriptapeCloudVectorStoreDriver(
            api_key=os.getenv("GT_CLOUD_API_KEY", ""),
            knowledge_base_id=kb_id,
        )

    # Uses this and all below to build the Rag Tool to get information from the KB
    def build_rag_engine(
        self, vector_store_driver: GriptapeCloudVectorStoreDriver
    ) -> RagEngine:
        return RagEngine(
            retrieval_stage=RetrievalRagStage(
                retrieval_modules=[
                    VectorStoreRetrievalRagModule(
                        vector_store_driver=vector_store_driver,
                    )
                ],
            ),
            response_stage=ResponseRagStage(
                response_modules=[TextChunksResponseRagModule()]
            ),
        )

    def build_rag_tool(self, engine: RagEngine) -> RagTool:
        return RagTool(
            description="Contains information about the textbook. Use it ONLY for context.",
            rag_engine=engine,
        )

    def make_rag_structure(
        self, vector_store: GriptapeCloudVectorStoreDriver
    ) -> Structure:
        if vector_store:
            tool = self.build_rag_tool(self.build_rag_engine(vector_store))
            use_rag_task = ToolTask(tool=tool)
            return Agent(tasks=[use_rag_task])
        errormsg = "No Vector Store"
        raise ValueError(errormsg)


if __name__ == "__main__":
    flow = SingleQuestion.create_statemachine(
        ["Comprehension"],
        {"p126-p129": "9efbb8ab-6a5e-4bca-aab0-7f7500bfb7b5"},
        (120, 150),
    )
    flow.send("start_up")
    # When incorporating into the main flow - we can just get the result of flow.generated_question and use that value onwards
    # TODO: Do any events need to be sent?
    print(flow.generated_question)