Squash merge dictify-states into main

app.py

@@ -11,7 +11,7 @@ from components.quizbowl.bonus import BonusInterface
 from components.quizbowl.tossup import TossupInterface
 from display.css_html_js import fonts_header, js_head, leaderboard_css
 from display.custom_css import css_bonus, css_pipeline, css_tossup
-from display.guide import GUIDE_MARKDOWN
+from display.guide import BUILDING_MARKDOWN, GUIDE_MARKDOWN, QUICKSTART_MARKDOWN
 from display.utils import AutoEvalColumn, fields
 
 # Constants
@@ -91,8 +91,6 @@ if __name__ == "__main__":
         theme=THEME,
         title="Quizbowl Bot",
     ) as demo:
-        with gr.Sidebar(width=400):
-            gr.Markdown(GUIDE_MARKDOWN)
         with gr.Row():
             gr.Markdown("## Welcome to Quizbowl Bot! This is a tool for creating and testing quizbowl agents.")
         with gr.Tabs() as gtab:
@@ -121,5 +119,11 @@ if __name__ == "__main__":
                     visible=True,
                 )
                 refresh_btn.click(fn=fetch_leaderboard_df, inputs=[], outputs=leaderboard_table)
+            with gr.Tab("❓ Help", id="help"):
+                with gr.Row():
+                    with gr.Column():
+                        gr.Markdown(QUICKSTART_MARKDOWN)
+                    with gr.Column():
+                        gr.Markdown(BUILDING_MARKDOWN)
 
     demo.queue(default_concurrency_limit=40).launch()

src/components/model_pipeline/model_pipeline.py

@@ -4,79 +4,74 @@ from loguru import logger
 
 from app_configs import UNSELECTED_VAR_NAME
 from components import commons
+from components import typed_dicts as td
 from components.model_pipeline.state_manager import (
     ModelStepUIState,
     PipelineState,
     PipelineStateManager,
     PipelineUIState,
+    TossupPipelineState,
+    TossupPipelineStateManager,
 )
 from components.model_step.model_step import ModelStepComponent
 from components.utils import make_state
-from workflows.structs import ModelStep, Workflow
+from workflows.structs import ModelStep, TossupWorkflow, Workflow
 from workflows.validators import WorkflowValidator
 
 
-def validate_simple_workflow(workflow: Workflow, required_output_variables: list[str]) -> Workflow:
-    """Validate the workflow."""
-    step = next(iter(workflow.steps.values()))
-    if not step.output_fields:
-        raise ValueError("No output fields found in the workflow")
-    output_field_names = {output.name for output in step.output_fields}
-    if not set(required_output_variables) <= output_field_names:
-        missing_vars = required_output_variables - output_field_names
-        raise ValueError(f"Missing required output variables: {missing_vars}")
-    return workflow
-
-
-def validate_complex_workflow(workflow: Workflow, required_output_variables: list[str]) -> Workflow:
-    """Validate the workflow."""
-    print("Validating complex workflow.")
-    return workflow
-    step = next(iter(workflow.steps.values()))
-    if not step.output_fields:
-        raise ValueError("No output fields found in the workflow")
-    output_field_names = {output.name for output in step.output_fields}
-    if not output_field_names <= set(required_output_variables):
-        missing_vars = output_field_names - set(required_output_variables)
-        raise ValueError(f"Missing required output variables: {missing_vars}")
-    return workflow
-
-
-def parse_yaml_workflow(yaml_str: str) -> Workflow:
-    """Parse a YAML workflow."""
-    workflow = yaml.safe_load(yaml_str)
-    return Workflow(**workflow)
-
-
-def update_workflow_from_code(yaml_str: str, ui_state: PipelineUIState) -> PipelineState:
-    """Update a workflow from a YAML string."""
-    workflow = parse_yaml_workflow(yaml_str)
-    ui_state = PipelineUIState.from_workflow(workflow)
-    return PipelineState(workflow=workflow, ui_state=ui_state)
-
-
 class PipelineInterface:
     """UI for the pipeline."""
 
     def __init__(
         self,
+        app: gr.Blocks,
         workflow: Workflow,
         ui_state: PipelineUIState | None = None,
         model_options: list[str] = None,
         simple: bool = False,
-        show_pipeline_selector: bool = False,
     ):
+        self.app = app
         self.model_options = model_options
         self.simple = simple
-        self.show_pipeline_selector = show_pipeline_selector
-        if not ui_state:
-            ui_state = PipelineUIState.from_workflow(workflow)
-        self.ui_state = make_state(ui_state)
-        self.pipeline_state = make_state(PipelineState(workflow=workflow, ui_state=ui_state))
+        ui_state = ui_state or PipelineUIState.from_workflow(workflow)
+
+        # Gradio States
+        self.workflow_state = make_state(workflow.model_dump())
         self.variables_state = make_state(workflow.get_available_variables())
-        self.model_selection_state = make_state({})
+        self.model_selection_state = make_state(workflow.get_model_selections())
+        self.pipeline_change = gr.State(False)
+
+        if isinstance(workflow, TossupWorkflow):
+            pipeline_state = TossupPipelineState(workflow=workflow, ui_state=ui_state)
+            self.sm = TossupPipelineStateManager()
+        else:
+            pipeline_state = PipelineState(workflow=workflow, ui_state=ui_state)
+            self.sm = PipelineStateManager()
+        self.pipeline_state = make_state(pipeline_state.model_dump())
+
+        def get_aux_states(pipeline_state_dict: td.PipelineStateDict):
+            """Get the auxiliary states for the pipeline."""
+            pipeline_state = self.sm.make_pipeline_state(pipeline_state_dict)
+            return (
+                pipeline_state.workflow.model_dump(),
+                pipeline_state.workflow.get_available_variables(),
+                pipeline_state.workflow.get_model_selections(),
+            )
 
-        self.sm = PipelineStateManager()
+        # Triggers for pipeline state changes
+        self.pipeline_state.change(
+            get_aux_states,
+            inputs=[self.pipeline_state],
+            outputs=[self.workflow_state, self.variables_state, self.model_selection_state],
+        )
+
+        self.workflow_state.change(
+            lambda x: logger.debug(f"Workflow state changed: {x}"),
+            inputs=[self.workflow_state],
+            outputs=[],
+        )
+
+        # IO Variables
         self.input_variables = workflow.inputs
         self.required_output_variables = list(workflow.outputs.keys())
 
@@ -108,13 +103,13 @@ class PipelineInterface:
             step_interface.on_model_step_change(
                 self.sm.update_model_step_state,
                 inputs=[self.pipeline_state, step_interface.model_step_state, step_interface.ui_state],
-                outputs=[self.pipeline_state, self.ui_state, self.variables_state, self.model_selection_state],
+                outputs=[self.pipeline_state],
             )
 
             step_interface.on_ui_change(
                 self.sm.update_model_step_ui,
                 inputs=[self.pipeline_state, step_interface.ui_state, gr.State(model_step.id)],
-                outputs=[self.pipeline_state, self.ui_state],
+                outputs=[self.pipeline_state],
             )
 
             if self.simple:
@@ -137,12 +132,20 @@ class PipelineInterface:
     def _assign_step_controls(self, buttons: tuple[gr.Button, gr.Button, gr.Button], position: int):
         up_button, down_button, remove_button = buttons
         position = gr.State(position)
-        up_button.click(self.sm.move_up, inputs=[self.ui_state, position], outputs=self.ui_state)
-        down_button.click(self.sm.move_down, inputs=[self.ui_state, position], outputs=self.ui_state)
+        up_button.click(
+            self.sm.move_up,
+            inputs=[self.pipeline_state, self.pipeline_change, position],
+            outputs=[self.pipeline_state, self.pipeline_change],
+        )
+        down_button.click(
+            self.sm.move_down,
+            inputs=[self.pipeline_state, self.pipeline_change, position],
+            outputs=[self.pipeline_state, self.pipeline_change],
+        )
         remove_button.click(
             self.sm.remove_step,
-            inputs=[self.pipeline_state, position],
-            outputs=[self.pipeline_state, self.ui_state, self.variables_state],
+            inputs=[self.pipeline_state, self.pipeline_change, position],
+            outputs=[self.pipeline_state, self.pipeline_change],
         )
 
     def _render_add_step_button(self, position: int):
