Init files

.gitignore

@@ -0,0 +1,197 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# UV
+#   Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#uv.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/latest/usage/project/#working-with-version-control
+.pdm.toml
+.pdm-python
+.pdm-build/
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+# Abstra
+# Abstra is an AI-powered process automation framework.
+# Ignore directories containing user credentials, local state, and settings.
+# Learn more at https://abstra.io/docs
+.abstra/
+
+# Visual Studio Code
+#  Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore 
+#  that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
+#  and can be added to the global gitignore or merged into this file. However, if you prefer, 
+#  you could uncomment the following to ignore the enitre vscode folder
+# .vscode/
+
+# Ruff stuff:
+.ruff_cache/
+
+# PyPI configuration file
+.pypirc
+
+# Cursor
+#  Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
+#  exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
+#  refer to https://docs.cursor.com/context/ignore-files
+.cursorignore
+.cursorindexingignore
+
+# ignore session dirs
+sessions/*

app.py

@@ -0,0 +1,23 @@
+
+import gradio as gr
+
+from src.frontend import (
+    chatbot,
+    continue_story_screen,
+    progress_screen,
+    sidebar,
+    start_screen,
+    story_information_widgets
+)
+
+demo = gr.Blocks()
+
+with demo:
+    start_screen.render()
+    story_information_widgets.render()
+    continue_story_screen.render()
+    sidebar.render()
+    chatbot.render()
+    progress_screen.render()
+
+demo.launch()

requirements.txt

@@ -0,0 +1,5 @@
+openai
+langgraph
+gradio
+igraph
+plotly

src/agents/base_agent.py

@@ -0,0 +1,40 @@
+
+from abc import ABC
+from typing import Dict, List, Literal
+
+from ..utils.init_openai import completions_with_backoff
+
+
+class BaseAgent(ABC):
+
+    def format_openai_message(
+        self,
+        system_message: str,
+        user_message: str
+    ) -> List[Dict[str, str]]:
+        return [
+            {
+                'role': 'system',
+                'content': system_message
+            },
+            {
+                'role': 'user',
+                'content': user_message
+            }
+        ]
+
+    def execute(
+        self,
+        model: Literal['openai'],
+        messages: List[Dict[str, str]],
+        temperature: float = 0,
+        max_tokens: int = 8192
+    ) -> str:
+        if model == 'openai':
+
+            response = completions_with_backoff(messages, temperature, max_tokens)
+            content = response.choices[0].message.content
+
+            return content
+        else:
+            raise ValueError(f'Model type {model} not supported (yet)')

src/agents/category_builder_agent.py

@@ -0,0 +1,46 @@
+
+import json
+
+from typing import List
+
+from .base_agent import BaseAgent
+from ..dataclasses.category import Category
+from ..prompts.category_builder_agent import system_prompt_template, user_prompt_template
+
+
+class CategoryBuilderAgent(BaseAgent):
+
+    def __init__(self):
+        self.model = 'openai'
+        self.temperature = 0.5
+
+    def get_categories(
+        self,
+        story_context: str,
+        category_context: str,
+        num_categories: int
+    ) -> List[Category]:
+        prompt_arguments = {
+            'story_context': story_context,
+            'category_context': category_context,
+            'num_categories': num_categories
+        }
+
+        system_message = system_prompt_template.format(**prompt_arguments)
+        user_message = user_prompt_template.format(**prompt_arguments)
+
+        messages = self.format_openai_message(system_message, user_message)
+
+        response = self.execute(self.model, messages, self.temperature)
+
+        # remove json tag from response
+        response = response.removeprefix('```json').removeprefix('```')
+        response = response.removesuffix('```json').removesuffix('```')
+
+        categories = json.loads(response)
+        assert type(categories) == list
+
+        return [
+            Category(**category)
+            for category in categories
+        ]

src/agents/coordinator.py

@@ -0,0 +1,190 @@
+
+import json
+import os
+
+from langgraph.graph import StateGraph, START, END
+from langgraph.graph.graph import CompiledGraph
+from typing import Dict, List, Literal, Tuple
+
+from src.agents.category_builder_agent import CategoryBuilderAgent
+from src.agents.evaluator_agent import EvaluatorAgent
+from src.agents.narrator_agent import NarratorAgent
+from src.agents.orchestrator_agent import OrchestratorAgent
+from src.agents.storybuilder_agent import StorybuilderAgent
+from src.dataclasses.category import Category
+from src.dataclasses.process_state import ProcessState
+from src.dataclasses.question import Question
+from src.utils.utils import get_session_dir, save_information, load_information
+
+num_questions_: int = None
+num_options_: int = None
+story_name: str = None
+states: Dict[int, Dict[str, str]] = dict()  # maps question_num to {str to str}
+questions: Dict[str, Question] = dict()  # maps uuid to Question
+state_question_map: Dict[int, Dict[str, str]] = dict() # maps question_num to {state_name to question_uuid}
+categories: List[Category] = []  # list of categories
+categories_seen: Dict[str, bool] = dict()  # maps category name to bool if category name was seen
+path_evaluations: Dict[str, List[str]] = dict()  # maps a path to an evaluation, including the reason
+
+
+def init_categories(
+    story_context: str,
+    category_context: str,
+    num_categories: int
+) -> None:
+    global categories, categories_seen
+
+    category_builder_agent = CategoryBuilderAgent()
+    categories = category_builder_agent.get_categories(
+        story_context, category_context, num_categories
+    )
+    categories_seen = {
+        category.name: False
+        for category in categories
+    }
+
+
+def user_node(state: ProcessState) -> ProcessState:
+    user_choice = state['user_choice']
+    question_uuid = state['foll_question_uuid']
+    question_num = state['question_num'] + 1
+
+    if user_choice:
+        foll_state_name = questions[question_uuid].options[user_choice]
+        foll_user_choices = [choice for choice in state['previous_user_choices']]
+        foll_user_choices.append(user_choice)
+    else:
+        foll_state_name = 'start'
+        foll_user_choices = []
+
+    return {
+        'question_num': question_num,
+        'state_name': foll_state_name,
+        'user_choice': user_choice,
+        'previous_state_names': state['previous_state_names'],
+        'previous_questions': state['previous_questions'],
+        'previous_user_choices': foll_user_choices,
+        'builder_instruction': '',
+        'foll_question_uuid': ''
+    }
+
+
+def init_graph(
+    story_context: str,
+    category_context: str,
+    num_questions: int,
+    num_options: int,
+    num_categories: int
+) -> CompiledGraph:
+    
+    global num_questions_, num_options_
+    num_questions_ = num_questions
+    num_options_ = num_options
+
+    orchestrator_agent = OrchestratorAgent(
+        story_context, category_context, num_questions
+    )
+    storybuilder_agent = StorybuilderAgent(
+        story_context, category_context, num_questions, num_options
+    )
+    narrator_agent = NarratorAgent(
+        story_context, category_context, num_questions, num_options
+    )
+
+    init_categories(story_context, category_context, num_categories)
+
+    builder = StateGraph(ProcessState)
+
+    builder.add_node('orchestrator', orchestrator_agent.run_agent)
+    builder.add_node('storybuilder', storybuilder_agent.run_agent)
+    builder.add_node('narrator', narrator_agent.run_agent)
+    builder.add_node('user', user_node)
+
+    def conditional_edge(state: ProcessState) -> Literal['storybuilder', END]:
+        question_num = state['question_num']
+        state_name = state['state_name']
+
+        if state_name not in states.get(question_num, dict()) or \
+            states[question_num][state_name] == 'Unexplored!':
+            return 'storybuilder'
+        else:
+            return END
+        
+    def last_action_edge(state: ProcessState) -> Literal['orchestrator', END]:
+        question_num = state['question_num']
+
+        if question_num > num_questions:
+            return END
+        else:
+            return 'orchestrator'
+
+    builder.add_edge(START, 'user')
+    builder.add_edge('user', 'orchestrator')
+    builder.add_edge('storybuilder', 'narrator')
+    builder.add_edge('narrator', END)
+    builder.add_conditional_edges('orchestrator', conditional_edge)
+    builder.add_conditional_edges('user', last_action_edge)
+    graph = builder.compile()
+
+    return graph
+
+def evaluation(state: ProcessState) -> Tuple[str, str]:
+
+    evaluator_agent = EvaluatorAgent()
+
+    previous_state_names = state['previous_state_names']
+    previous_questions = state['previous_questions']
+    previous_user_choices = state['previous_user_choices']
+
+    previous_questions.append(state['foll_question_uuid'])
+    previous_user_choices.append(state['user_choice'])
+
+    selected_category, reasoning = evaluator_agent.get_evaluation(
+        previous_state_names, previous_questions, previous_user_choices,
+        states, questions, categories
+    )
+
+    return selected_category, reasoning
+
+
+def save_coordinator() -> None:
+    data = [
+        states,
+        {k: v.dict() for k, v in questions.items()},
+        state_question_map,
+        [c.dict() for c in categories],
+        categories_seen,
+        path_evaluations,
+    ]
+    key_names = [
+        'states',
+        'questions',
+        'state_question_map',
+        'categories',
+        'categories_seen',
+        'path_evaluations'
+    ]
+
+    combined_data = {
+        key_name: data_
+        for key_name, data_ in zip(key_names, data)
+    }
+    filepath = os.path.join(get_session_dir(), story_name, 'data.json')
+    save_information(combined_data, filepath)
+
+
+def load_coordinator() -> None:
+    global states, questions, state_question_map, categories, categories_seen, path_evaluations
+
+    filepath = os.path.join(get_session_dir(), story_name, 'data.json')
+    combined_data = load_information(filepath)
+
+    states = {int(k): v for k, v in combined_data['states'].items()}
+    questions = {k: Question(**{
+            'question': v['question'],
+            'options': json.loads(v['options'].replace('\'', '\"'))
+        }) for k, v in combined_data['questions'].items()}
+    state_question_map = {int(k): v for k, v in combined_data['state_question_map'].items()}
+    categories = [Category(**c) for c in combined_data['categories']]
+    categories_seen = combined_data['categories_seen']
+    path_evaluations = combined_data['path_evaluations']

src/agents/evaluator_agent.py

@@ -0,0 +1,127 @@
+
+import json
+
+from typing import Dict, List, Tuple
+
+import src.agents.coordinator as C
+
+from .base_agent import BaseAgent
+from ..dataclasses.category import Category
+from ..prompts.evaluator_agent import system_prompt_template, user_prompt_template
+
+CATEGORY_FORMAT = """
+    {{
+        'name': {category_name},
+        'description': {category_description},
+        'traits': [
+            {category_traits}
+        ]
+    }}
+""".strip()
+
+
+class EvaluatorAgent(BaseAgent):
+
+    def __init__(self):
+        self.model = 'openai'
+        self.temperature = 0.2
+
+    def generate_path_key(
+        previous_questions: List[str],
+        previous_user_choices: List[str]
+    ) -> str:
+        return '-'.join([
+            f'{question}+{choice}'
+            for question, choice in zip(previous_questions, previous_user_choices)
+        ])
+
+    def is_encourage_categories(
+        self,
+        questions: Dict[str, str]
+    ) -> bool:
+        '''Attempts to ensure that each category has at least one story path to it.
+
+        Computes if the number of remaining paths equals to the number of unused categories.
+        Assumed to be a low chance as the number of questions increases, but exists as a checker.
+        '''
+        if all([v for _, v in C.categories_seen.items()]):
+            return False
+        
+        num_questions = len(questions)
+        random_question = list(questions.keys())[0]
+        num_options = len(C.questions[random_question].options)
+
+        rem_categories = sum([1 for _, v in C.categories_seen.items() if not v])
+        num_categories = len(C.categories)
+        return rem_categories >= num_options * num_questions - (num_categories - rem_categories)
+
+    def get_evaluation(
+        self,
+        previous_state_names: List[str],
+        previous_questions: List[str],
+        previous_user_choices: List[str],
+        states: Dict[int, Dict[str, str]],
+        questions: Dict[str, str],
+        categories: List[Category]
+    ) -> List[str]:
+        
+        complete_user_story = []
+
+        path_key = EvaluatorAgent.generate_path_key(previous_questions, previous_user_choices)
+
+        if path_key in C.path_evaluations:
+            return C.path_evaluations[path_key]
+
+        for i, (state_name, question_uuid, user_choice) in enumerate(zip(
+            previous_state_names, previous_questions, previous_user_choices
+        )):
+            scenario = states[i + 1][state_name]
+            question_str = questions[question_uuid].question
+
+            turn_content = f'`SCENARIO`: {scenario}\n' + \
+                f'`QUESTION`: {question_str}\n' + \
+                f'`CHOICE`: {user_choice}\n'
+            
+            complete_user_story.append(turn_content)
+        
+        complete_user_story = '\n\n'.join(complete_user_story)
+
+        is_encourage_categories = self.is_encourage_categories(questions)
+        category_contents = []
+        for i, category in enumerate(sorted(categories, key=lambda x: x.name)):
+            if is_encourage_categories and C.categories_seen[category.name]:
+                continue
+
+            category_content = CATEGORY_FORMAT.format(**{
+                'category_name': category.name,
+                'category_description': category.description,
+                'category_traits': category.traits
+            })
+            category_contents.append(category_content)
+
+        category_contents = '{\n' + ',\n'.join(category_contents) + '\n}'
+
+        prompt_arguments = {
+            'user_story': complete_user_story,
+            'category_contents': category_contents
+        }
+            
+        system_message = system_prompt_template.format(**prompt_arguments)
+        user_message = user_prompt_template.format(**prompt_arguments)
+
+        messages = self.format_openai_message(system_message, user_message)
+
+        response = self.execute(self.model, messages, self.temperature)
+
+        # remove json tag from response
+        response = response.removeprefix('```json').removeprefix('```')
+        response = response.removesuffix('```json').removesuffix('```')
+
+        selection_info = json.loads(response)
+
+        selected_category = selection_info['category']
+        reason = selection_info['reason']
+
+        C.path_evaluations[path_key] = [selected_category, reason]
+
+        return (selected_category, reason)

src/agents/narrator_agent.py

@@ -0,0 +1,98 @@
+
+import json
+
+from uuid import uuid4
+
+import src.agents.coordinator as C
+
+from .base_agent import BaseAgent
+from ..dataclasses.process_state import ProcessState
+from ..dataclasses.question import Question
+from ..prompts.narrator_agent import system_prompt_template, user_prompt_template
+
+class NarratorAgent(BaseAgent):
+
+    def __init__(
+        self,
+        story_context: str,
+        category_context: str,
+        num_questions: int,
+        num_options: int
+    ):
+        
+        self.prompt_arguments = {
+            'story_context': story_context,
+            'category_context': category_context,
+            'num_questions': num_questions,
+            'num_options': num_options
+        }
+
+        self.model = 'openai'
+        self.temperature = 0.5
+
+    def run_agent(self, state: ProcessState) -> ProcessState:
+
+        question_num = state['question_num']
+        state_name = state['state_name']
+
+        if state_name in C.state_question_map.get(question_num, dict()):
+            question_uuid = C.state_question_map[question_num][state_name]
+
+        else:
+            question_uuid = str(uuid4())
+            if question_num not in C.state_question_map:
+                C.state_question_map[question_num] = dict()
+            C.state_question_map[question_num][state_name] = question_uuid
+
+            self.prompt_arguments['user_story'] = '\n'.join([
+                C.states[i + 1][state]
+                for i, state in enumerate(state['previous_state_names'][:-1])
+            ])
+            self.prompt_arguments['scenario'] = C.states[question_num][state_name]
+            self.prompt_arguments['builder_instruction'] = state['builder_instruction']
+            if question_num in C.states:
+                self.prompt_arguments['existing_scenarios'] = '\n'.join([
+                    '{\n' + f'scenario_name: {state_name}\n' + \
+                        f'scenario: {scenario}' + '}\n'
+                    for state_name, scenario in C.states.get(question_num + 1, dict()).items()
+                ])
+            else:
+                self.prompt_arguments['existing_scenarios'] = ''
+
+            system_message = system_prompt_template.format(**self.prompt_arguments)
+            user_message = user_prompt_template.format(**self.prompt_arguments)
+
+            messages = self.format_openai_message(system_message, user_message)
+
+            response = self.execute(self.model, messages, self.temperature)
+
+            # remove json tag from response
+            response = response.removeprefix('```json').removeprefix('```')
+            response = response.removesuffix('```json').removesuffix('```')
+
+            question_info = json.loads(response)
+
+            if isinstance(question_info, list):
+                question_info = question_info[0]
+
+            question = Question(**question_info)
+            for foll_state_name in question.options.values():
+                if question_num + 1 not in C.states:
+                    C.states[question_num + 1] = dict()
+                C.states[question_num + 1][foll_state_name] = 'Unexplored!'
+
+            C.questions[question_uuid] = question
+
+        foll_questions = [question for question in state['previous_questions']]
+        foll_questions.append(question_uuid)
+
+        return {
+            'question_num': state['question_num'],
+            'state_name': state['state_name'],
+            'user_choice': state['user_choice'],
+            'previous_state_names': state['previous_state_names'],
+            'previous_questions': foll_questions,
+            'previous_user_choices': state['previous_user_choices'],
+            'builder_instruction': state['builder_instruction'],
+            'foll_question_uuid': question_uuid
+        }

src/agents/orchestrator_agent.py

@@ -0,0 +1,85 @@
+
+import src.agents.coordinator as C
+
+from .base_agent import BaseAgent
+from .category_builder_agent import CategoryBuilderAgent
+from ..dataclasses.process_state import ProcessState
+from ..prompts.orchestrator_agent import (
+    start_system_prompt_template,
+    start_user_prompt_template,
+    coordinate_system_prompt_template,
+    coordinate_user_prompt_template,
+    get_phase_message
+)
+
+class OrchestratorAgent(BaseAgent):
+
+    def __init__(
+        self,
+        story_context: str,
+        category_context: str,
+        num_questions: int
+    ):
+        # build prompt argument
+        self.prompt_arguments = {
+            'story_context': story_context,
+            'category_context': category_context,
+            'num_questions': num_questions
+        }
+
+        # openai calls
+        self.model = 'openai'
+        self.temperature = 0.2
+
+    def run_agent(self, state: ProcessState) -> ProcessState:
+
+        question_num = state['question_num']
+        state_name = state['state_name']
+
+        # check which system and user prompts to build
+        if question_num == 1:
+            system_message = start_system_prompt_template.format(**self.prompt_arguments)
+            user_message = start_user_prompt_template.format(**self.prompt_arguments)
+        else:
+            progress = question_num / self.prompt_arguments['num_questions']
+            self.prompt_arguments['phase_message'] = get_phase_message(progress)
+            self.prompt_arguments['user_story'] = '\n'.join([
+                C.states[i + 1][state]
+                for i, state in enumerate(state['previous_state_names'])
+            ])
+            system_message = coordinate_system_prompt_template.format(**self.prompt_arguments)
+            user_message = coordinate_user_prompt_template.format(**self.prompt_arguments)
+
+        # check if the state has been generated before
+        if state_name in C.states.get(question_num, dict()) and C.states[question_num][state_name] != 'Unexplored!':
+            foll_state_names = [name for name in state['previous_state_names']]
+            foll_state_names.append(state['state_name'])
+
+            builder_instruction = ''
+
+            foll_question_uuid = C.state_question_map[question_num][state_name]
+            previous_questions = [question_uuid for question_uuid in state['previous_questions']]
+            if state['foll_question_uuid'] or question_num == 1:
+                if state['foll_question_uuid']:
+                    previous_questions.append(state['foll_question_uuid'])
+                else:
+                    previous_questions.append(C.state_question_map[1]['start'])
+            
+        else:
+            foll_state_names = state['previous_state_names']
+            
+            messages = self.format_openai_message(system_message, user_message)
+            builder_instruction = self.execute(self.model, messages, self.temperature)
+            foll_question_uuid = state['foll_question_uuid']
+            previous_questions = [question_uuid for question_uuid in state['previous_questions']]
+
+        return {
+            'question_num': state['question_num'],
+            'state_name': state['state_name'],
+            'user_choice': state['user_choice'],
+            'previous_state_names': foll_state_names,
+            'previous_questions': previous_questions,
+            'previous_user_choices': state['previous_user_choices'],
+            'builder_instruction': builder_instruction,
+            'foll_question_uuid': foll_question_uuid
+        }

src/agents/storybuilder_agent.py

@@ -0,0 +1,69 @@
+
+import src.agents.coordinator as C
+
+from .base_agent import BaseAgent
+from ..dataclasses.process_state import ProcessState
+from ..prompts.storybuilder_agent import system_prompt_template, user_prompt_template
+
+class StorybuilderAgent(BaseAgent):
+
+    def __init__(
+        self,
+        story_context: str,
+        category_context: str,
+        num_questions: int,
+        num_options: int
+    ):
+        
+        self.prompt_arguments = {
+            'story_context': story_context,
+            'category_context': category_context,
+            'num_questions': num_questions,
+            'num_options': num_options
+        }
+
+        self.model = 'openai'
+        self.temperature = 0.5
+
+    def run_agent(self, state: ProcessState) -> ProcessState:
+
+        question_num = state['question_num']
+        state_name = state['state_name']
+
+        self.prompt_arguments['builder_instruction'] = state['builder_instruction']
+        self.prompt_arguments['user_story'] = '\n'.join([
+            C.states[i + 1][state]
+            for i, state in enumerate(state['previous_state_names'])
+        ])
+
+        if question_num == 1:
+            self.prompt_arguments['question'] = ''
+            self.prompt_arguments['response'] = ''
+        else:
+            self.prompt_arguments['question'] = C.questions[state['previous_questions'][-1]].question
+            self.prompt_arguments['response'] = state['user_choice']
+
+        system_message = system_prompt_template.format(**self.prompt_arguments)
+        user_message = user_prompt_template.format(**self.prompt_arguments)
+
+        messages = self.format_openai_message(system_message, user_message)
+
+        scenario = self.execute(self.model, messages, self.temperature)
+
+        if question_num not in C.states:
+            C.states[question_num] = dict()
+        C.states[question_num][state_name] = scenario
+
+        foll_state_names = [name for name in state['previous_state_names']]
+        foll_state_names.append(state_name)
+
+        return {
+            'question_num': state['question_num'],
+            'state_name': state['state_name'],
+            'user_choice': state['user_choice'],
+            'previous_state_names': foll_state_names,
+            'previous_questions': state['previous_questions'],
+            'previous_user_choices': state['previous_user_choices'],
+            'builder_instruction': state['builder_instruction'],
+            'foll_question_uuid': state['foll_question_uuid']
+        }

src/dataclasses/category.py

@@ -0,0 +1,13 @@
+
+from dataclasses import asdict, dataclass
+from typing import List
+
+
+@dataclass
+class Category:
+    name: str
+    description: str
+    traits: List[str]
+
+    def dict(self):
+        return {k: str(v) for k, v in asdict(self).items()}

src/dataclasses/process_state.py

@@ -0,0 +1,17 @@
+
+from typing import List, TypedDict
+
+
+class ProcessState(TypedDict):
+    question_num: int
+
+    state_name: str
+    user_choice: str
+
+    previous_state_names: List[str]
+    previous_questions: List[str]
+    previous_user_choices: List[str]
+
+    builder_instruction: str
+
+    foll_question_uuid: str

src/dataclasses/question.py

@@ -0,0 +1,12 @@
+
+from dataclasses import asdict, dataclass
+from typing import Dict
+
+
+@dataclass
+class Question:
+    question: str
+    options: Dict[str, str]
+    
+    def dict(self):
+        return {k: str(v) for k, v in asdict(self).items()}

src/frontend/chatbot.py

@@ -0,0 +1,207 @@
+
+import gradio as gr
+
+from langgraph.graph.graph import CompiledGraph
+from typing import Dict, List, Union
+
+import src.agents.coordinator as C
+
+from src.dataclasses.process_state import ProcessState
+
+graph: CompiledGraph = None
+MAX_OPTIONS = 16
+state: ProcessState = {
+    'question_num': 0,
+    'state_name': 'start',
+    'user_choice': '',
+    'previous_state_names': [],
+    'previous_questions': [],
+    'previous_user_choices': [],
+    'builder_instruction': '',
+    'foll_question_uuid': ''
+}
+state_history: List[ProcessState] = []
+num_questions_: int = None
+
+chatbot = gr.Chatbot(
+    type='messages',
+    key='chatbot',
+    preserved_by_key='key',
+    visible=False
+)
+option_buttons = [
+    gr.Button(
+        value='',
+        visible=False,
+        render=False,
+        key=f'option_button_{i}',
+        preserved_by_key='key'
+    )
+    for i in range(1, MAX_OPTIONS + 1)
+]
+restart_button = gr.Button(
+    'Explore other storylines!',
+    visible=False,
+    interactive=True,
+    key='restart_button',
+    preserved_by_key='key'
+)
+button_row = gr.Row(
+    key='button_row',
+    preserved_by_key='key'
+)
+
+def init_graph(
+    story_context: str,
+    categories_context: str,
+    num_questions: int,
+    num_options: int,
+    num_categories: int
+) -> None:
+    global graph, option_buttons, num_questions_
+
+    num_questions_ = num_questions
+    graph = C.init_graph(
+        story_context,
+        categories_context,
+        num_questions,
+        num_options,
+        num_categories
+    )
+
+    for i in range(num_options, MAX_OPTIONS):
+        option_buttons[i].unrender()
+    
+    option_buttons = option_buttons[:num_options]
+
+
+def on_user_response(
+    user_choice: str,
+    history: List[Dict[str, str]]
+) -> Dict[Union[gr.Button, gr.Chatbot], Union[List[Dict[str, str]], gr.update]]:
+    state['user_choice'] = user_choice.replace(' ', '_')
+
+    user_message = [{'role': 'user', 'content': user_choice}]
+    updated_history = history + user_message
+
+    return {chatbot: updated_history} | {
+        button: gr.update(visible=False)
+        for button in option_buttons
+    }
+
+
+def control_screen_widgets() -> List[Union[gr.Chatbot, gr.Row, gr.Button]]:
+    return [chatbot, button_row, restart_button]
+
+
+def control_screen(
+    is_visible: bool
+) -> Dict[Union[gr.Chatbot, gr.Row, gr.Button], Union[gr.update, gr.Row]]:
+    return {
+        chatbot: gr.update(visible=is_visible),
+        button_row: gr.Row(visible=is_visible),
+        restart_button: gr.update(visible=is_visible)
+    }
+
+
+def on_chatbot_response(
+    history: List[Dict[str, str]]
+) -> Dict[Union[gr.Chatbot, gr.Button], Union[List[Dict[str, str]], gr.update]]:
+    global state
+
+    state_history.append({key: val for key, val in state.items()})
+
+    question_num = state['question_num']
+    if question_num < num_questions_:
+        state = graph.invoke(state)
+        
+        question_num = state['question_num']
+        state_name = state['state_name']
+        scenario = C.states[question_num][state_name]
+        question = C.questions[state['foll_question_uuid']]
+        question_str = question.question.replace('_', ' ').strip('\"').strip('\'')
+
+        text_to_user = scenario + '\n' + question_str
+        bot_message = [{'role': 'assistant', 'content': text_to_user}]
+        updated_history = history + bot_message
+
+        options = [
+            option.replace('_', ' ').strip('\"').strip('\'')
+            for option in question.options.keys()
+        ]
+
+        button_updates = {
+            button: gr.update(value=option, visible=True)
+            for option, button in zip(options, option_buttons)
+        }
+        button_updates[restart_button] = gr.update(visible=False)
+    else:
+        selected_category, reason = C.evaluation(state)
+        description, traits = None, None
+        for category in C.categories:
+            if category.name != selected_category:
+                continue
+            description = category.description
+            traits = category.traits
+            C.categories_seen[selected_category] = True
+        traits_string = traits
+        bot_messages = [
+            {'role': 'assistant', 'content': f'We think that you are closest to **{selected_category}**!'},
+            {'role': 'assistant', 'content': f'## {selected_category}\n\n{description}\n\n{traits_string}'},
+            {'role': 'assistant', 'content': reason}
+        ]
+        updated_history = history + bot_messages
+        button_updates = {
+            button: gr.update(visible=False)
+            for button in option_buttons
+        }
+        button_updates[restart_button] = gr.update(visible=True)
+    
+    C.save_coordinator()
+
+    return {chatbot: gr.update(value=updated_history, visible=True)} | button_updates
+
+
+def on_restart_button_click()-> Dict[
+    Union[gr.Chatbot, gr.Button], Union[List[Dict[str, str]], gr.update]
+]:
+    global state
+
+    state = {
+        'question_num': 0,
+        'state_name': 'start',
+        'user_choice': '',
+        'previous_state_names': [],
+        'previous_questions': [],
+        'previous_user_choices': [],
+        'builder_instruction': '',
+        'foll_question_uuid': ''
+    }
+
+    return {restart_button: gr.update(visible=False)} | \
+        on_chatbot_response([])
+
+
+def render():
+    chatbot.render()
+
+    button_row.render()
+    with button_row:
+        for button in option_buttons:
+            button.unrender()
+            button.render()
+            button.click(
+                fn=on_user_response,
+                inputs=[button, chatbot],
+                outputs=option_buttons + [chatbot]
+            ).then(
+                fn=on_chatbot_response,
+                inputs=[chatbot],
+                outputs=option_buttons + [chatbot, restart_button]
+            )
+    restart_button.render()
+    restart_button.click(
+        fn=on_restart_button_click,
+        inputs=None,
+        outputs=option_buttons + [chatbot, restart_button]
+    )

src/frontend/continue_story_screen.py

@@ -0,0 +1,134 @@
+
+import gradio as gr
+import os
+
+from typing import Dict, List, Tuple, Union
+
+import src.agents.coordinator as C
+
+from src.frontend.chatbot import (
+    chatbot,
+    init_graph,
+    on_chatbot_response,
+    option_buttons,
+    restart_button
+)
+from src.frontend import sidebar
+from src.utils.utils import load_information, transform_story_name
+
+information_text = gr.Text(
+    value='Select which story to continue below!',
+    interactive=False,
+    key='continue_story_user_information',
+    preserved_by_key='key',
+    visible=False,
+    label=None
+)
+session_dir = os.path.join(
+    os.path.dirname(__file__),
+    '..', '..', 'sessions'
+)
+story_widgets: List[Tuple[gr.Row, gr.Button, gr.Text]] = []
+story_dir_mapper = dict()
+for i, dirname in enumerate(sorted(os.listdir(session_dir))):
+    if dirname.startswith('.'):
+        continue
+    dirpath = os.path.join(session_dir, dirname)
+    if not os.listdir(dirpath):
+        continue
+    story_information_filepath = os.path.join(dirpath, 'story.json')
+    story_information_dict = load_information(story_information_filepath)
+    story_name = story_information_dict['story_name']
+    story_context = story_information_dict['story_context']
+    row = gr.Row(
+        visible=False,
+        key=f'continue_row_{i}',
+        preserved_by_key='key'
+    )
+    button = gr.Button(
+        value=story_name,
+        visible=False,
+        key=f'continue_button_{i}',
+        preserved_by_key='key',
+        scale=1
+    )
+    text = gr.Text(
+        value=story_context,
+        interactive=False,
+        visible=False,
+        key=f'continue_story_{i}',
+        preserved_by_key='key',
+        scale=3
+    )
+    story_widgets.append((row, button, text))
+    story_dir_mapper[story_name] = dirpath
+
+
+def get_widgets() -> List[Union[gr.Text, gr.Button, gr.Row]]:
+    widgets = [information_text]
+    for row, button, text in story_widgets:
+        widgets.append(row)
+        widgets.append(button)
+        widgets.append(text)
+    return widgets
+
+
+def get_widgets_updates(
+    is_visible: bool
+) -> Dict[Union[gr.Slider, gr.Text, gr.Button], gr.update]:
+    updates = dict()
+    for widget in get_widgets():
+        if isinstance(widget, gr.Row):
+            updates[widget] = gr.Row(visible=is_visible)
+        else:
+            updates[widget] = gr.update(visible=is_visible)
+    return updates
+
+
+def on_button_click(
+    story_name: str
+) -> Dict[Union[gr.Text, gr.Slider, gr.Button, gr.Chatbot], gr.update]:
+    
+    dirpath = story_dir_mapper[story_name]
+    story_name_ = transform_story_name(story_name)
+
+    story_information_filepath = os.path.join(dirpath, 'story.json')
+    story_information_dict = load_information(story_information_filepath)
+    story_context = story_information_dict['story_context']
+    categories_context = story_information_dict['categories_context']
+    num_questions = story_information_dict['num_questions']
+    num_options = story_information_dict['num_options']
+    num_categories = story_information_dict['num_categories']
+
+    init_graph(
+        story_context,
+        categories_context,
+        num_questions,
+        num_options,
+        num_categories
+    )
+
+    C.story_name = story_name_
+    C.load_coordinator()
+
+    chatbot_updates = on_chatbot_response([])
+
+    return chatbot_updates | get_widgets_updates(False) | sidebar.view_screen()
+
+
+def render():
+    information_text.render()
+
+    for row, button, text in story_widgets:
+        row.render()
+        with row:
+            button.render()
+            text.render()
+        button.click(
+            fn=on_button_click,
+            inputs=[button],
+            outputs=get_widgets() + \
+                [chatbot, restart_button] + \
+                option_buttons + \
+                sidebar.get_widgets()
+        )

src/frontend/progress_screen.py

@@ -0,0 +1,240 @@
+
+import gradio as gr
+import pandas as pd
+import plotly.graph_objects as go
+
+from igraph import Graph, EdgeSeq
+from typing import Dict, List, Union
+
+import src.agents.coordinator as C
+
+categories_row = gr.Row(
+    visible=False,
+    key='categories_row',
+    preserved_by_key='key'
+)
+graph_row = gr.Row(
+    visible=False,
+    key='graph_row',
+    preserved_by_key='key'
+)
+
+category_buttons = [
+    gr.Button(
+        value='',
+        visible=False,
+        key=f'category_button_{i + 1}',
+        preserved_by_key='key'
+    )
+    for i in range(16)
+]
+category_text = gr.Text(
+    value='Select a category!',
+    visible=False,
+    interactive=False,
+    key='category_text',
+    preserved_by_key='key'
+)
+graph_plot = gr.Plot(
+    visible=False,
+    key='graph_plot',
+    preserved_by_key='key'
+)
+
+
+def control_screen(is_visible: bool) -> Dict[gr.Row, gr.Row]:
+    return {
+        categories_row: gr.Row(visible=is_visible),
+        graph_row: gr.Row(visible=is_visible)
+    }
+
+
+def update_categories() -> Dict[Union[gr.Button, gr.Text], gr.update]:
+    button_updates = {
+        button: gr.update(
+            value=category.name if C.categories_seen[category.name] \
+                else '???',
+            visible=True,
+            interactive=C.categories_seen[category.name],
+            variant='primary' if C.categories_seen[category.name] \
+                else 'secondary'
+        )
+        for button, category in zip(category_buttons, C.categories)
+    }
+
+    return {
+        category_text: gr.update(visible=True, value='Select a category!')
+    } | button_updates
+
+
+def on_category_click(category_name: str) -> str:
+    for category in C.categories:
+        if category.name != category_name:
+            continue
+        return category.name + '\n\n' + \
+            category.description + '\n\n' + \
+            category.traits
+
+
+def update_graph():
+    n_vertices = sum([
+        len(states_)
+        for level, states_ in C.states.items()
+        if level <= C.num_questions_
+    ])
+
+    graph = Graph(directed=True)
+    nodes = []
+    node_attributes = {'description': []}
+    for level in range(1, C.num_questions_ + 1):
+        for name, description in C.states[level].items():
+            node_name = f'{name} (Stage {level})'
+            nodes.append(node_name)
+            node_attributes['description'].append(description)
+    
+    edges = []
+    edge_attributes = {
+        'question': [],
+        'option': []
+    }
+    for level in range(1, C.num_questions_):
+        for state, question_uuid in C.state_question_map[level].items():
+            question = C.questions[question_uuid]
+            question_str = question.question
+            options = question.options
+            for option, foll_state in options.items():
+                edge = (
+                    f'{state} (Stage {level})',
+                    f'{foll_state} (Stage {level + 1})'
+                )
+                edges.append(edge)
+                edge_attributes['question'].append(question_str)
+                edge_attributes['option'].append(option)
+
+    graph.add_vertices(nodes, attributes=node_attributes)
+    graph.add_edges(edges, attributes=edge_attributes)
+    layout = graph.layout('rt')
+
+    # adapted from https://plotly.com/python/tree-plots/
+    position = {k: layout[k] for k in range(n_vertices)}
+    Y = [layout[k][1] for k in range(n_vertices)]
+    M = max(Y)
+
+    E = [e.tuple for e in graph.es] # list of edges
+
+    L = len(position)
+    Xn = [position[k][0] for k in range(L)]
+    Yn = [2*M-position[k][1] for k in range(L)]
+    Xe = []
+    Ye = []
+
+    # for labelling edges
+    X_edge_nodes = []
+    Y_edge_nodes = []
+    for edge in E:
+        Xe+=[position[edge[0]][0],position[edge[1]][0], None]
+        Ye+=[2*M-position[edge[0]][1],2*M-position[edge[1]][1], None]
+
+        X_edge_nodes.append((position[edge[0]][0] + position[edge[1]][0]) / 2)
+        Y_edge_nodes.append((2*M-position[edge[0]][1] + 2*M-position[edge[1]][1]) / 2)
+
+    node_labels = [
+        node.replace('_', ' ') + '\n\n' + \
+        description
+        for node, description in zip(nodes, node_attributes['description'])
+    ]
+    node_labels = pd.DataFrame(node_labels, columns=['label'])
+    node_labels['label'] = node_labels['label'].str.wrap(30)\
+        .apply(lambda x: x.replace('\n', '<br>'))
+    node_labels = node_labels['label'].to_list()
+    edge_labels = [
+        question.replace('_', ' ') + '\n\n' + option.replace('_', ' ')
+        for question, option in zip(
+            edge_attributes['question'], edge_attributes['option']
+        )
+    ]
+    edge_labels = pd.DataFrame(edge_labels, columns=['label'])
+    edge_labels['label'] = edge_labels['label'].str.wrap(30)\
+        .apply(lambda x: x.replace('\n', '<br>'))
+    edge_labels = edge_labels['label'].to_list()
+
+    fig = go.Figure()
+    fig.add_trace(go.Scatter(
+        x=Xe, y=Ye,
+        mode='lines',
+        line=dict(color='rgb(210,210,210)', width=1),
+    ))
+    fig.add_trace(go.Scatter(
+        x=Xn, y=Yn,
+        mode='markers',
+        marker=dict(
+            symbol='circle-dot', size=18, color='#6175c1',
+            line=dict(color='rgb(50,50,50)', width=1)
+        ),
+        text=node_labels,
+        hoverinfo='text',
+        opacity=0.8
+    ))
+
+    fig.add_trace(go.Scatter(
+        x=X_edge_nodes, y=Y_edge_nodes,
+        mode='markers',
+        marker=dict(
+            symbol='circle-dot', size=0, color="#42c744",
+            line=dict(color='rgb(50,50,50)', width=0)
+        ),
+        text=edge_labels,
+        hoverinfo='text',
+        opacity=0
+    ))
+
+    axis = dict(
+        showline=False,
+        zeroline=False,
+        showgrid=False,
+        showticklabels=False,
+    )
+
+    fig.update_layout(
+        showlegend=False,
+        xaxis=axis,
+        yaxis=axis
+    )
+
+    return gr.Plot(fig, visible=True)
+
+
+def control_screen_widgets() -> List[Union[gr.Row, gr.Text]]:
+    return [categories_row, category_text, graph_row, graph_plot] + \
+        category_buttons
+
+
+def control_screen(
+    is_visible: bool
+) -> Dict[Union[gr.Plot, gr.Row, gr.Button], Union[gr.update, gr.Row]]:
+    row_updates = {
+        categories_row: gr.Row(visible=is_visible),
+        graph_row: gr.Row(visible=is_visible)
+    }
+    graph_update = {
+        graph_plot: update_graph() if is_visible else gr.update(visible=False)
+    }
+    category_button_updates = update_categories()
+
+    return row_updates | graph_update | category_button_updates
+
+
+def render():
+    categories_row.render()
+    with categories_row:
+        for button in category_buttons:
+            button.render()
+            button.click(
+                fn=on_category_click,
+                inputs=[button],
+                outputs=[category_text]
+            )
+    category_text.render()
+    graph_row.render()
+    with graph_row:
+        graph_plot.render()

src/frontend/sidebar.py

@@ -0,0 +1,108 @@
+
+import gradio as gr
+
+from typing import Dict, List, Union
+
+from src.frontend import chatbot, progress_screen
+
+sidebar = gr.Sidebar(
+    open=True,
+    visible=False,
+    key='sidebar',
+    preserved_by_key='key'
+)
+chatbot_button = gr.Button(
+    value='Viewing Chatbot!',
+    variant='secondary',
+    visible=False,
+    interactive=False,
+    key='chatbot_button',
+    preserved_by_key='key'
+)
+progress_button = gr.Button(
+    value='View Your Story Progress!',
+    variant='primary',
+    visible=False,
+    interactive=True,
+    key='progress_button',
+    preserved_by_key='key'
+)
+
+
+def get_widgets() -> List[Union[gr.Sidebar, gr.Button]]:
+    return [sidebar, chatbot_button, progress_button]
+
+
+def view_screen() -> Dict[Union[gr.Sidebar, gr.Button], gr.update]:
+    return {
+        widget: gr.update(visible=True)
+        for widget in get_widgets()
+    }
+
+
+def on_chatbot_button_click():
+    chatbot_screen_updates = chatbot.control_screen(True)
+    progress_screen_updates = progress_screen.control_screen(False)
+
+    chatbot_button_update = gr.update(
+        value='Viewing Chatbot!',
+        variant='secondary',
+        interactive=False
+    )
+    progress_button_update = gr.update(
+        value='View Your Story Progress!',
+        variant='primary',
+        interactive=True
+    )
+
+    return {
+        chatbot_button: chatbot_button_update,
+        progress_button: progress_button_update
+    } | \
+        chatbot_screen_updates | progress_screen_updates
+
+
+def on_progress_button_click():
+    chatbot_screen_updates = chatbot.control_screen(False)
+    progress_screen_updates = progress_screen.control_screen(True)
+
+    chatbot_button_update = gr.update(
+        value='Come & Explore!',
+        variant='primary',
+        interactive=True
+    )
+    progress_button_update = gr.update(
+        value='Viewing Your Story Progress!',
+        variant='secondary',
+        interactive=False
+    )
+
+    return {
+        chatbot_button: chatbot_button_update,
+        progress_button: progress_button_update
+    } | \
+        chatbot_screen_updates | progress_screen_updates
+
+
+def render():
+    sidebar.render()
+    with sidebar:
+        chatbot_button.render()
+        progress_button.render()
+
+    chatbot_button.click(
+        fn=on_chatbot_button_click,
+        inputs=[],
+        outputs=[chatbot_button, progress_button] + \
+            chatbot.control_screen_widgets() + \
+            progress_screen.control_screen_widgets()
+    )
+
+    progress_button.click(
+        fn=on_progress_button_click,
+        inputs=[],
+        outputs=[chatbot_button, progress_button] + \
+            chatbot.control_screen_widgets() + \
+            progress_screen.control_screen_widgets()
+    )
+

src/frontend/start_screen.py

@@ -0,0 +1,98 @@
+
+import gradio as gr
+
+from typing import Dict, List, Union
+
+from src.frontend import (
+    continue_story_screen,
+    story_information_widgets
+)
+from src.utils.init_openai import init_client
+
+button_row = gr.Row(
+    key='start_screen_button_row',
+    preserved_by_key='key'
+)
+new_story_button = gr.Button(
+    value='Begin a new adventure!',
+    visible=True,
+    interactive=True,
+    key='new_story_button',
+    preserved_by_key='key'
+)
+continue_story_button = gr.Button(
+    value='Continue an existing adventure!',
+    visible=True,
+    interactive=True,
+    key='continue_story_button',
+    preserved_by_key='key'
+)
+api_key_textbox = gr.Text(
+    type='password',
+    label='OpenAI API Key',
+    placeholder='Enter your OpenAI API key',
+    interactive=True,
+    visible=True,
+    key='api_key_textbox',
+    preserved_by_key='key'
+)
+
+
+def get_widgets() -> List[Union[gr.Text, gr.Button, gr.Row]]:
+    return [
+        button_row,
+        new_story_button,
+        continue_story_button,
+        api_key_textbox
+    ]
+
+
+def get_wigets_updates(
+    is_visible: bool = False
+) -> Dict[Union[gr.Text, gr.Button, gr.Row], Union[gr.update, gr.Row]]:
+    return {
+        widget: gr.Row(visible=is_visible) if isinstance(widget, gr.Row) else \
+            gr.update(visible=is_visible)
+        for widget in get_widgets()
+    }
+
+
+def on_submit_new_story(
+    api_key: str
+) -> Dict[Union[gr.Text, gr.Slider, gr.Button], Union[gr.Row, gr.update]]:
+    if init_client(api_key):
+        return get_wigets_updates(False) | \
+            story_information_widgets.get_widgets_updates(True)
+    return get_wigets_updates(True) | \
+        story_information_widgets.get_widgets_updates(False)
+
+
+def on_submit_continue_story(
+    api_key: str
+) -> Dict[Union[gr.Text, gr.Slider, gr.Button], Union[gr.Row, gr.update]]:
+    if init_client(api_key):
+        return get_wigets_updates(False) | \
+            continue_story_screen.get_widgets_updates(True)
+    return get_wigets_updates(True) | \
+        continue_story_screen.get_widgets_updates(False)
+
+
+def render():
+    api_key_textbox.render()
+    button_row.render()
+    with button_row:
+        new_story_button.render()
+        continue_story_button.render()
+
+    new_story_button.click(
+        fn=on_submit_new_story,
+        inputs=[api_key_textbox],
+        outputs=get_widgets() + \
+            story_information_widgets.get_widgets()
+    )
+    continue_story_button.click(
+        fn=on_submit_continue_story,
+        inputs=[api_key_textbox],
+        outputs=get_widgets() + \
+            continue_story_screen.get_widgets()
+    )

src/frontend/story_information_widgets.py

@@ -0,0 +1,210 @@
+
+import gradio as gr
+import json
+import os
+
+from typing import Dict, List, Union
+
+import src.agents.coordinator as C
+
+from src.frontend.chatbot import (
+    MAX_OPTIONS,
+    chatbot,
+    init_graph,
+    on_chatbot_response,
+    option_buttons,
+    restart_button
+)
+from src.frontend import sidebar
+from src.utils.utils import (
+    get_session_dir,
+    save_information,
+    transform_story_name
+)
+
+story_name_textbox = gr.Text(
+    placeholder='Give your adventure a unique name!',
+    label='Story Name',
+    interactive=True,
+    key='story_name_textbox',
+    preserved_by_key='key',
+    visible=False
+)
+story_context_textbox = gr.Text(
+    placeholder='What kind of story would you like to adventure on?',
+    label='Story Context',
+    interactive=True,
+    key='story_context_textbox',
+    preserved_by_key='key',
+    visible=False
+)
+categories_context_textbox = gr.Text(
+    placeholder='What is the theme of the categories you want to have?',
+    label='Categories Context',
+    interactive=True,
+    key='categories_context_textbox',
+    preserved_by_key='key',
+    visible=False
+)
+num_questions_slider = gr.Slider(
+    minimum=1,
+    maximum=10,
+    value=5,
+    step=1,
+    label='How many questions would you want your journey to have?',
+    interactive=True,
+    visible=False,
+    key='num_questions_slider',
+    preserved_by_key='key'
+)
+num_options_slider = gr.Slider(
+    minimum=2,
+    maximum=4,
+    value=2,
+    step=1,
+    label='How many options would you like at each turn?',
+    interactive=True,
+    visible=False,
+    key='num_options_slider',
+    preserved_by_key='key'
+)
+num_categories_slider = gr.Slider(
+    minimum=2,
+    maximum=MAX_OPTIONS,
+    value=2,
+    step=1,
+    label='How many categories / endings would you like to have?',
+    interactive=True,
+    visible=False,
+    key='num_categories_slider',
+    preserved_by_key='key'
+)
+story_information_submit_button = gr.Button(
+    visible=False,
+    interactive=False,
+    key='story_information_submit_button',
+    preserved_by_key='key'
+)
+
+def get_widgets() -> List[Union[gr.Text, gr.Slider, gr.Button]]:
+    return [
+        story_name_textbox,
+        story_context_textbox,
+        categories_context_textbox,
+        num_questions_slider,
+        num_options_slider,
+        num_categories_slider,
+        story_information_submit_button
+    ]
+
+
+def get_widgets_updates(
+    is_visible: bool
+) -> Dict[Union[gr.Slider, gr.Text, gr.Button], gr.update]:
+    return {
+        widget: gr.update(visible=is_visible)
+        for widget in get_widgets()
+    }
+
+
+def check_story_name(story_name: str) -> bool:
+    story_name_dir = transform_story_name(story_name)
+    return not story_name_dir in os.listdir(get_session_dir())
+
+
+def on_text_change(
+    story_name: str,
+    story_context: str,
+    categories_context: str
+) -> gr.update:
+    if all([story_name, story_context, categories_context]) and \
+            check_story_name(story_name):
+        return gr.update(interactive=True)
+    return gr.update(interactive=False)
+
+
+def save_story_information(
+    story_name: str,
+    story_context: str,
+    categories_context: str,
+    num_questions: int,
+    num_options: int,
+    num_categories: int
+):
+    story_name_dir = transform_story_name(story_name)
+    story_dirpath = os.path.join(get_session_dir(), story_name_dir)
+    os.mkdir(story_dirpath)
+
+    story_information_filepath = os.path.join(story_dirpath, 'story.json')
+    story_information = {
+        'story_name': story_name,
+        'story_context': story_context,
+        'categories_context': categories_context,
+        'num_questions':  num_questions,
+        'num_options': num_options,
+        'num_categories': num_categories
+    }
+    save_information(story_information, story_information_filepath)
+
+
+def on_submit(
+    story_name: str,
+    story_context: str,
+    categories_context: str,
+    num_questions: int,
+    num_options: int,
+    num_categories: int
+) -> Dict[Union[gr.Text, gr.Slider, gr.Button, gr.Chatbot], gr.update]:
+    
+    save_story_information(
+        story_name,
+        story_context,
+        categories_context,
+        num_questions,
+        num_options,
+        num_categories
+    )
+    C.story_name = transform_story_name(story_name)
+
+    init_graph(
+        story_context,
+        categories_context,
+        num_questions,
+        num_options,
+        num_categories
+    )
+
+    chatbot_updates = on_chatbot_response([])
+
+    return chatbot_updates | get_widgets_updates(False) | sidebar.view_screen()
+
+
+def render():
+
+    for widget in get_widgets():
+        widget.render()
+
+        if not isinstance(widget, gr.Text):
+            continue
+        widget.change(
+            fn=on_text_change,
+            inputs=[
+                story_name_textbox,
+                story_context_textbox,
+                categories_context_textbox
+            ],
+            outputs=[story_information_submit_button]
+        )
+
+    story_information_submit_button.click(
+        fn=on_submit,
+        inputs=[
+            widget
+            for widget in get_widgets()
+            if not isinstance(widget, gr.Button)
+        ],
+        outputs=get_widgets() + \
+            [chatbot, restart_button] + \
+            option_buttons + \
+            sidebar.get_widgets()
+    )

src/prompts/category_builder_agent.py

@@ -0,0 +1,43 @@
+
+system_prompt_template = """
+You are a creative writer.
+Your task is to create categories based on a given context.
+Your categories should be different from one another.
+Keep the categories fun, light and interesting. Be as creative as possible.
+
+You will be given the following:
+1. `story context`, which describe the story we are about to take the user on,
+2. `category context`, which describe what the general theme of the categories should be,
+3. the `number of categories`
+
+You can find the provided information later, where
+1. `story context` can be found enclosed in <story> tags,
+2. `category context` can be found enclosed in <category> tags,
+3. `number of categories` can be found enclosed in <number> tags.
+
+Ensure that the categories you generate are related to the story context.
+When you are ready, output **all** the categories in **JSON** format.
+For each category, they should have the following fields:
+1. `name`, which describes the name of the category,
+2. `description`: which gives a short description (at most three lines) of the category,
+3. `traits`: a sequence of traits of the category. Separate the traits with a comma on the same line. Ensure that the traits have some relation to the story context.
+
+Ensure that the focus is still on the `category context`.
+A useful guideline is that the `name` of the category directly follows from the `category context`.
+And then the `story context` can be used in the `description` and `traits`.
+You only need to output the JSON information, no elaboration is required.
+""".strip()
+
+user_prompt_template = """
+<story>
+{story_context}
+</story>
+
+<category>
+{category_context}
+</category>
+
+<number>
+{num_categories}
+</number>
+""".strip()

src/prompts/evaluator_agent.py

@@ -0,0 +1,38 @@
+
+system_prompt_template = """
+You are an expert profiler, regardless of the theme.
+We have just set our user through a round of questions, and they have provided us their responses.
+Our task now is to allocate them into one category based on the scenario they were given, questions they were asked, and choices they made.
+
+You will be given the following:
+1. `user story`, where for each turn, you are given the SCENARIO, QUESTION and USER CHOICE,
+2. `categories`, where for each category, you are given the name, description and traits of the category.
+
+You can find the provided information later, where
+1. `user story` can be found enclosed in <user> tags,
+2. `categories` can be found enclosed in <categories> tags,
+
+Read the user's story and determine which category best fits the user based on their choices.
+
+When you are ready, output your answer in the below JSON format.
+Follow the format strictly, as any deviation will result in errors later on.
+You only need to output your answer, no elaboration is required.
+`category` below refers to the category name of the selected category. Output the corresponding category name exactly.
+For your reason, replace terms like \'the user\' with \'your\' or \'you\'. We will be passing your reason to the user.
+```json
+{{
+    "category":
+    "reason":
+}}
+```
+""".strip()
+
+user_prompt_template = """
+<user>
+{user_story}
+</user>
+
+<categories>
+{category_contents}
+</categories>
+""".strip()

src/prompts/narrator_agent.py

@@ -0,0 +1,97 @@
+
+system_prompt_template = """
+You are a creative writer for a story, in charge of writing the lines users will view.
+The story you are making is dependent on our user's choices. Do **not** rush to make a complete story.
+
+We are building a turn-based call and response story, where we will give users a scenario, and they will respond with what they would choose to do.
+Based on our user's response, and the context they give us, we will craft what the next scenario will be.
+
+Your task is to craft a question for the user to answer based on the scenario provided to you, and the potential responses to the question you make.
+Your question should allow the user to respond with actionable options.
+Focus on creating a question that is in line with the given scenario and storyline so far.
+Keep in mind the context of the categories we will assign the user to in the end.
+There should be some harmony and cohesion between the story's context and the categories' context.
+
+You will be given the following:
+1. `story context`, which describes the theme of the story we want to build for our users,
+2. `categories context`, which describes the theme of the categories we will assign our users to later,
+3. `user story`, which is the current story shown to the user,
+4. `scenario`, which is the scenario that your question should be based on,
+5. `builder instruction`, which describes the suggested pace of the story,
+6. `number of options`, which is the number of options the user has to respond to the question you pose. You will need to generate the options as well.
+7. `existing next scenarios`, which is the current list of scenarios that follow this scenario. This will be empty if we are seeing the current scenario for the first time. You will be given each scenario in the below format.
+```
+{{
+    scenario_name:
+    scenario:
+}}
+```
+
+You can find the provided information later, where
+1. `story context` can be found enclosed in <story> tags,
+2. `categories context` can be found enclosed in <category> tags.
+3. `user story` can be found enclosed in <user> tags,
+4. `scenario`, can be found enclosed in <scenario> tags,
+5. `builder instruction`, can be found enclosed in <builder> tags,
+6. `number of options`, can be found enclosed in <options> tags,
+7. `existing next scenarios`, can be found enclosed in <existing> tags
+
+First, think about a question to ask the user. This question should not be long winded.
+Instead, aim for short questions that are easy for the users to respond to, but also compelling enough to keep them engaged.
+Next, think about potential responses to the question.
+The number of options you provide **must** be exactly the number of options provided, no more and no less.
+Lastly, think about what suitable next scenarios each of the options can lead to.
+Prioritise using **existing** scenarios, where possible.
+Prioritise using the **same** scenario for multiple options, where possible.
+If you want to use the same or existing scenario, you **must** use the exact same `scenario_name`.
+If you want to create a new scenario, you **must** use a different `scenario_name` than the ones provided.
+Although you do not know the final categories, ensure you provide options that are distinct enough such that the users can definitely be binned into different categories.
+
+When you are ready, output your answer in the below JSON format.
+Follow the format strictly, as any deviation will result in errors later on.
+For your question, replace any spaces with the underbar character.
+You only need to output your answer, no elaboration is required.
+If there are angular brackets, it indicates to replace those fields with your answer.
+Where there are ellipsis, it means to continue following the format.
+```json
+{{
+    "question":
+    "options": {{
+        "<option_1>": "<following scenario name>",
+        ...
+    }}
+}}
+```
+
+You are reminded to ensure that the number of options is exactly the number of options provided.
+""".strip()
+
+user_prompt_template = """
+<story>
+{story_context}
+</story>
+
+<category>
+{category_context}
+</category>
+
+<user>
+{user_story}
+</user>
+
+<scenario>
+{scenario}
+</scenario>
+
+<builder>
+{builder_instruction}
+</builder>
+
+<options>
+{num_options}
+</options>
+
+<existing>
+{existing_scenarios}
+</existing>
+""".strip()

src/prompts/orchestrator_agent.py

@@ -0,0 +1,108 @@
+
+start_system_prompt_template = """
+You are an orchestrator of a story board.
+Your task is to provide instructions to your creatives about the flow and urgency of the story.
+This means whether the creatives should start thinking about wrapping up the story, or if they have some time to explore a certain scenario.
+
+We are building a turn-based call and response story, where we will give users a scenario, and they will respond with what they would choose to do.
+Based on our user's response, and the context they give us, we will craft what the next scenario will be.
+
+Your task is **not** to think about what to craft next.
+Instead, your task is to guide the agent that is doing the creative work.
+Your goal is to guide the story such that it does to feel too rushed or too drawn out, whilst ensuring a compelling and interesting story.
+
+In this first step, we have just begun the process.
+You will be given the following:
+1. `story context`, which describes the theme of the story we want to build for our users,
+2. `number of questions`, which describes the number of turns we have with our users. Use this as a gauge to determine if the given context is too lengthy or too short.
+
+You can find the provided information later, where
+1. `story context` can be found enclosed in <story> tags,
+2. `number of questions` can be found enclosed in <number> tags.
+
+Think about the context and the number of questions.
+We both know that the creative agent will begin with some kind of introductory question.
+But what we want to focus on is how fast this phase should be.
+
+Your output is the instruction to our creative agent on the tempo of the story.
+When the creative agent reads your instruction, they should know how fast to craft the scenario to give the users.
+Keep the user's overall experience the focus.
+
+Once you are ready, output your instruction to our creative agent.
+No elaboration is required.
+""".strip()
+
+start_user_prompt_template = """
+<story>
+{story_context}
+</story>
+
+<number>
+{num_questions}
+</number>
+""".strip()
+
+coordinate_system_prompt_template = """
+You are an orchestrator of a story board.
+Your task is to provide instructions to your creatives about the flow and urgency of the story.
+This means whether the creatives should start thinking about wrapping up the story, or if they have some time to explore a certain scenario.
+
+We are building a turn-based call and response story, where we will give users a scenario, and they will respond with what they would choose to do.
+Based on our user's response, and the context they give us, we will craft what the next scenario will be.
+
+Your task is **not** to think about what to craft next.
+Instead, your task is to guide the agent that is doing the creative work.
+Your goal is to guide the story such that it does to feel too rushed or too drawn out, whilst ensuring a compelling and interesting story.
+
+In this phase, we are some way into our story.
+You will be given the following:
+1. `story context`, which describes the theme of the story we want to build for our users,
+2. `number of questions`, which describes the number of turns we have with our users. Use this as a gauge to determine if the given context is too lengthy or too short.
+3. `phase message`: which is a pre-written phrase related to how far into the story the user is in. You do not need to calculate this, we will do it for you,
+4. `user story`, which is the current story shown to the user.
+
+You can find the provided information later, where
+1. `story context` can be found enclosed in <story> tags,
+2. `number of questions` can be found enclosed in <number> tags,
+3. `phase message` can be found enclosed in <phase> tags,
+4. `user story` can be found enclosed in <user> tags.
+
+Think about the context and phase.
+Our creative agent is getting ready to craft the next scenario to give to the user.
+Our focus on is how fast scenario moves the story along.
+
+Your output is the instruction to our creative agent on the tempo of the story.
+When the creative agent reads your instruction, they should know how fast to craft the scenario to give the users.
+Keep the user's overall experience the focus.
+
+Once you are ready, output your instruction to our creative agent.
+No elaboration is required.
+""".strip()
+
+coordinate_user_prompt_template = """
+<story>
+{story_context}
+</story>
+
+<number>
+{num_questions}
+</number>
+
+<phase>
+{phase_message}
+</phase>
+
+<user>
+{user_story}
+</user>
+""".strip()
+
+def get_phase_message(progress: float) -> str:
+    if progress < 0.2:
+        return 'We are still in the early phases of the story!'
+    if progress <= 0.5:
+        return 'We are approaching the halfway point, it\'s time to build towards a climax!'
+    if progress < 0.8:
+        return 'We are approaching the climax of the story! Let\'s quicken the pace.'
+    
+    return 'The story is approaching the end! Let\'s wrap this up!'

src/prompts/storybuilder_agent.py

@@ -0,0 +1,66 @@
+
+system_prompt_template = """
+You are a creative writer for a story.
+The story you are making is dependent on our user's choices. Do **not** rush to make a complete story.
+
+We are building a turn-based call and response story, where we will give users a scenario, and they will respond with what they would choose to do.
+Based on our user's response, and the context they give us, we will craft what the next scenario will be.
+
+Your task is to craft a scenario for the user to think about.
+Your task is **not** to craft the question to ask to the user.
+Focus on creating a scenario that flows well with the story so far, and is in line with the story context.
+Keep in mind the context of the categories we will assign the user to in the end.
+There should be some harmony and cohesion between the story's context and the categories' context.
+
+You will be given the following:
+1. `story context`, which describes the theme of the story we want to build for our users,
+2. `categories context`, which describes the theme of the categories we will assign our users to later,
+3. `builder's instruction`, a set of instructions with respect to the tempo of the scenario you are going to create.
+4. `user story`, which is the current story shown to the user,
+5. `question`, which is the last question asked to the user, if we just started the process, this is empty.
+6. `response`, which is the response to the question shown to the user. If we just started the process, this is empty.
+
+You can find the provided information later, where
+1. `story context` can be found enclosed in <story> tags,
+2. `categories context` can be found enclosed in <category> tags.
+3. `builder's instruction` can be found enclosed in <builder> tags.
+4. `user story` can be found enclosed in <user> tags,
+5. `question`, can be found enclosed in <question> tags,
+6. `response`, can be found enclosed in <response> tags.
+
+First, think about a fitting scenario that naturally follows the story the user is shown.
+Next, think about some specifics of the scenario.
+Limit the number of words in your scenario to be at most 50 words.
+Feel free to use named characters and have them recur throughout this journey.
+Keep the scenario light, fun and compelling for the user's to follow.
+
+When you are ready, output **only** the scenario.
+You are reminded that the scenario should **not** include the question to be asked to the user.
+You do not need to provide any explanation.
+""".strip()
+
+user_prompt_template = """
+<story>
+{story_context}
+</story>
+
+<category>
+{category_context}
+</category>
+
+<builder>
+{builder_instruction}
+</builder>
+
+<user>
+{user_story}
+</user>
+
+<question>
+{question}
+</question>
+
+<response>
+{response}
+</response>
+""".strip()

src/utils/init_openai.py

@@ -0,0 +1,80 @@
+
+import random
+import time
+
+from openai import OpenAI, RateLimitError
+from openai.types.chat.chat_completion import ChatCompletion
+from typing import Any, Dict
+
+client: OpenAI = None
+MODEL_ = 'gpt-4o'
+
+
+def init_client(api_key: str) -> bool:
+    global client
+    try:
+        client = OpenAI(api_key=api_key)
+    except:
+        return False
+    return True
+
+
+# from OpenAI's website
+# define a retry decorator
+def __retry_with_exponential_backoff(
+    func,
+    initial_delay: float = 1,
+    exponential_base: float = 2,
+    jitter: bool = True,
+    max_retries: int = 10,
+    errors: tuple = (RateLimitError,),
+):
+    """Retry a function with exponential backoff."""
+
+    def wrapper(*args, **kwargs):
+        # Initialize variables
+        num_retries = 0
+        delay = initial_delay
+
+        # Loop until a successful response or max_retries is hit or an exception is raised
+        while True:
+            try:
+                return func(*args, **kwargs)
+
+            # Retry on specified errors
+            except errors as e:
+                # Increment retries
+                num_retries += 1
+
+                # Check if max retries has been reached
+                if num_retries > max_retries:
+                    raise Exception(
+                        f"Maximum number of retries ({max_retries}) exceeded."
+                    )
+
+                # Increment the delay
+                delay *= exponential_base * (1 + jitter * random.random())
+
+                # Sleep for the delay
+                time.sleep(delay)
+
+            # Raise exceptions for any errors not specified
+            except Exception as e:
+                raise e
+
+    return wrapper
+
+
+@__retry_with_exponential_backoff
+def completions_with_backoff(
+    messages: Dict[str, Any],
+    temperature: float = 0,
+    max_tokens: int = 8192
+) -> ChatCompletion:
+    kwargs = {
+        'model': MODEL_,
+        'messages': messages,
+        'temperature': temperature,
+        'max_tokens': max_tokens
+    }
+    return client.chat.completions.create(**kwargs)

src/utils/utils.py

@@ -0,0 +1,38 @@
+
+import json
+import os
+import uuid
+
+from typing import Any
+
+
+def get_session_dir() -> str:
+    return os.path.join(
+        os.path.dirname(__file__),
+        '..', '..', 'sessions'
+    )
+
+
+def save_information(
+    information: Any,
+    filepath: str
+) -> None:
+    
+    with open(filepath, 'w', encoding='utf-8') as f:
+        json.dump(information, f, ensure_ascii=False, indent=4)
+
+
+def load_information(filepath: str) -> Any:
+    with open(filepath, 'r', encoding='utf-8') as f:
+        content = json.load(f)
+    return content
+
+
+def transform_story_name(
+    story_name: str,
+    is_inverse: bool = False
+) -> str:
+    if is_inverse:
+        return ' '.join([word.capitalize() for word in story_name.split(' ')])
+    else:
+        return story_name.lower().replace(' ', '_')