from __future__ import annotations import json import logging import os import random from abc import abstractmethod from pathlib import Path from typing import TYPE_CHECKING, cast import requests from dotenv import load_dotenv from griptape.artifacts import ListArtifact, TextArtifact from griptape.configs import Defaults from griptape.configs.drivers import ( OpenAiDriversConfig, ) from griptape.drivers import ( GriptapeCloudVectorStoreDriver, LocalStructureRunDriver, 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.memory.structure import ConversationMemory from griptape.rules import Rule, Ruleset from griptape.structures import Agent, Workflow from griptape.tasks import CodeExecutionTask, StructureRunTask, ToolTask from griptape.tools import RagTool from parsers import UWConfigParser from statemachine import State, StateMachine from statemachine.factory import StateMachineMetaclass logger = logging.getLogger(__name__) logging.getLogger("griptape").setLevel(logging.ERROR) if TYPE_CHECKING: from griptape.structures import Structure from griptape.tools import BaseTool from statemachine.event import Event load_dotenv() # Sets max tokens and OpenAI as the driver. Defaults.drivers_config = OpenAiDriversConfig( prompt_driver=OpenAiChatPromptDriver(model="gpt-4o", max_tokens=4096) ) def custom_dict_merge(dict1: dict, dict2: dict) -> dict: result = dict1.copy() for key, value in dict2.items(): if key in result and isinstance(result[key], list) and isinstance(value, list): result[key] = result[key] + value else: result[key] = value return result class UWBaseMachine(StateMachine): """Base class for a machine. Attributes: config_file (Path): The path to the configuration file. config (dict): The configuration data. outputs_to_user (list[str]): Outputs to return to the user. """ def __init__(self, config_file: Path, **kwargs) -> None: self.config_parser = UWConfigParser(config_file) self.config = self.config_parser.parse() self._structures = {} self.vector_stores = {} # Store here in case needs multiple uses self.question_list: list = [] # For the parameters necessary from the user self.page_range: tuple = () self.question_number: int = 0 self.taxonomy: list = [] # To track give up self.give_up_count = 0 self.current_question_count = 0 # To keep vector stores on track self.kb_ids = {} self.state_status: dict[str, bool] = {} for key in self.state_transitions: self.state_status[key] = False def on_event(event: BaseEvent) -> None: """Takes in griptape events from eventbus and fixes them.""" print(f"Received Griptape event: {json.dumps(event.to_dict(), indent=2)}") try: self.send( "process_event", event_={"type": "griptape_event", "value": event.to_dict()}, ) except Exception as e: errormsg = f"Would not allow process_event to be sent. Check to see if it is defined in the config.yaml. Error:{e}" raise ValueError(errormsg) from e EventBus.clear_event_listeners() EventBus.add_event_listener( EventListener(on_event, event_types=[FinishStructureRunEvent]), ) super().__init__() @property def available_events(self) -> list[str]: return self.current_state.transitions.unique_events @property @abstractmethod def tools(self) -> dict[str, BaseTool]: """Returns the Tools for the machine.""" ... @property def _current_state_config(self) -> dict: return self.config["states"][self.current_state_value] @classmethod def from_definition( # noqa: C901, PLR0912 cls, definition: dict, **extra_kwargs ) -> UWBaseMachine: try: states_instances = {} for state_id, state_kwargs in definition["states"].items(): # These are the relevant states that need GOAP. states_instances[state_id] = State(**state_kwargs, value=state_id) except Exception as e: errormsg = f"""Error in state definition: {e}. """ raise ValueError(errormsg) from e events = {} state_transitions = {} for event_name, transitions in definition["events"].items(): for transition_data in transitions: try: source_name = transition_data["from"] source = states_instances[source_name] target = states_instances[transition_data["to"]] relevance = "" if "relevance" in transition_data: relevance = transition_data["relevance"] if source_name not in state_transitions: state_transitions[source_name] = {event_name: relevance} else: state_transitions[source_name][event_name] = relevance except Exception as e: errormsg = f"Error:{e}. Please check your transitions to be sure each transition has a source and destination." raise ValueError(errormsg) from e transition = source.to( target, event=event_name, cond=transition_data.get("cond"), unless=transition_data.get("unless"), on=transition_data.get("on"), internal=transition_data.get("internal"), ) if event_name in events: events[event_name] |= transition else: events[event_name] = transition for state_id, state in states_instances.items(): if state_id not in ("end", "start"): transition = state.to( state, event="process_event", on=f"on_event_{state_id}", internal=True, ) if "process_event" in events: events["process_event"] |= transition else: events["process_event"] = transition attrs_mapper = { **extra_kwargs, **states_instances, **events, "state_transitions": state_transitions, } return cast( UWBaseMachine, StateMachineMetaclass(cls.__name__, (cls,), attrs_mapper)(**extra_kwargs), ) @classmethod def from_config_file( cls, config_file: Path, **extra_kwargs, ) -> UWBaseMachine: """Creates a StateMachine class from a configuration file""" config_parser = UWConfigParser(config_file) config = config_parser.parse() extra_kwargs["config_file"] = config_file definition_states = { state_id: { "initial": state_value.get("initial", False), "final": state_value.get("final", False), } for state_id, state_value in config["states"].items() } definition_events = { event_name: list(event_value["transitions"]) for event_name, event_value in config["events"].items() } definition = {"states": definition_states, "events": definition_events} return cls.from_definition(definition, **extra_kwargs) @abstractmethod def start_machine(self) -> None: """Starts the machine.""" ... def reset_structures(self) -> None: """Resets the structures.""" self._structures = {} def on_enter_state(self, source: State, state: State, event: Event) -> None: print(f"Transitioning from {source} to {state} with event {event}") def get_structure(self, structure_id: str) -> Structure: global_structure_config = self.config["structures"][structure_id] state_structure_config = self._current_state_config.get("structures", {}).get( structure_id, {} ) structure_config = custom_dict_merge( global_structure_config, state_structure_config ) if structure_id not in self._structures: # Initialize Structure with all the expensive setup structure = Agent( id=structure_id, conversation_memory=ConversationMemory(), ) self._structures[structure_id] = structure # Create a new clone with state-specific stuff structure = self._structures[structure_id] structure = Agent( id=structure.id, prompt_driver=structure.prompt_driver, conversation_memory=structure.conversation_memory, rulesets=[ *self._get_structure_rulesets(structure_config.get("ruleset_ids", [])), ], ) print(f"Structure: {structure_id}") for ruleset in structure.rulesets: for rule in ruleset.rules: print(f"Rule: {rule.value}") return structure def _get_structure_rulesets(self, ruleset_ids: list[str]) -> list[Ruleset]: ruleset_configs = [ self.config["rulesets"][ruleset_id] for ruleset_id in ruleset_ids ] # Convert ruleset configs to Rulesets return [ Ruleset( name=ruleset_config["name"], rules=[Rule(rule) for rule in ruleset_config["rules"]], ) for ruleset_config in ruleset_configs ] def retrieve_vector_stores(self) -> None: base_url = "https://cloud.griptape.ai/api/" kb_url = f"{base_url}/knowledge-bases" headers = {"Authorization": f"Bearer {os.getenv('GT_CLOUD_API_KEY')}"} response = requests.get(url=kb_url, headers=headers) response.raise_for_status() all_kbs = {} if response.status_code == 200: data = response.json() for kb in data["knowledge_bases"]: name = kb["name"] kb_id = kb["knowledge_base_id"] if "KB_section" in name: all_kbs[name] = kb_id else: raise ValueError(response.status_code) self.kb_ids = all_kbs # ALL METHODS RELATING TO THE WORKFLOW AND PIPELINE ARE BELOW THIS LINE # This is the overarching workflow. Creates a workflow with get_single_question x amount of times. def get_questions_workflow(self) -> Workflow: workflow = Workflow(id="create_question_workflow") # How many questions still need to be created? for _ in range(self.question_number - len(self.question_list)): task = StructureRunTask( structure_run_driver=LocalStructureRunDriver( create_structure=self.get_single_question ), child_ids=["end_task"], ) # Create X amount of workflows to run for X amount of questions needed. workflow.add_task(task) end_task = CodeExecutionTask(id="end_task", on_run=self.end_workflow) workflow.add_task(end_task) return workflow # Ends the get_questions_workflow. Compiles all workflow outputs into one output. def end_workflow(self, task: CodeExecutionTask) -> ListArtifact: parent_outputs = task.parent_outputs questions = [] for outputs in parent_outputs.values(): if outputs.type == "InfoArtifact": continue questions.append(outputs) return ListArtifact(questions) # Generates one workflow to create a single question. def get_single_question(self) -> Workflow: question_generator = Workflow(id="single_question") taxonomy = random.choice(self.taxonomy) taxonomyprompt = { "Knowledge": "Generate a quiz question based ONLY on this information: {{parent_outputs['information_task']}}, then write the answer to the question. The interrogative verb for the question should be randomly chosen from: 'define', 'list', 'state', 'identify','label'.", "Comprehension": "Generate a quiz question based ONLY on this information: {{parent_outputs['information_task']}}, 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'.", "Application": "Generate a quiz question based ONLY on this information: {{parent_outputs['information_task']}}, 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?'", } pages, driver = self.get_vector_store_id_from_page() get_information = StructureRunTask( id="information_task", input="What is the information in KB?", structure_run_driver=LocalStructureRunDriver( create_structure=lambda: self.make_rag_structure(driver) ), child_ids=["get_question"], ) # Get KBs and select it, assign it to the structure or create the structure right here. # Rules for subject matter expert: return only a json with question and answer as keys. generate_q_task = StructureRunTask( id="get_question", input=taxonomyprompt[taxonomy], structure_run_driver=LocalStructureRunDriver( create_structure=lambda: self.get_structure("subject_matter_expert") ), parent_ids=["information_task"], ) get_question_code_task = CodeExecutionTask( id="get_only_question", on_run=self.get_question_for_wrong_answers, parent_ids=["get_question"], child_ids=["wrong_answers"], ) get_separated_answer_code_task = CodeExecutionTask( id="get_separated_answer", on_run=self.get_separated_answer_for_wrong_answers, parent_ids=["get_question"], child_ids=["wrong_answers"], ) generate_wrong_answers = StructureRunTask( id="wrong_answers", input="""Write and return three incorrect answers for this question: {{parent_outputs['get_separated_question']}}. The correct answer to the question is: {{parent_outputs['get_separated_answer']}}, and incorrect answers should have similar sentence structure to the correct answer. Write the incorrect answers from this information: {{parent_outputs['information_task']}}""", structure_run_driver=LocalStructureRunDriver( create_structure=lambda: self.get_structure("wrong_answers_generator") ), parent_ids=["get_only_question", "information_task"], ) compile_task = CodeExecutionTask( id="compile_task", input=f"{pages}, {taxonomy}", on_run=self.single_question_last_task, parent_ids=["wrong_answers", "get_question"], ) question_generator.add_tasks( get_information, generate_q_task, get_question_code_task, get_separated_answer_code_task, generate_wrong_answers, compile_task, ) return question_generator # Task to separate the Question into a string def get_question_for_wrong_answers(self, task: CodeExecutionTask) -> TextArtifact: parent_outputs = task.parent_outputs question = parent_outputs["get_question"].value question = json.loads(question)["Question"] return TextArtifact(question) # Task to separate the Answer into a string def get_separated_answer_for_wrong_answers( self, task: CodeExecutionTask ) -> TextArtifact: parent_outputs = task.parent_outputs answer = parent_outputs["get_question"].value print(answer) answer = json.loads(answer)["Answer"] return TextArtifact(answer) # Combines all the outputs into one dictionary that represents the question def single_question_last_task(self, task: CodeExecutionTask) -> TextArtifact: parent_outputs = task.parent_outputs wrong_answers = parent_outputs["wrong_answers"].value # Output is a list wrong_answers = wrong_answers.split("\n") question_and_answer = parent_outputs["get_question"].value # Output is a json try: question_and_answer = json.loads(question_and_answer) except: question_and_answer = question_and_answer.split("\n")[1:] question_and_answer = "".join(question_and_answer) question_and_answer = json.loads(question_and_answer) inputs = task.input.value.split(",") question = { "Question": question_and_answer["Question"], "Answer": question_and_answer["Answer"], "Wrong Answers": wrong_answers, "Page": inputs[0], "Taxonomy": inputs[1], } return TextArtifact(question) # These are helper methods # Picks the KB from the dictionary def get_vector_store_id_from_page( self, ) -> tuple[str, GriptapeCloudVectorStoreDriver]: 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}" 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)