@@ -153,8 +156,8 @@ class PipelineInterface:
             add_step_btn = gr.Button("➕ Add Step", elem_classes="add-step-button")
         add_step_btn.click(
             self.sm.add_step,
-            inputs=[self.pipeline_state, gr.State(position)],
-            outputs=[self.pipeline_state, self.ui_state, self.variables_state],
+            inputs=[self.pipeline_state, self.pipeline_change, gr.State(position)],
+            outputs=[self.pipeline_state, self.pipeline_change],
         )
         return add_step_btn
 
@@ -183,11 +186,9 @@ class PipelineInterface:
                     )
                     dropdowns[output_field] = dropdown
 
-        def update_choices(available_variables):
+        def update_choices(available_variables: list[str]):
             """Update the choices for the dropdowns"""
-            return [
-                gr.update(choices=available_variables, value=None, selected=None) for dropdown in dropdowns.values()
-            ]
+            return [gr.update(choices=available_variables, value=None, selected=None) for _ in dropdowns.values()]
 
         self.variables_state.change(
             update_choices,
@@ -196,16 +197,15 @@ class PipelineInterface:
         )
         return dropdowns
 
-    def validate_workflow(self, state: PipelineState) -> PipelineState:
+    def validate_workflow(self, state_dict: td.PipelineStateDict):
         """Validate the workflow."""
         try:
-            if self.simple:
-                workflow = validate_simple_workflow(state.workflow, self.required_output_variables)
-            else:
-                workflow = validate_complex_workflow(state.workflow, self.required_output_variables)
-            state.workflow = workflow
-            return state
+            state = self.sm.make_pipeline_state(state_dict)
+            WorkflowValidator().validate(state.workflow)
         except ValueError as e:
+            logger.exception(e)
+            state_dict_str = yaml.dump(state_dict, default_flow_style=False, indent=2)
+            logger.error(f"Could not validate workflow: \n{state_dict_str}")
             raise gr.Error(e)
 
     def _render_pipeline_header(self):
@@ -220,35 +220,35 @@ class PipelineInterface:
         gr.Markdown(f"* Input Variables: {input_variables_str}")
         gr.Markdown(f"* Output Variables: {output_variables_str}")
 
-        # if not self.simple:
-        #     self._render_add_step_button(0)
-
     def render(self):
         """Render the pipeline UI."""
         # Create a placeholder for all the step components
         self.all_components = []
 
-        # self.pipeline_state.change(
-        #     lambda x, y: print(f"Pipeline state changed! UI:\n{x}\n\n Data:\n{y}"),
-        #     inputs=[self.ui_state, self.pipeline_state],
-        #     outputs=[],
-        # )
-
         self._render_pipeline_header()
 
         # Function to render all steps
-        @gr.render(inputs=[self.pipeline_state, self.ui_state], concurrency_limit=1, concurrency_id="render_steps")
-        def render_steps(state: PipelineState, ui_state: PipelineUIState):
+        @gr.render(
+            triggers=[self.app.load, self.pipeline_change.change],
+            inputs=[self.pipeline_state],
+            concurrency_limit=1,
+            concurrency_id="render_steps",
+        )
+        def render_steps(pipeline_state: td.PipelineStateDict, evt: gr.EventData):
             """Render all steps in the pipeline"""
-            logger.info(f"Rerender triggered! Current UI State:{ui_state.model_dump()}")
-            workflow = state.workflow
+            logger.info(
+                f"Rerender triggered! \nInput Pipeline's UI State:{pipeline_state.get('ui_state')}\n Event: {evt.target} {evt._data}"
+            )
+            pipeline_state = self.sm.make_pipeline_state(pipeline_state)
+            ui_state = pipeline_state.ui_state
+            workflow = pipeline_state.workflow
             components = []
 
             step_objects = []  # Reset step objects list
             for i, step_id in enumerate(ui_state.step_ids):
                 step_data = workflow.steps[step_id]
                 step_ui_state = ui_state.steps[step_id]
-                available_variables = self.sm.get_all_variables(state, step_id)
+                available_variables = pipeline_state.get_available_variables(step_id)
                 sub_components = self._render_step(step_data, step_ui_state, available_variables, i, ui_state.n_steps)
                 step_objects.append(sub_components)
 
@@ -258,11 +258,13 @@ class PipelineInterface:
             self._render_add_step_button(-1)
 
         @gr.render(
+            triggers=[self.variables_state.change, self.app.load],
             inputs=[self.variables_state, self.pipeline_state],
             concurrency_limit=1,
             concurrency_id="render_output_fields",
         )
-        def render_output_fields(available_variables, pipeline_state):
+        def render_output_fields(available_variables: list[str], pipeline_state_dict: td.PipelineStateDict):
+            pipeline_state = self.sm.make_pipeline_state(pipeline_state_dict)
             logger.info(f"Rerendering output panel: {available_variables} {pipeline_state.workflow}")
             self._render_output_panel(available_variables, pipeline_state)
 
@@ -285,14 +287,14 @@ class PipelineInterface:
         # components.append(config_accordion)
 
         self.config_output.blur(
-            fn=update_workflow_from_code,
-            inputs=[self.config_output, self.ui_state],
-            outputs=[self.pipeline_state],
+            fn=self.sm.update_workflow_from_code,
+            inputs=[self.config_output, self.pipeline_change],
+            outputs=[self.pipeline_state, self.pipeline_change],
         )
 
         # Connect the export button to show the workflow JSON
         self.add_triggers_for_pipeline_export([export_btn.click], self.pipeline_state)
