Bugfix: Buzzer update on pipeline change.

src/components/model_pipeline/model_pipeline.py

@@ -161,8 +161,9 @@ class PipelineInterface:
         )
         return add_step_btn
 
-    def _render_output_panel(self, available_variables: list[str], pipeline_state: PipelineState):
+    def _render_output_panel(self, pipeline_state: PipelineState):
         dropdowns = {}
+        available_variables = pipeline_state.workflow.get_available_variables()
         variable_options = [UNSELECTED_VAR_NAME] + [v for v in available_variables if v not in self.input_variables]
         with gr.Column(elem_classes="step-accordion control-panel"):
             commons.get_panel_header(
@@ -258,15 +259,15 @@ 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],
+            triggers=[self.workflow_state.change, self.app.load],
+            inputs=[self.pipeline_state],
             concurrency_limit=1,
             concurrency_id="render_output_fields",
         )
-        def render_output_fields(available_variables: list[str], pipeline_state_dict: td.PipelineStateDict):
+        def render_output_fields(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)
+            logger.info(f"Rerendering output panel: {pipeline_state.workflow}")
+            self._render_output_panel(pipeline_state)
 
         export_btn = gr.Button("Export Pipeline", elem_classes="export-button")
         # components.append(export_btn)

src/components/model_pipeline/tossup_pipeline.py

@@ -13,6 +13,22 @@ from .model_pipeline import PipelineInterface, PipelineState, PipelineUIState
 from .state_manager import PipelineStateManager, TossupPipelineStateManager
 
 
+def get_probs_model_name(workflow: TossupWorkflow, answer_var: str | None = None) -> str | None:
+    if answer_var is None:
+        answer_var = workflow.outputs["answer"]
+    if answer_var is None or answer_var == UNSELECTED_VAR_NAME:
+        return None
+    step_id = answer_var.split(".")[0]
+    return workflow.steps[step_id].get_full_model_name()
+
+
+def is_logprobs_supported(workflow: TossupWorkflow, answer_var: str | None = None) -> bool:
+    model_name = get_probs_model_name(workflow, answer_var)
+    if model_name is None:
+        return True
+    return AVAILABLE_MODELS[model_name].get("logprobs", False)
+
+
 def toggleable_slider(
     value, minimum, maximum, step, toggle_value=False, label=None, info=None, min_width=200, scale=1
 ):
@@ -65,32 +81,70 @@ class TossupPipelineInterface(PipelineInterface):
                 gr.update(value="AND", interactive=True),
                 gr.update(visible=False),
             )
-        step_id = answer_var.split(".")[0]
-        model_name = state.workflow.steps[step_id].get_full_model_name()
-        model_config = AVAILABLE_MODELS[model_name]
-        is_model_with_logprobs = model_config.get("logprobs", False)
+        logprobs_supported = is_logprobs_supported(state.workflow, answer_var)
         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"
+        tokens_prob_threshold = tokens_prob if logprobs_supported else None
+        method = buzzer.method if logprobs_supported else "AND"
         state.workflow.buzzer = Buzzer(
             method=method,
             confidence_threshold=buzzer.confidence_threshold,
             prob_threshold=tokens_prob_threshold,
         )
+        model_name = get_probs_model_name(state.workflow, answer_var)
         return (
             state.model_dump(),
-            gr.update(interactive=is_model_with_logprobs),
-            gr.update(value=method, interactive=is_model_with_logprobs),
+            gr.update(interactive=logprobs_supported),
+            gr.update(value=method, interactive=logprobs_supported),
             gr.update(
                 value=tiny_styled_warning(
-                    f"{model_name} does not support `logprobs`. The probability slider will be disabled."
+                    f"<code>{model_name}</code> does not support <code>logprobs</code>. The probability slider will be disabled."
                 ),
-                visible=not is_model_with_logprobs,
+                visible=not logprobs_supported,
             ),
         )
 
-    def _render_output_panel(self, available_variables: list[str], pipeline_state: TossupPipelineState):
+    def _render_buzzer_panel(
+        self, buzzer: Buzzer, prob_slider_supported: bool, selected_model_name: str | None = None
+    ):
+        with gr.Row(elem_classes="control-panel"):
+            self.confidence_slider = gr.Slider(
+                minimum=0.0,
+                maximum=1.0,
+                value=buzzer.confidence_threshold,
+                step=0.01,
+                label="Confidence",
+                elem_classes="slider-container",
+                show_reset_button=False,
+            )
+            value = buzzer.method if prob_slider_supported else "AND"
+            self.buzzer_method_dropdown = gr.Dropdown(
+                choices=["AND", "OR"],
+                value=value,
+                label="Method",
+                interactive=prob_slider_supported,
+                min_width=80,
+                scale=0,
+            )
+            self.prob_slider = gr.Slider(
+                value=buzzer.prob_threshold or 0.0,
+                interactive=prob_slider_supported,
+                label="Probability",
+                minimum=0.0,
+                maximum=1.0,
+                step=0.001,
+                elem_classes="slider-container",
+                show_reset_button=False,
+            )
+        display_html = ""
+        if selected_model_name is not None:
+            display_html = tiny_styled_warning(
+                f"<code>{selected_model_name}</code> does not support <code>logprobs</code>. The probability slider will be disabled."
+            )
+        self.buzzer_warning_display = gr.HTML(display_html, visible=not prob_slider_supported)
+
+    def _render_output_panel(self, pipeline_state: TossupPipelineState):
         dropdowns = {}