-        export_btn.click(self.validate_workflow, inputs=[self.pipeline_state], outputs=[self.pipeline_state]).success(
+        export_btn.click(self.validate_workflow, inputs=[self.pipeline_state], outputs=[]).success(
             fn=lambda: gr.update(visible=True, open=True), outputs=[self.config_accordion]
         )
 
@@ -301,7 +303,7 @@ class PipelineInterface:
             triggers,
             self.validate_workflow,
             inputs=[input_pipeline_state],
-            outputs=[self.pipeline_state],
+            outputs=[],
         ).success(
             fn=self.sm.get_formatted_config,
             inputs=[self.pipeline_state, gr.State("yaml")],

src/components/model_pipeline/state_manager.py

@@ -1,137 +1,26 @@
 import json
-from typing import Any, Literal
+from typing import Literal
 
-import gradio as gr
 import yaml
-from loguru import logger
-from pydantic import BaseModel, Field
 
+from app_configs import UNSELECTED_VAR_NAME
+from components import typed_dicts as td
 from components import utils
+from components.structs import ModelStepUIState, PipelineState, PipelineUIState, TossupPipelineState
 from workflows.factory import create_new_llm_step
-from workflows.structs import ModelStep, TossupWorkflow, Workflow
-
-
-def make_step_id(step_number: int):
-    """Make a step id from a step name."""
-    if step_number < 26:
-        return chr(ord("A") + step_number)
-    else:
-        # For more than 26 steps, use AA, AB, AC, etc.
-        first_char = chr(ord("A") + (step_number // 26) - 1)
-        second_char = chr(ord("A") + (step_number % 26))
-        return f"{first_char}{second_char}"
-
-
-def make_step_number(step_id: str):
-    """Make a step number from a step id."""
-    if len(step_id) == 1:
-        return ord(step_id) - ord("A")
-    else:
-        return (ord(step_id[0]) - ord("A")) * 26 + (ord(step_id[1]) - ord("A")) + 1
-
-
-class ModelStepUIState(BaseModel):
-    """Represents the UI state for a model step component."""
-
-    expanded: bool = True
-    active_tab: Literal["model-tab", "inputs-tab", "outputs-tab"] = "model-tab"
-
-    def update(self, key: str, value: Any) -> "ModelStepUIState":
-        """Update the UI state."""
-        new_state = self.model_copy(update={key: value})
-        return new_state
-
-
-class PipelineUIState(BaseModel):
-    """Represents the UI state for a pipeline component."""
-
-    step_ids: list[str] = Field(default_factory=list)
-    steps: dict[str, ModelStepUIState] = Field(default_factory=dict)
-
-    def model_post_init(self, __context: utils.Any) -> None:
-        if not self.steps and self.step_ids:
-            self.steps = {step_id: ModelStepUIState() for step_id in self.step_ids}
-        return super().model_post_init(__context)
-
-    def get_step_position(self, step_id: str):
-        """Get the position of a step in the pipeline."""
-        return next((i for i, step in enumerate(self.step_ids) if step == step_id), None)
-
-    @property
-    def n_steps(self) -> int:
-        """Get the number of steps in the pipeline."""
-        return len(self.step_ids)
-
-    @classmethod
-    def from_workflow(cls, workflow: Workflow):
-        """Create a pipeline UI state from a workflow."""
-        return PipelineUIState(
-            step_ids=list(workflow.steps.keys()),
-            steps={step_id: ModelStepUIState() for step_id in workflow.steps.keys()},
-        )
-
-
-class PipelineState(BaseModel):
-    """Represents the state for a pipeline component."""
-
-    workflow: Workflow
-    ui_state: PipelineUIState
-
-    def insert_step(self, position: int, step: ModelStep) -> "PipelineState":
-        if step.id in self.workflow.steps:
-            raise ValueError(f"Step {step.id} already exists in pipeline")
-
-        # Validate position
-        if position != -1 and (position < 0 or position > self.n_steps):
-            raise ValueError(f"Invalid position: {position}. Must be between 0 and {self.n_steps} or -1")
-
-        self.workflow.steps[step.id] = step
-
-        self.ui_state = self.ui_state.model_copy()
-        self.ui_state.steps[step.id] = ModelStepUIState()
-        if position == -1:
-            self.ui_state.step_ids.append(step.id)
-        else:
-            self.ui_state.step_ids.insert(position, step.id)
-        return self
-
-    def remove_step(self, position: int) -> "PipelineState":
-        step_id = self.ui_state.step_ids.pop(position)
-        self.workflow.steps.pop(step_id)
-        self.ui_state = self.ui_state.model_copy()
-        self.ui_state.steps.pop(step_id)
-        self.update_output_variables_mapping()
-        return self
-
-    def update_output_variables_mapping(self) -> "PipelineState":
-        available_variables = set(self.available_variables)
-        for output_field in self.workflow.outputs:
-            if self.workflow.outputs[output_field] not in available_variables:
-                self.workflow.outputs[output_field] = None
-        return self
-
-    @property
-    def available_variables(self) -> list[str]:
-        return self.workflow.get_available_variables()
-
-    @property
-    def n_steps(self) -> int:
-        return len(self.workflow.steps)
-
-    def get_new_step_id(self) -> str:
-        """Get a step ID for a new step."""
-        if not self.workflow.steps:
-            return "A"
-        else:
-            last_step_number = max(map(make_step_number, self.workflow.steps.keys()))
-            return make_step_id(last_step_number + 1)
+from workflows.structs import Buzzer, ModelStep, TossupWorkflow, Workflow
 
 
 class PipelineStateManager:
     """Manages a pipeline of multiple steps."""
 
-    def get_formatted_config(self, state: PipelineState, format: Literal["json", "yaml"] = "yaml"):
+    def make_pipeline_state(self, state_dict: td.PipelineStateDict) -> PipelineState:
+        """Make a state from a state dictionary."""
+        return PipelineState(**state_dict)
+
+    def get_formatted_config(self, state_dict: td.PipelineStateDict, format: Literal["json", "yaml"] = "yaml") -> str:
         """Get the full pipeline configuration."""
+        state = self.make_pipeline_state(state_dict)
         config = state.workflow.model_dump(exclude_defaults=True)
         if isinstance(state.workflow, TossupWorkflow):
             buzzer_config = state.workflow.buzzer.model_dump(exclude_defaults=False)
@@ -141,65 +30,121 @@ class PipelineStateManager:
         else:
             return json.dumps(config, indent=4, sort_keys=False)
 
-    def count_state(self):
-        return gr.State(len(self.steps))
-
-    def add_step(self, state: PipelineState, position: int = -1, name=""):
+    def add_step(
+        self, state_dict: td.PipelineStateDict, pipeline_change: bool, position: int = -1, name=""
+    ) -> td.PipelineStateDict:
         """Create a new step and return its state."""
+        state = self.make_pipeline_state(state_dict)
         step_id = state.get_new_step_id()
         step_name = name or f"Step {state.n_steps + 1}"
         new_step = create_new_llm_step(step_id=step_id, name=step_name)
         state = state.insert_step(position, new_step)
-        return state, state.ui_state, state.available_variables
+        return state.model_dump(), not pipeline_change
 
-    def remove_step(self, state: PipelineState, position: int):
+    def remove_step(
+        self, state_dict: td.PipelineStateDict, pipeline_change: bool, position: int
+    ) -> td.PipelineStateDict:
         """Remove a step from the pipeline."""
+        state = self.make_pipeline_state(state_dict)
         if 0 <= position < state.n_steps:
             state = state.remove_step(position)
         else:
             raise ValueError(f"Invalid step position: {position}")
-        return state, state.ui_state, state.available_variables
+        return state.model_dump(), not pipeline_change
 
-    def move_up(self, ui_state: PipelineUIState, position: int):
-        """Move a step up in the pipeline."""
-        utils.move_item(ui_state.step_ids, position, "up")
-        return ui_state.model_copy()
+    def _move_step(
+        self, state_dict: td.PipelineStateDict, position: int, direction: Literal["up", "down"]
+    ) -> tuple[td.PipelineStateDict, bool]:
+        state = self.make_pipeline_state(state_dict)
+        old_order = list(state.ui_state.step_ids)
+        utils.move_item(state.ui_state.step_ids, position, direction)
+        return state.model_dump(), old_order != list(state.ui_state.step_ids)
 
-    def move_down(self, ui_state: PipelineUIState, position: int):
+    def move_up(self, state_dict: td.PipelineStateDict, pipeline_change: bool, position: int) -> td.PipelineStateDict:
+        """Move a step up in the pipeline."""
+        new_state_dict, change = self._move_step(state_dict, position, "up")
+        if change:
+            pipeline_change = not pipeline_change
+        return new_state_dict, pipeline_change
+
+    def move_down(
+        self, state_dict: td.PipelineStateDict, pipeline_change: bool, position: int
+    ) -> td.PipelineStateDict:
         """Move a step down in the pipeline."""
-        utils.move_item(ui_state.step_ids, position, "down")
-        return ui_state.model_copy()
-
-    def update_model_step_state(self, state: PipelineState, model_step: ModelStep, ui_state: ModelStepUIState):
-        """Update a step in the pipeline."""
-        state.workflow.steps[model_step.id] = model_step.model_copy()
-        state.ui_state.steps[model_step.id] = ui_state.model_copy()
-        state.ui_state = state.ui_state.model_copy()
-        state.update_output_variables_mapping()
-        model_selections = {step_id: step.get_full_model_name() for step_id, step in state.workflow.steps.items()}
-        return state, state.ui_state, state.available_variables, model_selections
-
-    def update_output_variables(self, state: PipelineState, target: str, produced_variable: str):
-        if produced_variable == "Choose variable...":
+        new_state_dict, change = self._move_step(state_dict, position, "down")
+        if change:
+            pipeline_change = not pipeline_change
+        return new_state_dict, pipeline_change
+
+    def update_model_step_state(
+        self, state_dict: td.PipelineStateDict, model_step: ModelStep, ui_state: ModelStepUIState
+    ) -> td.PipelineStateDict:
+        """Update a particular model step in the pipeline."""
+        state = self.make_pipeline_state(state_dict)
+        state = state.update_step(model_step, ui_state)
+        return state.model_dump()
+
+    def update_output_variables(
+        self, state_dict: td.PipelineStateDict, target: str, produced_variable: str
+    ) -> td.PipelineStateDict:
+        if produced_variable == UNSELECTED_VAR_NAME:
             produced_variable = None
         """Update the output variables for a step."""
-        state.workflow.outputs.update({target: produced_variable})
-        return state
+        state = self.make_pipeline_state(state_dict)
+        state.workflow.outputs[target] = produced_variable
+        return state.model_dump()
 
-    def update_model_step_ui(self, state: PipelineState, step_ui: ModelStepUIState, step_id: str):
+    def update_model_step_ui(
+        self, state_dict: td.PipelineStateDict, step_ui: ModelStepUIState, step_id: str
+    ) -> td.PipelineStateDict:
         """Update a step in the pipeline."""
+        state = self.make_pipeline_state(state_dict)
         state.ui_state.steps[step_id] = step_ui.model_copy()
-        return state, state.ui_state
+        return state.model_dump()
 
-    def get_all_variables(self, state: PipelineState, model_step_id: str | None = None) -> list[str]:
+    def get_all_variables(self, state_dict: td.PipelineStateDict, model_step_id: str | None = None) -> list[str]:
         """Get all variables from all steps."""
-        available_variables = state.available_variables
-        if model_step_id is None:
-            return available_variables
-        else:
-            prefix = f"{model_step_id}."
-            return [var for var in available_variables if not var.startswith(prefix)]
-
-    def get_pipeline_config(self):
-        """Get the full pipeline configuration."""
-        return self.workflow
+        return self.make_pipeline_state(state_dict)
+
+    def parse_yaml_workflow(self, yaml_str: str) -> Workflow:
+        """Parse a YAML workflow."""
+        workflow = yaml.safe_load(yaml_str)
+        return Workflow(**workflow)
+
+    def update_workflow_from_code(self, yaml_str: str) -> td.PipelineStateDict:
+        """Update a workflow from a YAML string."""
+        workflow = self.parse_yaml_workflow(yaml_str)
+        return PipelineState.from_workflow(workflow).model_dump()
+
+
+class TossupPipelineStateManager(PipelineStateManager):
+    """Manages a tossup pipeline state."""
+
+    def make_pipeline_state(self, state_dict: td.PipelineStateDict) -> TossupPipelineState:
+        """Make a state from a state dictionary."""
+        return TossupPipelineState(**state_dict)
+
+    def parse_yaml_workflow(self, yaml_str: str) -> TossupWorkflow:
+        """Parse a YAML workflow."""
+        workflow = yaml.safe_load(yaml_str)
+        return TossupWorkflow(**workflow)
+
+    def update_workflow_from_code(self, yaml_str: str, change_state: bool) -> tuple[td.PipelineStateDict, bool]:
+        """Update a workflow from a YAML string."""
+        workflow = self.parse_yaml_workflow(yaml_str)
+        return TossupPipelineState.from_workflow(workflow).model_dump(), not change_state
+
+    def update_buzzer(
+        self,
+        state_dict: td.TossupPipelineStateDict,
+        confidence_threshold: float,
+        method: str,
+        tokens_prob: float | None,
+    ) -> td.TossupPipelineStateDict:
+        """Update the buzzer."""
+        state = self.make_pipeline_state(state_dict)
+        prob_threshold = float(tokens_prob) if tokens_prob and tokens_prob > 0 else None
+        state.workflow.buzzer = Buzzer(
+            method=method, confidence_threshold=confidence_threshold, prob_threshold=prob_threshold
+        )
+        return state.model_dump()

src/components/model_pipeline/tossup_pipeline.py

@@ -4,10 +4,13 @@ from loguru import logger
 
 from app_configs import AVAILABLE_MODELS, UNSELECTED_VAR_NAME
 from components import commons
+from components.structs import TossupPipelineState
+from components.typed_dicts import TossupPipelineStateDict
 from display.formatting import tiny_styled_warning
 from workflows.structs import Buzzer, TossupWorkflow
 
 from .model_pipeline import PipelineInterface, PipelineState, PipelineUIState
+from .state_manager import PipelineStateManager, TossupPipelineStateManager
 
 
 def toggleable_slider(
@@ -30,48 +33,34 @@ def toggleable_slider(
     return checkbox, slider
 
 
-class TossupPipelineState(PipelineState):
-    workflow: TossupWorkflow
-
-
 class TossupPipelineInterface(PipelineInterface):
     def __init__(
         self,
+        app: gr.Blocks,
         workflow: TossupWorkflow,
         ui_state: PipelineUIState | None = None,
         model_options: list[str] = None,
         simple: bool = False,
-        show_pipeline_selector: bool = False,
         defaults: dict = {},
     ):
-        super().__init__(workflow, ui_state, model_options, simple, show_pipeline_selector)
+        super().__init__(app, workflow, ui_state, model_options, simple)
         self.defaults = defaults
 
-    def update_buzzer(
-        self,
-        state: TossupPipelineState,
-        confidence_threshold: float,
-        method: str,
-        tokens_prob: float | None,
-    ):
-        """Update the buzzer."""
-
-        prob_threshold = float(tokens_prob) if tokens_prob and tokens_prob > 0 else None
-        state.workflow.buzzer = state.workflow.buzzer.model_copy(
-            update={
-                "method": method,
-                "confidence_threshold": confidence_threshold,
-                "prob_threshold": prob_threshold,
-            }
+        self.pipeline_state.change(
+            lambda x: logger.debug(
+                f"Pipeline state changed. Type: {type(x)}. Has buzzer info: {x['workflow']['buzzer'] if isinstance(x, dict) else 'N/A'}"
+            ),
+            inputs=[self.pipeline_state],
         )
-        Buzzer.model_validate(state.workflow.buzzer)
-        return state
 
-    def update_prob_slider(self, state: TossupPipelineState, answer_var: str, tokens_prob: float | None):
+    def update_prob_slider(
+        self, state_dict: TossupPipelineStateDict, answer_var: str, tokens_prob: float | None
+    ) -> tuple[TossupPipelineStateDict, dict, dict, dict]:
         """Update the probability slider based on the answer variable."""
+        state = TossupPipelineState(**state_dict)
         if answer_var == UNSELECTED_VAR_NAME:
             return (
-                state,
+                state.model_dump(),
                 gr.update(interactive=True),
                 gr.update(value="AND", interactive=True),
                 gr.update(visible=False),
@@ -83,14 +72,13 @@ class TossupPipelineInterface(PipelineInterface):
         buzzer = state.workflow.buzzer
         tokens_prob_threshold = tokens_prob if is_model_with_logprobs else None
         method = buzzer.method if is_model_with_logprobs else "AND"
-        state = self.update_buzzer(
-            state,
-            confidence_threshold=buzzer.confidence_threshold,
+        state.workflow.buzzer = Buzzer(
             method=method,
-            tokens_prob=tokens_prob_threshold,
+            confidence_threshold=buzzer.confidence_threshold,
+            prob_threshold=tokens_prob_threshold,
         )
         return (
-            state,
+            state.model_dump(),
             gr.update(interactive=is_model_with_logprobs),
             gr.update(value=method, interactive=is_model_with_logprobs),
             gr.update(
@@ -156,11 +144,9 @@ class TossupPipelineInterface(PipelineInterface):
                 )
             self.buzzer_warning_display = gr.HTML(visible=False)
 
-        def update_choices(available_variables):
+        def update_choices(available_variables: list[str]):
             """Update the choices for the dropdowns"""
-            return [
-                gr.update(choices=available_variables, value=None, selected=None) for dropdown in dropdowns.values()
-            ]
+            return [gr.update(choices=available_variables, value=None, selected=None) for _ in dropdowns.values()]
 
         self.variables_state.change(
             update_choices,
@@ -170,17 +156,12 @@ class TossupPipelineInterface(PipelineInterface):
 
         gr.on(
             triggers=[
-                self.confidence_slider.input,
+                self.confidence_slider.release,
                 self.buzzer_method_dropdown.input,
-                self.prob_slider.input,
-            ],
-            fn=self.update_buzzer,
-            inputs=[
-                self.pipeline_state,
-                self.confidence_slider,
-                self.buzzer_method_dropdown,
-                self.prob_slider,
+                self.prob_slider.release,
             ],
+            fn=self.sm.update_buzzer,
+            inputs=[self.pipeline_state, self.confidence_slider, self.buzzer_method_dropdown, self.prob_slider],
             outputs=[self.pipeline_state],
         )
 

src/components/model_step/model_step.py

@@ -5,7 +5,8 @@ import gradio as gr
 from gradio.components import FormComponent
 
 from app_configs import UNSELECTED_VAR_NAME
-from components.model_pipeline.state_manager import ModelStepUIState, PipelineState, PipelineStateManager
+from components.model_pipeline.state_manager import ModelStepUIState, PipelineStateManager
+from components.typed_dicts import PipelineStateDict
 from utils import get_full_model_name
 from workflows.structs import ModelStep
 
@@ -454,12 +455,12 @@ class ModelStepComponent(FormComponent):
         new_label = _make_accordion_label(new_model_step)
         return new_model_step, gr.update(label=new_label)
 
-    def refresh_variable_dropdowns(self, pipeline_state: PipelineState):
+    def refresh_variable_dropdowns(self, pipeline_state_dict: PipelineStateDict):
         # TODO: Fix this. Not sure why this is needed.
         """Refresh the variable dropdown options in all input rows."""
         variable_choices = []
         if self.pipeline_sm is not None:
-            variable_choices = self.pipeline_sm.get_all_variables(pipeline_state)
+            variable_choices = self.pipeline_sm.get_all_variables(pipeline_state_dict)
 
         for _, fields, _ in self.input_rows:
             _, inp_var, _ = fields

src/components/quizbowl/bonus.py

@@ -9,6 +9,7 @@ from loguru import logger
 from app_configs import UNSELECTED_PIPELINE_NAME
 from components import commons
 from components.model_pipeline.model_pipeline import PipelineInterface, PipelineState, PipelineUIState
+from components.typed_dicts import PipelineStateDict
 from display.formatting import styled_error
 from submission import submit
 from workflows.qb_agents import QuizBowlBonusAgent
@@ -128,6 +129,7 @@ class BonusInterface:
             self.pipeline_selector = commons.get_pipeline_selector([])
             self.load_btn = gr.Button("⬇️ Import Pipeline", variant="secondary")
         self.pipeline_interface = PipelineInterface(
+            self.app,
             workflow,
             simple=simple,
             model_options=list(self.model_options.keys()),
@@ -214,24 +216,28 @@ class BonusInterface:
         names = [UNSELECTED_PIPELINE_NAME] + populate.get_pipeline_names("bonus", profile)
         return gr.update(choices=names, value=UNSELECTED_PIPELINE_NAME)
 
-    def load_pipeline(self, model_name: str, profile: gr.OAuthProfile | None) -> tuple[str, PipelineState]:
+    def load_pipeline(
+        self, model_name: str, pipeline_change: bool, profile: gr.OAuthProfile | None
+    ) -> tuple[str, PipelineStateDict, bool, dict]:
         try:
-            pipeline_state = populate.load_pipeline("bonus", model_name, profile)
-            if pipeline_state is None:
-                return UNSELECTED_PIPELINE_NAME, gr.skip(), gr.update(visible=False)
-            return UNSELECTED_PIPELINE_NAME, pipeline_state, gr.update(visible=True)
+            workflow = populate.load_workflow("bonus", model_name, profile)
+            if workflow is None:
+                logger.warning(f"Could not load workflow for {model_name}")
+                return UNSELECTED_PIPELINE_NAME, gr.skip(), gr.skip(), gr.update(visible=False)
+            pipeline_state_dict = PipelineState.from_workflow(workflow).model_dump()
+            return UNSELECTED_PIPELINE_NAME, pipeline_state_dict, not pipeline_change, gr.update(visible=True)
         except Exception as e:
             error_msg = styled_error(f"Error loading pipeline: {str(e)}")
-            return UNSELECTED_PIPELINE_NAME, gr.skip(), gr.update(visible=True, value=error_msg)
+            return UNSELECTED_PIPELINE_NAME, gr.skip(), gr.skip(), gr.update(visible=True, value=error_msg)
 
     def single_run(
         self,
         question_id: int,
-        pipeline_state: PipelineState,
+        state_dict: PipelineStateDict,
     ) -> tuple[str, Any, Any]:
         """Run the agent in bonus mode."""
         try:
-            # Validate inputs
+            pipeline_state = PipelineState(**state_dict)
             question_id = int(question_id - 1)
             if not self.ds or question_id < 0 or question_id >= len(self.ds):
                 return "Invalid question ID or dataset not loaded", None, None
@@ -263,9 +269,10 @@ class BonusInterface:
                 gr.update(visible=True, value=error_msg),
             )
 
-    def evaluate(self, pipeline_state: PipelineState, progress: gr.Progress = gr.Progress()):
+    def evaluate(self, state_dict: PipelineStateDict, progress: gr.Progress = gr.Progress()):
         """Evaluate the bonus questions."""
         try:
+            pipeline_state = PipelineState(**state_dict)
             # Validate inputs
             if not self.ds or not self.ds.num_rows:
                 return "No dataset loaded", None, None
@@ -307,9 +314,14 @@ class BonusInterface:
             return gr.skip(), gr.update(visible=True, value=error_msg)
 
     def submit_model(
-        self, model_name: str, description: str, pipeline_state: PipelineState, profile: gr.OAuthProfile = None
+        self,
+        model_name: str,
+        description: str,
+        state_dict: PipelineStateDict,
+        profile: gr.OAuthProfile = None,
     ):
         """Submit the model output."""
+        pipeline_state = PipelineState(**state_dict)
         return submit.submit_model(model_name, description, pipeline_state.workflow, "bonus", profile)
 
     def _setup_event_listeners(self):
@@ -328,20 +340,19 @@ class BonusInterface:
             outputs=[self.pipeline_selector],
         )
 
-        self.new_loaded_pipeline_state = gr.State(value=None)
+        pipeline_state = self.pipeline_interface.pipeline_state
+        pipeline_change = self.pipeline_interface.pipeline_change
         self.load_btn.click(
             fn=self.load_pipeline,
-            inputs=[self.pipeline_selector],
-            outputs=[self.pipeline_selector, self.new_loaded_pipeline_state, self.error_display],
-        )
-        self.pipeline_interface.add_triggers_for_pipeline_export(
-            [self.new_loaded_pipeline_state.change], self.new_loaded_pipeline_state
+            inputs=[self.pipeline_selector, pipeline_change],
+            outputs=[self.pipeline_selector, pipeline_state, pipeline_change, self.error_display],
         )
+        self.pipeline_interface.add_triggers_for_pipeline_export([pipeline_state.change], pipeline_state)
 
         self.run_btn.click(
             self.pipeline_interface.validate_workflow,
             inputs=[self.pipeline_interface.pipeline_state],
-            outputs=[self.pipeline_interface.pipeline_state],
+            outputs=[],
         ).success(
             self.single_run,
             inputs=[

src/components/quizbowl/populate.py

@@ -4,7 +4,7 @@ import gradio as gr
 from loguru import logger
 
 from app_configs import UNSELECTED_PIPELINE_NAME
-from components.model_pipeline.model_pipeline import PipelineState, PipelineUIState
+from components.structs import TossupWorkflow, Workflow
 from display.formatting import styled_error
 from submission import submit
 
@@ -24,7 +24,9 @@ def get_pipeline_names(competition_type: str, profile: gr.OAuthProfile | None) -
     return all_names
 
 
-def load_pipeline(competition_type: str, model_name: str, profile: gr.OAuthProfile | None) -> Optional[PipelineState]:
+def load_workflow(
+    competition_type: str, model_name: str, profile: gr.OAuthProfile | None
+) -> Workflow | TossupWorkflow | None:
     if not model_name or model_name == UNSELECTED_PIPELINE_NAME:
         return None
     username, model_name = model_name.split("/")
@@ -35,4 +37,5 @@ def load_pipeline(competition_type: str, model_name: str, profile: gr.OAuthProfi
         workflow = submission.workflow
     else:
         raise gr.Error("Authentication required. Please log in to view your submissions.")
-    return PipelineState(workflow=workflow, ui_state=PipelineUIState.from_workflow(workflow))
+
+    return workflow

src/components/quizbowl/tossup.py

@@ -11,6 +11,7 @@ from app_configs import UNSELECTED_PIPELINE_NAME
 from components import commons
 from components.model_pipeline.model_pipeline import PipelineInterface, PipelineState, PipelineUIState
 from components.model_pipeline.tossup_pipeline import TossupPipelineInterface, TossupPipelineState
+from components.typed_dicts import PipelineStateDict, TossupPipelineStateDict
 from display.formatting import styled_error
 from submission import submit
 from workflows.qb_agents import QuizBowlTossupAgent, TossupResult
@@ -174,13 +175,7 @@ def validate_model_step(model_step: ModelStep):
 class TossupInterface:
     """Gradio interface for the Tossup mode."""
 
-    def __init__(
-        self,
-        app: gr.Blocks,
-        dataset: Dataset,
-        model_options: dict,
-        defaults: dict,
-    ):
+    def __init__(self, app: gr.Blocks, dataset: Dataset, model_options: dict, defaults: dict):
         """Initialize the Tossup interface."""
         logger.info(f"Initializing Tossup interface with dataset size: {len(dataset)}")
         self.ds = dataset
@@ -196,6 +191,7 @@ class TossupInterface:
             self.pipeline_selector = commons.get_pipeline_selector([])
             self.load_btn = gr.Button("⬇️ Import Pipeline", variant="secondary")
         self.pipeline_interface = TossupPipelineInterface(
+            self.app,
             workflow,
             simple=simple,
             model_options=list(self.model_options.keys()),
@@ -255,9 +251,10 @@ class TossupInterface:
 
         self._setup_event_listeners()
 
-    def validate_workflow(self, pipeline_state: PipelineState):
+    def validate_workflow(self, state_dict: TossupPipelineStateDict):
         """Validate the workflow."""
         try:
+            pipeline_state = TossupPipelineState(**state_dict)
             validate_workflow(pipeline_state.workflow)
         except Exception as e:
             raise gr.Error(f"Error validating workflow: {str(e)}")
@@ -293,20 +290,25 @@ class TossupInterface:
         names = [UNSELECTED_PIPELINE_NAME] + populate.get_pipeline_names("tossup", profile)
         return gr.update(choices=names, value=UNSELECTED_PIPELINE_NAME)
 
-    def load_pipeline(self, model_name: str, profile: gr.OAuthProfile | None) -> tuple[str, PipelineState]:
+    def load_pipeline(
+        self, model_name: str, pipeline_change: bool, profile: gr.OAuthProfile | None
+    ) -> tuple[str, PipelineStateDict, bool, dict]:
         try:
-            pipeline_state = populate.load_pipeline("tossup", model_name, profile)
-            if pipeline_state is None:
-                return UNSELECTED_PIPELINE_NAME, gr.skip(), gr.update(visible=False)
-            return UNSELECTED_PIPELINE_NAME, pipeline_state, gr.update(visible=True)
+            workflow = populate.load_workflow("tossup", model_name, profile)
+            if workflow is None:
+                logger.warning(f"Could not load workflow for {model_name}")
+                return UNSELECTED_PIPELINE_NAME, gr.skip(), gr.skip(), gr.update(visible=False)
+            pipeline_state_dict = TossupPipelineState.from_workflow(workflow).model_dump()
+            return UNSELECTED_PIPELINE_NAME, pipeline_state_dict, not pipeline_change, gr.update(visible=True)
         except Exception as e:
+            logger.exception(e)
             error_msg = styled_error(f"Error loading pipeline: {str(e)}")
-            return UNSELECTED_PIPELINE_NAME, gr.skip(), gr.update(visible=True, value=error_msg)
+            return UNSELECTED_PIPELINE_NAME, gr.skip(), gr.skip(), gr.update(visible=True, value=error_msg)
 
     def single_run(
         self,
         question_id: int,
-        pipeline_state: PipelineState,
+        state_dict: TossupPipelineStateDict,
         early_stop: bool = True,
     ) -> tuple[str, Any, Any]:
         """Run the agent in tossup mode with a system prompt."""
@@ -316,6 +318,7 @@ class TossupInterface:
             if not self.ds or question_id < 0 or question_id >= len(self.ds):
                 return "Invalid question ID or dataset not loaded", None, None
             example = self.ds[question_id]
+            pipeline_state = TossupPipelineState(**state_dict)
             outputs = self.get_model_outputs(example, pipeline_state, early_stop)
 
             # Process results and prepare visualization data
@@ -343,13 +346,13 @@ class TossupInterface:
                 gr.update(visible=True, value=error_msg),
             )
 
-    def evaluate(self, pipeline_state: PipelineState, progress: gr.Progress = gr.Progress()):
+    def evaluate(self, state_dict: TossupPipelineStateDict, progress: gr.Progress = gr.Progress()):
         """Evaluate the tossup questions."""
         try:
             # Validate inputs
             if not self.ds or not self.ds.num_rows:
                 return "No dataset loaded", None, None
-
+            pipeline_state = TossupPipelineState(**state_dict)
             buzz_counts = 0
             correct_buzzes = 0
             token_positions = []
@@ -389,9 +392,14 @@ class TossupInterface:
             )
 
     def submit_model(
-        self, model_name: str, description: str, pipeline_state: PipelineState, profile: gr.OAuthProfile = None
+        self,
+        model_name: str,
+        description: str,
+        state_dict: TossupPipelineStateDict,
+        profile: gr.OAuthProfile = None,
     ):
         """Submit the model output."""
+        pipeline_state = TossupPipelineState(**state_dict)
         return submit.submit_model(model_name, description, pipeline_state.workflow, "tossup", profile)
 
     def _setup_event_listeners(self):
@@ -408,20 +416,19 @@ class TossupInterface:
             outputs=[self.pipeline_selector],
         )
 
-        self.new_loaded_pipeline_state = gr.State(value=None)
+        pipeline_state = self.pipeline_interface.pipeline_state
+        pipeline_change = self.pipeline_interface.pipeline_change
         self.load_btn.click(
             fn=self.load_pipeline,
-            inputs=[self.pipeline_selector],
-            outputs=[self.pipeline_selector, self.new_loaded_pipeline_state, self.error_display],
-        )
-        self.pipeline_interface.add_triggers_for_pipeline_export(
-            [self.new_loaded_pipeline_state.change], self.new_loaded_pipeline_state
+            inputs=[self.pipeline_selector, pipeline_change],
+            outputs=[self.pipeline_selector, pipeline_state, pipeline_change, self.error_display],
         )
+        self.pipeline_interface.add_triggers_for_pipeline_export([pipeline_state.change], pipeline_state)
 
         self.run_btn.click(
             self.pipeline_interface.validate_workflow,
             inputs=[self.pipeline_interface.pipeline_state],
-            outputs=[self.pipeline_interface.pipeline_state],
+            outputs=[],
         ).success(
             self.single_run,
             inputs=[
@@ -454,9 +461,3 @@ class TossupInterface:
             ],
             outputs=[self.submit_status],
         )
-
-        self.hidden_input.change(
-            fn=update_tossup_plot,
-            inputs=[self.hidden_input, self.output_state],
-            outputs=[self.confidence_plot],
-        )

src/components/structs.py

@@ -0,0 +1,191 @@
+from typing import Any, Literal
+
+from pydantic import BaseModel, Field, model_validator
+
+from workflows.structs import ModelStep, TossupWorkflow, Workflow
+
+
+def make_step_id(step_number: int):
+    """Make a step id from a step name."""
+    if step_number < 26:
+        return chr(ord("A") + step_number)
+    else:
+        # For more than 26 steps, use AA, AB, AC, etc.
+        first_char = chr(ord("A") + (step_number // 26) - 1)
+        second_char = chr(ord("A") + (step_number % 26))
+        return f"{first_char}{second_char}"
+
+
+def make_step_number(step_id: str):
+    """Make a step number from a step id."""
+    if len(step_id) == 1:
+        return ord(step_id) - ord("A")
+    else:
+        return (ord(step_id[0]) - ord("A")) * 26 + (ord(step_id[1]) - ord("A")) + 1
+
+
+class ModelStepUIState(BaseModel):
+    """Represents the UI state for a model step component."""
+
+    expanded: bool = True
+    active_tab: Literal["model-tab", "inputs-tab", "outputs-tab"] = "model-tab"
+
+    class Config:
+        frozen = True
+
+    def update(self, key: str, value: Any) -> "ModelStepUIState":
+        """Update the UI state."""
+        return self.model_copy(update={key: value})
+
+
+class PipelineUIState(BaseModel):
+    """Represents the UI state for a pipeline component."""
+
+    step_ids: list[str] = Field(default_factory=list)
+    steps: dict[str, ModelStepUIState] = Field(default_factory=dict)
+
+    def model_post_init(self, __context: Any) -> None:
+        if not self.steps and self.step_ids:
+            self.steps = {step_id: ModelStepUIState() for step_id in self.step_ids}
+        return super().model_post_init(__context)
+
+    def get_step_position(self, step_id: str):
+        """Get the position of a step in the pipeline."""
+        return next((i for i, step in enumerate(self.step_ids) if step == step_id), None)
+
+    @property
+    def n_steps(self) -> int:
+        """Get the number of steps in the pipeline."""
+        return len(self.step_ids)
+
+    @classmethod
+    def from_workflow(cls, workflow: Workflow):
+        """Create a pipeline UI state from a workflow."""
+        return PipelineUIState(
+            step_ids=list(workflow.steps.keys()),
+            steps={step_id: ModelStepUIState() for step_id in workflow.steps.keys()},
+        )
+
+    @classmethod
+    def from_pipeline_state(cls, pipeline_state: "PipelineState"):
+        """Create a pipeline UI state from a pipeline state."""
+        return cls.from_workflow(pipeline_state.workflow)
+
+    # Update methods
+
+    def insert_step(self, step_id: str, position: int = -1) -> "PipelineUIState":
+        """Insert a step into the pipeline at the given position."""
+        if position == -1:
+            position = len(self.step_ids)
+        self.step_ids.insert(position, step_id)
+        steps = self.steps | {step_id: ModelStepUIState()}
+        return self.model_copy(update={"step_ids": self.step_ids, "steps": steps})
+
+    def remove_step(self, step_id: str) -> "PipelineUIState":
+        """Remove a step from the pipeline."""
+        self.step_ids.remove(step_id)
+        self.steps.pop(step_id)
+        return self.model_copy(update={"step_ids": self.step_ids, "steps": self.steps})
+
+    def update_step(self, step_id: str, ui_state: ModelStepUIState) -> "PipelineUIState":
+        """Update a step in the pipeline."""
+        if step_id not in self.steps:
+            raise ValueError(f"Step {step_id} not found in pipeline")
+        return self.model_copy(update={"steps": self.steps | {step_id: ui_state}})
+
+
+class PipelineState(BaseModel):
+    """Represents the state for a pipeline component."""
+
+    workflow: Workflow
+    ui_state: PipelineUIState
+
+    @classmethod
+    def from_workflow(cls, workflow: Workflow):
+        """Create a pipeline state from a workflow."""
+        return cls(workflow=workflow, ui_state=PipelineUIState.from_workflow(workflow))
+
+    def update_workflow(self, workflow: Workflow) -> "PipelineState":
+        return self.model_copy(update={"workflow": workflow})
+
+    def insert_step(self, position: int, step: ModelStep) -> "PipelineState":
+        if step.id in self.workflow.steps:
+            raise ValueError(f"Step {step.id} already exists in pipeline")
+
+        # Validate position
+        if position != -1 and (position < 0 or position > self.n_steps):
+            raise ValueError(f"Invalid position: {position}. Must be between 0 and {self.n_steps} or -1")
+
+        # Create a new workflow with updated steps
+        workflow = self.workflow.add_step(step)
+
+        self.ui_state = self.ui_state.insert_step(step.id, position)
+
+        # Return a new PipelineState with the updated workflow
+        return self.model_copy(update={"workflow": workflow, "ui_state": self.ui_state})
+
+    def remove_step(self, position: int) -> "PipelineState":
+        step_id = self.ui_state.step_ids.pop(position)
+
+        workflow = self.workflow.remove_step(step_id)
+        self.ui_state = self.ui_state.remove_step(step_id)
+
+        # Return a new PipelineState with the updated workflow
+        updated_outputs = self.get_output_variables_updates(workflow)
+        return self.model_copy(update={"workflow": workflow, "outputs": updated_outputs})
+
+    def update_step(self, step: ModelStep, ui_state: ModelStepUIState | None = None) -> "PipelineState":
+        """Update a step in the pipeline."""
+        if step.id not in self.workflow.steps:
+            raise ValueError(f"Step {step.id} not found in pipeline")
+        steps = self.workflow.steps | {step.id: step}
+        workflow = self.workflow.model_copy(update={"steps": steps})
+        update = {"workflow": workflow, "outputs": self.get_output_variables_updates(workflow)}
+        if ui_state is not None:
+            update["ui_state"] = self.ui_state.update_step(step.id, ui_state)
+        return self.model_copy(update=update)
+
+    def get_output_variables_updates(self, new_workflow: Workflow) -> dict[str, str | None]:
+        available_variables = set(self.available_variables)
+        updated_outputs = new_workflow.outputs.copy()
+        for output_field in updated_outputs:
+            if updated_outputs[output_field] not in available_variables:
+                updated_outputs[output_field] = None
+        return updated_outputs
+
+    def update_output_variables_mapping(self) -> "PipelineState":
+        updated_outputs = self.get_output_variables_updates(self.workflow)
+
+        # Create a new workflow with updated outputs
+        workflow = self.workflow.model_copy(update={"outputs": updated_outputs})
+
+        # Return a new PipelineState with the updated workflow
+        return self.model_copy(update={"workflow": workflow})
+
+    def get_available_variables(self, model_step_id: str | None = None) -> list[str]:
+        """Get all variables from all steps."""
+        available_variables = self.available_variables
+        if model_step_id is None:
+            return available_variables
+        prefix = f"{model_step_id}."
+        return [var for var in available_variables if not var.startswith(prefix)]
+
+    @property
+    def available_variables(self) -> list[str]:
+        return self.workflow.get_available_variables()
+
+    @property
+    def n_steps(self) -> int:
+        return len(self.workflow.steps)
+
+    def get_new_step_id(self) -> str:
+        """Get a step ID for a new step."""
+        if not self.workflow.steps:
+            return "A"
+        else:
+            last_step_number = max(map(make_step_number, self.workflow.steps.keys()))
+            return make_step_id(last_step_number + 1)
+
+
+class TossupPipelineState(PipelineState):
+    workflow: TossupWorkflow

src/components/typed_dicts.py

@@ -0,0 +1,64 @@
+from typing import Any, Dict, List, Literal, Optional, TypedDict, Union
+
+
+# TypedDicts for workflows/structs.py
+class InputFieldDict(TypedDict):
+    name: str
+    description: str
+    variable: str
+    func: Optional[str]
+
+
+class OutputFieldDict(TypedDict):
+    name: str
+    type: Literal["str", "int", "float", "bool", "list[str]", "list[int]", "list[float]", "list[bool]"]
+    description: str
+    func: Optional[str]
+
+
+class ModelStepDict(TypedDict):
+    id: str
+    name: str
+    model: str
+    provider: str
+    call_type: Literal["llm", "search", "python_func"]
+    temperature: Optional[float]
+    system_prompt: str
+    input_fields: List[InputFieldDict]
+    output_fields: List[OutputFieldDict]
+
+
+class WorkflowDict(TypedDict):
+    inputs: List[str]
+    outputs: Dict[str, Optional[str]]
+    steps: Dict[str, ModelStepDict]
+
+
+class BuzzerDict(TypedDict):
+    method: Literal["AND", "OR"]
+    confidence_threshold: float
+    prob_threshold: Optional[float]
+
+
+class TossupWorkflowDict(WorkflowDict):
+    buzzer: BuzzerDict
+
+
+# TypedDicts for components/model_pipeline/state_manager.py
+class ModelStepUIStateDict(TypedDict):
+    expanded: bool
+    active_tab: Literal["model-tab", "inputs-tab", "outputs-tab"]
+
+
+class PipelineUIStateDict(TypedDict):
+    step_ids: List[str]
+    steps: Dict[str, ModelStepUIStateDict]
+
+
+class PipelineStateDict(TypedDict):
+    workflow: WorkflowDict
+    ui_state: PipelineUIStateDict
+
+
+class TossupPipelineStateDict(PipelineStateDict):
+    workflow: TossupWorkflowDict

src/display/guide.py

@@ -2,7 +2,9 @@
 
 GUIDE_MARKDOWN = """
 # 🎯 Quizbowl Bot Guide
+"""
 
+QUICKSTART_MARKDOWN = """
 ## Quick Start
 1. Choose between Tossup or Bonus mode
 2. Design your pipeline
@@ -35,7 +37,9 @@ GUIDE_MARKDOWN = """
   - `answer`: Your predicted answer
   - `confidence`: Score between 0-1
   - `explanation`: Brief justification for human collaboration
+"""
 
+BUILDING_MARKDOWN = """
 ## Building Your First Pipeline
 
 ### 1. Simple Pipeline (Recommended for First Submission)

src/submission/structs.py

@@ -37,6 +37,7 @@ class Submission(BaseModel):
     description: str = Field(description="Detailed description of what the submission does")
     competition_type: CompetitionType = Field(description="Type of competition (tossup or bonus)")
     submission_type: SubmissionType = Field(description="Format of the submission (python file or workflow)")
+    # TODO: Make workflow as json / yaml string instead of Workflow object
     workflow: Optional[Workflow] = Field(default=None, description="Optional workflow definition stored as JSON dict")
     code: Optional[str] = Field(default=None, description="Optional code content for python file submissions")
     status: SubmissionStatus = Field(description="Current status of the submission")

src/workflows/structs.py

@@ -1,4 +1,5 @@
 # %%
+from copy import deepcopy
 from enum import Enum
 from typing import Any, Literal, Optional
 
@@ -48,6 +49,9 @@ class InputField(BaseModel):
     # function to call on the input before passing it to the model
     func: str | None = None
 
+    class Config:
+        frozen = True
+
 
 class OutputField(BaseModel):
     """
@@ -70,6 +74,9 @@ class OutputField(BaseModel):
     # function to call on the output string from the model
     func: str | None = None
 
+    class Config:
+        frozen = True
+
 
 class CallType(str, Enum):
     LLM = "llm"
@@ -120,6 +127,7 @@ class ModelStep(BaseModel):
         return [f"{self.id}.{field.name}" for field in self.output_fields if field.name]
 
     def update(self, update: dict[str, Any]) -> "ModelStep":
+        """Returns a new copy with the updated properties."""
         return self.model_copy(update=update)
 
     def update_property(self, field: str, value: Any) -> "ModelStep":
@@ -157,11 +165,16 @@ class ModelStep(BaseModel):
         Returns:
             A new ModelStep with the updated fields.
         """
-        new_step = self.model_copy()
-        fields = new_step.input_fields if field_type == "input" else new_step.output_fields
-        new_field = ModelStep.create_new_field(field_type, input_var)
-        fields.insert(index + 1, new_field) if index != -1 else fields.append(new_field)
-        return new_step
+        if field_type == "input":
+            fields = deepcopy(self.input_fields)
+            new_field = ModelStep.create_new_field(field_type, input_var)
+            fields.insert(index + 1, new_field) if index != -1 else fields.append(new_field)
+            return self.model_copy(update={"input_fields": fields})
+        else:
+            fields = deepcopy(self.output_fields)
+            new_field = ModelStep.create_new_field(field_type)
+            fields.insert(index + 1, new_field) if index != -1 else fields.append(new_field)
+            return self.model_copy(update={"output_fields": fields})
 
     def delete_field(self, field_type: FieldType, index: int) -> "ModelStep":
         """
@@ -174,10 +187,10 @@ class ModelStep(BaseModel):
         Returns:
             A new ModelStep with the updated fields.
         """
-        new_step = self.model_copy()
-        fields = new_step.input_fields if field_type == "input" else new_step.output_fields
+        fields = self.input_fields if field_type == "input" else self.output_fields
+        fields = deepcopy(fields)
         fields.pop(index)
-        return new_step
+        return self.model_copy(update={"input_fields": fields} if field_type == "input" else {"output_fields": fields})
 
 
 class Workflow(BaseModel):
@@ -242,6 +255,22 @@ class Workflow(BaseModel):
             variables.update(self.get_step_variables(step.id))
         return list(variables)
 
+    def get_model_selections(self) -> dict[str, str]:
+        """Get all model selections for all steps."""
+        return {step_id: step.get_full_model_name() for step_id, step in self.steps.items()}
+
+    # Step update method
+
+    def add_step(self, step: ModelStep) -> "Workflow":
+        """Add a step to the workflow."""
+        steps = self.steps | {step.id: step}
+        return self.model_copy(update={"steps": steps})
+
+    def remove_step(self, step_id: str) -> "Workflow":
+        """Remove a step from the workflow."""
+        self.steps.pop(step_id)
+        return self.model_copy(update={"steps": self.steps})
+
 
 class BuzzerMethod(str, Enum):
     AND = "AND"
@@ -257,6 +286,7 @@ class Buzzer(BaseModel):
 
     class Config:
         use_enum_values = True
+        frozen = True
 
     def run(self, confidence: float, prob: float | None = None, logprob: float | None = None) -> bool:
         """Run the buzzer logic."""
@@ -285,3 +315,7 @@ class TossupWorkflow(Workflow):
     """Workflow specialized for tossup questions with buzzing capability."""
 
     buzzer: Buzzer
+
+    def update_buzzer(self, buzzer: Buzzer) -> "TossupWorkflow":
+        """Update the buzzer."""
+        return self.model_copy(update={"buzzer": buzzer})