+        available_variables = pipeline_state.workflow.get_available_variables()
         variable_options = [UNSELECTED_VAR_NAME] + [v for v in available_variables if v not in self.input_variables]
         with gr.Column(elem_classes="step-accordion control-panel"):
             commons.get_panel_header(
@@ -98,7 +152,8 @@ class TossupPipelineInterface(PipelineInterface):
             )
             with gr.Row(elem_classes="output-fields-row"):
                 for output_field in self.required_output_variables:
-                    value = pipeline_state.workflow.outputs.get(output_field, UNSELECTED_VAR_NAME)
+                    value = pipeline_state.workflow.outputs.get(output_field)
+                    value = value or UNSELECTED_VAR_NAME
                     dropdown = gr.Dropdown(
                         label=output_field,
                         value=value,
@@ -117,42 +172,19 @@ class TossupPipelineInterface(PipelineInterface):
                 header="Buzzer settings:",
                 subheader="Set your thresholds for confidence and output tokens probability.",
             )
-            with gr.Row(elem_classes="control-panel"):
-                self.confidence_slider = gr.Slider(
-                    minimum=0.0,
-                    maximum=1.0,
-                    value=self.defaults.get("confidence_threshold", 0.85),
-                    step=0.01,
-                    label="Confidence",
-                    elem_classes="slider-container",
-                )
-                self.buzzer_method_dropdown = gr.Dropdown(
-                    choices=["AND", "OR"],
-                    value=self.defaults.get("buzzer_method", "AND"),
-                    label="Method",
-                    interactive=True,
-                    min_width=80,
-                    scale=0,
-                )
-                self.prob_slider = gr.Slider(
-                    value=self.defaults.get("logits_prob", 0.0),
-                    label="Probability",
-                    minimum=0.0,
-                    maximum=1.0,
-                    step=0.001,
-                    elem_classes="slider-container",
-                )
-            self.buzzer_warning_display = gr.HTML(visible=False)
-
-        def update_choices(available_variables: list[str]):
-            """Update the choices for the dropdowns"""
-            return [gr.update(choices=available_variables, value=None, selected=None) for _ in dropdowns.values()]
-
-        self.variables_state.change(
-            update_choices,
-            inputs=[self.variables_state],
-            outputs=list(dropdowns.values()),
-        )
+            logprobs_supported = is_logprobs_supported(pipeline_state.workflow)
+            selected_model_name = get_probs_model_name(pipeline_state.workflow)
+            self._render_buzzer_panel(pipeline_state.workflow.buzzer, logprobs_supported, selected_model_name)
+
+        # def update_choices(available_variables: list[str]):
+        #     """Update the choices for the dropdowns"""
+        #     return [gr.update(choices=available_variables, value=None, selected=None) for _ in dropdowns.values()]
+
+        # self.variables_state.change(
+        #     update_choices,
+        #     inputs=[self.variables_state],
+        #     outputs=list(dropdowns.values()),
+        # )
 
         gr.on(
             triggers=[

src/components/structs.py

@@ -128,40 +128,19 @@ class PipelineState(BaseModel):
         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})
+        ui_state = self.ui_state.remove_step(step_id)
+        return self.model_copy(update={"workflow": workflow, "ui_state": ui_state})
 
     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)}
+        workflow = self.workflow.update_step(step)
+        update = {"workflow": 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

src/workflows/structs.py

@@ -269,7 +269,24 @@ class Workflow(BaseModel):
     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})
+        workflow = self.model_copy(update={"steps": self.steps})
+        workflow.refresh_output_variables()
+        return workflow
+
+    def update_step(self, step: ModelStep) -> "Workflow":
+        """Update a step in the workflow."""
+        self.steps[step.id] = step
+        steps = self.steps | {step.id: step}
+        workflow = self.model_copy(update={"steps": steps})
+        workflow.refresh_output_variables()
+        return workflow
+
+    # Output variables
+    def refresh_output_variables(self) -> "Workflow":
+        """Refresh the output variables for the workflow."""
+        produced_variables = self.get_available_variables()
+        self.outputs = {k: (v if v in produced_variables else None) for k, v in self.outputs.items()}
+        return self
 
 
 class BuzzerMethod(str, Enum):