Added more validation layers

src/components/model_pipeline/model_pipeline.py

@@ -8,12 +8,11 @@ from components import typed_dicts as td
 from components.model_pipeline.state_manager import (
     BasePipelineValidator,
     PipelineStateManager,
-    TossupPipelineStateManager,
 )
 from components.model_step.model_step import ModelStepComponent
-from components.structs import ModelStepUIState, PipelineState, PipelineUIState, TossupPipelineState
+from components.structs import ModelStepUIState, PipelineState, PipelineUIState
 from components.utils import make_state
-from workflows.structs import ModelStep, TossupWorkflow, Workflow
+from workflows.structs import ModelStep, Workflow
 from workflows.validators import WorkflowValidationError, WorkflowValidator
 
 from .state_manager import get_output_panel_state
@@ -24,6 +23,8 @@ DEFAULT_MAX_TEMPERATURE = 5.0
 class PipelineInterface:
     """UI for the pipeline."""
 
+    state_manager_cls = PipelineStateManager
+
     def __init__(
         self,
         app: gr.Blocks,
@@ -43,17 +44,14 @@ class PipelineInterface:
         self.workflow_state = make_state(workflow.model_dump())
         self.variables_state = make_state(workflow.get_available_variables())
         self.output_panel_state = make_state(get_output_panel_state(workflow))
+        self.ui_validator = validator
 
         # Maintains the toggle state change for pipeline changes through user input.
         self.pipeline_change = gr.State(False)
 
-        if isinstance(workflow, TossupWorkflow):
-            pipeline_state = TossupPipelineState(workflow=workflow, ui_state=ui_state)
-            self.sm = TossupPipelineStateManager(validator)
-        else:
-            pipeline_state = PipelineState(workflow=workflow, ui_state=ui_state)
-            self.sm = PipelineStateManager(validator)
-        self.pipeline_state = make_state(pipeline_state.model_dump())
+        self.sm = self.state_manager_cls(validator)
+        pipeline_state_dict = self.sm.create_pipeline_state_dict(workflow=workflow, ui_state=ui_state)
+        self.pipeline_state = make_state(pipeline_state_dict)
 
         def get_aux_states(pipeline_state_dict: td.PipelineStateDict):
             """Get the auxiliary states for the pipeline."""
@@ -164,20 +162,15 @@ class PipelineInterface:
         )
         return add_step_btn
 
-    def validate_workflow(self, state_dict: td.PipelineStateDict):
+    def validate_workflow_ui(self, state_dict: td.PipelineStateDict):
         """Validate the workflow."""
         try:
             state = self.sm.make_pipeline_state(state_dict)
-            validator = WorkflowValidator(
-                max_temperature=self.config.get("max_temperature", 10),
-            )
-            if not validator.validate(state.workflow):
-                raise WorkflowValidationError(validator.errors)
+            self.ui_validator(state.workflow)
+            return gr.update(visible=False)
         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)
+            return gr.update(visible=True, value=str(e))
 
     def _render_pipeline_header(self):
         # Add Step button at top
@@ -270,9 +263,8 @@ class PipelineInterface:
         )
 
         # Connect the export button to show the workflow JSON
-        self.add_triggers_for_pipeline_export([export_btn.click], self.pipeline_state, scroll=True)
-        export_btn.click(self.validate_workflow, inputs=[self.pipeline_state], outputs=[]).success(
-            fn=lambda: gr.update(visible=True, open=True), outputs=[self.config_accordion]
+        self.add_triggers_for_pipeline_export(
+            [export_btn.click], self.pipeline_state, scroll=True, expand_accordion=True
         )
 
     def render(self):
@@ -315,21 +307,30 @@ class PipelineInterface:
 
         self._render_pipeline_preview()
 
-    def add_triggers_for_pipeline_export(self, triggers: list, input_pipeline_state: gr.State, scroll: bool = False):
+    def add_triggers_for_pipeline_export(
+        self,
+        triggers: list,
+        input_pipeline_state: gr.State,
+        scroll: bool = False,
+        expand_accordion: bool = False,
+    ):
         js = None
         if scroll:
             js = "() => {document.querySelector('.pipeline-preview').scrollIntoView({behavior: 'smooth'})}"
 
         # TODO: modify this validate workflow to user input level and not executable label.
         # (workflows that can be converted to UI interface, no logical validation)
-        gr.on(
+        event = gr.on(
             triggers,
-            self.validate_workflow,
+            self.validate_workflow_ui,
             inputs=[input_pipeline_state],
-            outputs=[],
+            outputs=[self.error_display],
         ).success(
             fn=self.sm.get_formatted_config,
             inputs=[self.pipeline_state, gr.State("yaml")],
             outputs=[self.config_output, self.error_display],
             js=js,
         )
+
+        if expand_accordion:
+            event.then(fn=lambda: gr.update(visible=True, open=True), outputs=[self.config_accordion])

src/components/model_pipeline/state_manager.py

@@ -11,7 +11,7 @@ from pydantic import BaseModel, ValidationError
 from app_configs import UNSELECTED_VAR_NAME
 from components import typed_dicts as td
 from components import utils
-from components.structs import ModelStepUIState, PipelineState, TossupPipelineState
+from components.structs import ModelStepUIState, PipelineState, PipelineUIState, TossupPipelineState
 from envs import DOCS_REPO_BRANCH, DOCS_REPO_URL
 from workflows.factory import create_new_llm_step
 from workflows.structs import Buzzer, BuzzerMethod, ModelStep, TossupWorkflow, Workflow
@@ -72,6 +72,10 @@ class PipelineStateManager:
         """Make a state from a state dictionary."""
         return self.pipeline_state_cls(**state_dict)
 
+    def create_pipeline_state_dict(self, workflow: Workflow, ui_state: PipelineUIState) -> td.PipelineStateDict:
+        """Create a pipeline state from a workflow."""
+        return self.pipeline_state_cls(workflow=workflow, ui_state=ui_state).model_dump()
+
     def add_step(
         self, state_dict: td.PipelineStateDict, pipeline_change: bool, position: int = -1, name=""
     ) -> td.PipelineStateDict:
@@ -212,6 +216,7 @@ class PipelineStateManager:
         """Update a workflow from a YAML string."""
         try:
             workflow = self.parse_yaml_workflow(yaml_str, strict=True)
+            logger.debug(f"Validator: {self.validator}")
             self.validator and self.validator(workflow)
             state = self.pipeline_state_cls.from_workflow(workflow)
             return state.model_dump(), not change_state, gr.update(visible=False)
@@ -229,6 +234,11 @@ class TossupPipelineStateManager(PipelineStateManager):
     def make_pipeline_state(self, state_dict: td.PipelineStateDict) -> TossupPipelineState:
         return super().make_pipeline_state(state_dict)
 
+    def create_pipeline_state_dict(
+        self, workflow: TossupWorkflow, ui_state: PipelineUIState
+    ) -> td.TossupPipelineStateDict:
+        return super().create_pipeline_state_dict(workflow, ui_state)
+
     def update_workflow_from_code(
         self, yaml_str: str, change_state: bool
     ) -> tuple[td.TossupPipelineStateDict, bool, dict]:

src/components/model_pipeline/tossup_pipeline.py

@@ -9,7 +9,7 @@ from display.formatting import tiny_styled_warning
 from workflows.structs import Buzzer, TossupWorkflow
 
 from .model_pipeline import PipelineInterface
-from .state_manager import BasePipelineValidator
+from .state_manager import BasePipelineValidator, TossupPipelineStateManager
 
 
 def toggleable_slider(
@@ -33,6 +33,10 @@ def toggleable_slider(
 
 
 class TossupPipelineInterface(PipelineInterface):
+    """UI for the tossup pipeline."""
+
+    state_manager_cls = TossupPipelineStateManager
+
     def __init__(
         self,
         app: gr.Blocks,

src/components/quizbowl/validation.py

@@ -1,10 +1,10 @@
 from typing import Literal
 
-from app_configs import CONFIGS
+from app_configs import AVAILABLE_MODELS, CONFIGS
 from components.structs import PipelineState, TossupPipelineState
 from components.typed_dicts import PipelineStateDict, TossupPipelineStateDict
 from workflows.structs import TossupWorkflow, Workflow
-from workflows.validators import WorkflowValidator
+from workflows.validators import WorkflowValidationError, WorkflowValidator
 
 
 def validate_workflow(
@@ -83,3 +83,14 @@ class UserInputWorkflowValidator:
                     "\nDon't remove the keys from the 'outputs' field in the workflow. Only update their values."
                     f"\nMake sure you have values set for all the outputs: {default_str}"
                 )
+
+        # Validate the workflow
+        allowed_model_names = AVAILABLE_MODELS.keys()
+        self.validator = WorkflowValidator(allowed_model_names=allowed_model_names)
+        try:
+            self.validator.validate(workflow, allow_empty=True)
+        except WorkflowValidationError as e:
+            error_msg_total = f"Found {len(e.errors)} errors in the workflow:\n"
+            error_msg_list = [f"- {err.message}" for err in e.errors]
+            error_msg = error_msg_total + "\n".join(error_msg_list)
+            raise ValueError(error_msg)

src/workflows/validators.py

@@ -15,10 +15,14 @@ MAX_DESCRIPTION_LENGTH = 200
 MAX_SYSTEM_PROMPT_LENGTH = 4000
 MAX_TEMPERATURE = 10.0
 
+from loguru import logger
+
 
 class ValidationErrorType(Enum):
     """Types of validation errors that can occur"""
 
+    INPUTS = "inputs"
+    OUTPUTS = "outputs"
     STEP = "step"
     DAG = "dag"
     VARIABLE = "variable"
@@ -38,6 +42,17 @@ class ValidationError:
     step_id: Optional[str] = None
     field_name: Optional[str] = None
 
+    def __str__(self):
+        subject = ""
+        if self.step_id:
+            subject = f"Model step '{self.step_id}'"
+        if self.field_name:
+            if self.step_id:
+                subject = f"Field '{self.step_id}.{self.field_name}'"
+            else:
+                subject = f"Field '{self.field_name}'"
+        return f"{self.error_type.value}: {subject} - {self.message}"
+
 
 class WorkflowValidationError(ValueError):
     """Base class for workflow validation errors"""
@@ -83,6 +98,9 @@ class WorkflowValidator:
         max_field_name_length: int = MAX_FIELD_NAME_LENGTH,
         max_description_length: int = MAX_DESCRIPTION_LENGTH,
         max_system_prompt_length: int = MAX_SYSTEM_PROMPT_LENGTH,
+        allowed_model_names: Optional[list[str]] = None,
+        required_input_vars: Optional[list[str]] = None,
+        required_output_vars: Optional[list[str]] = None,
     ):
         self.errors: list[ValidationError] = []
         self.workflow: Optional[Workflow] = None
@@ -91,99 +109,82 @@ class WorkflowValidator:
         self.max_field_name_length = max_field_name_length
         self.max_description_length = max_description_length
         self.max_system_prompt_length = max_system_prompt_length
+        self.required_input_vars = required_input_vars
+        self.required_output_vars = required_output_vars
+        self.allowed_model_names = set(allowed_model_names) if allowed_model_names else None
+
+    def validate(self, workflow: Workflow, allow_empty: bool = False) -> bool:
+        validated = self._validate(workflow, allow_empty)
+        if not validated:
+            raise WorkflowValidationError(self.errors)
+        return True
 
-    def validate(self, workflow: Workflow) -> bool:
-        """Main validation entry point"""
+    def _validate(self, workflow: Workflow, allow_empty: bool = False) -> bool:
+        """Main validation entry point
+        Args:
+            workflow: The workflow to validate.
+            allow_empty: If True, empty workflow is allowed. This flag is used to validate the intermediate states while User edits the workflow.
+        """
         self.errors = []
         self.workflow = workflow
 
         # Basic workflow validation
-        if not self._validate_workflow_basic(workflow):
+        if not self._validate_workflow_basic(workflow, allow_empty):
             return False
 
         # If it's a single-step workflow, use simple validation
         if len(workflow.steps) == 1:
-            return self.validate_simple_workflow(workflow)
+            return self.validate_simple_workflow(workflow, allow_empty)
 
         # Otherwise use complex validation
-        return self.validate_complex_workflow(workflow)
+        return self.validate_complex_workflow(workflow, allow_empty)
 
-    def validate_simple_workflow(self, workflow: Workflow) -> bool:
-        """Validates a single-step workflow"""
-        if not self.workflow:
-            return False
+    def _validate_required_inputs(self, workflow: Workflow, allow_empty: bool = False) -> bool:
+        """Validates that the workflow has the correct inputs"""
 
-        # Get the single step
-        step = next(iter(workflow.steps.values()))
-
-        # Validate the step itself
-        if not self._validate_step(step):
+        required_input_vars = self.required_input_vars or []
+        input_vars = set(workflow.inputs)
+        for req_var in required_input_vars:
+            if req_var in input_vars:
+                continue
+            self.errors.append(
+                ValidationError(ValidationErrorType.INPUTS, f"Workflow must have '{req_var}' as an input")
+            )
             return False
 
-        # Validate input variables
-        for input_var in workflow.inputs:
+        for input_var in input_vars:
             if not self._is_valid_external_input(input_var):
                 self.errors.append(
                     ValidationError(ValidationErrorType.VARIABLE, f"Invalid input variable format: {input_var}")
                 )
                 return False
-
-        # Validate output variables
-        for output_name, output_var in workflow.outputs.items():
-            if not output_var:
-                self.errors.append(
-                    ValidationError(ValidationErrorType.VARIABLE, f"Missing output variable for {output_name}")
-                )
-                return False
-
-            # Check if output variable references a valid step output
-            if not self._is_valid_variable_reference(output_var):
-                self.errors.append(
-                    ValidationError(ValidationErrorType.VARIABLE, f"Invalid output variable reference: {output_var}")
-                )
-                return False
-
-            # Verify the output field exists in the step
-            _, field_name = _parse_variable_reference(output_var)
-            if not any(field.name == field_name for field in step.output_fields):
-                self.errors.append(
-                    ValidationError(
-                        ValidationErrorType.VARIABLE,
-                        f"Output field '{field_name}' not found in step '{step.id}'",
-                        step.id,
-                        field_name,
-                    )
-                )
-                return False
-
         return True
 
-    def validate_complex_workflow(self, workflow: Workflow) -> bool:
-        """Validates a multi-step workflow"""
-        if not self.workflow:
-            return False
-
-        # Validate each step
-        for step in workflow.steps.values():
-            if not self._validate_step(step):
-                return False
+    def _validate_required_outputs(self, workflow: Workflow, allow_empty: bool = False) -> bool:
+        """Validates that the workflow has the correct outputs"""
 
-        # Validate input variables
-        for input_var in workflow.inputs:
-            if not self._is_valid_external_input(input_var):
-                self.errors.append(
-                    ValidationError(ValidationErrorType.VARIABLE, f"Invalid input variable format: {input_var}")
-                )
-                return False
+        required_output_vars = self.required_output_vars or []
+        output_vars = set(workflow.outputs)
+        for req_var in required_output_vars:
+            if req_var in output_vars:
+                continue
+            self.errors.append(
+                ValidationError(ValidationErrorType.OUTPUTS, f"Workflow must produce '{req_var}' as an output")
+            )
+            return False
 
         # Validate output variables
         for output_name, output_var in workflow.outputs.items():
+            logger.debug(f"Output name: {output_name}, Output var: {output_var}")
             if not output_var:
+                if allow_empty:
+                    continue
                 self.errors.append(
                     ValidationError(ValidationErrorType.VARIABLE, f"Missing output variable for {output_name}")
                 )
                 return False
 
+            # Check if output variable references a valid step output
             if not self._is_valid_variable_reference(output_var):
                 self.errors.append(
                     ValidationError(ValidationErrorType.VARIABLE, f"Invalid output variable reference: {output_var}")
@@ -192,9 +193,10 @@ class WorkflowValidator:
 
             # Verify the output field exists in the referenced step
             step_id, field_name = _parse_variable_reference(output_var)
+            logger.debug(f"Step ID: {step_id}, Field name: {field_name}, Workflow steps: {workflow.steps.keys()}")
             if step_id not in workflow.steps:
                 self.errors.append(
-                    ValidationError(ValidationErrorType.VARIABLE, f"Referenced step '{step_id}' not found")
+                    ValidationError(ValidationErrorType.VARIABLE, f"Referenced model step '{step_id}' not found")
                 )
                 return False
 
@@ -203,12 +205,57 @@ class WorkflowValidator:
                 self.errors.append(
                     ValidationError(
                         ValidationErrorType.VARIABLE,
-                        f"Output field '{field_name}' not found in step '{step_id}'",
+                        f"Output field '{field_name}' not found in model step '{step_id}'",
                         step_id,
                         field_name,
                     )
                 )
                 return False
+        return True
+
+    def validate_input_outputs(self, workflow: Workflow, allow_empty: bool = False) -> bool:
+        """Validates the input and output variables"""
+
+        self._validate_required_inputs(workflow, allow_empty)
+        self._validate_required_outputs(workflow, allow_empty)
+
+        # Check for atleast one input
+        if not workflow.inputs:
+            self.errors.append(
+                ValidationError(ValidationErrorType.GENERAL, "Workflow must contain at least one input")
+            )
+
+        # Check for atleast one output
+        if not workflow.outputs:
+            self.errors.append(
+                ValidationError(ValidationErrorType.GENERAL, "Workflow must contain at least one output")
+            )
+
+        return len(self.errors) == 0
+
+    def validate_simple_workflow(self, workflow: Workflow, allow_empty: bool = False) -> bool:
+        """Validates a single-step workflow"""
+        if not self.workflow:
+            return False
+
+        # Get the single step
+        step = next(iter(workflow.steps.values()))
+
+        # Validate the step itself
+        if not self._validate_step(step, allow_empty):
+            return False
+
+        return True
+
+    def validate_complex_workflow(self, workflow: Workflow, allow_empty: bool = False) -> bool:
+        """Validates a multi-step workflow"""
+        if not self.workflow:
+            return False
+
+        # Validate each step
+        for step in workflow.steps.values():
+            if not self._validate_step(step, allow_empty):
+                return False
 
         dep_graph = create_step_dep_graph(workflow)
         if cycle_step_id := detect_cycles(dep_graph):
@@ -237,28 +284,17 @@ class WorkflowValidator:
 
         return True
 
-    def _validate_workflow_basic(self, workflow: Workflow) -> bool:
+    def _validate_workflow_basic(self, workflow: Workflow, allow_empty: bool = False) -> bool:
         """Validates basic workflow properties"""
-        # Check for atleast one input
-        if not workflow.inputs:
-            self.errors.append(
-                ValidationError(ValidationErrorType.GENERAL, "Workflow must contain at least one input")
-            )
-            return False
 
-        if not workflow.outputs:
-            self.errors.append(
-                ValidationError(ValidationErrorType.GENERAL, "Workflow must contain at least one output")
-            )
+        # Check the workflow inputs and outputs
+        if not self.validate_input_outputs(workflow, allow_empty):
             return False
 
-        for output_var in workflow.outputs.values():
-            if output_var is None:
-                self.errors.append(ValidationError(ValidationErrorType.GENERAL, "Output variable cannot be None"))
-                return False
-
         # Check for empty workflow
         if not workflow.steps:
+            if allow_empty:
+                return True
             self.errors.append(ValidationError(ValidationErrorType.GENERAL, "Workflow must contain at least one step"))
             return False
 
@@ -271,10 +307,26 @@ class WorkflowValidator:
                 return False
         return True
 
-    def _validate_step(self, step: ModelStep) -> bool:
+    def _validate_step(self, step: ModelStep, allow_empty: bool = False) -> bool:
         """Validates a single step"""
         # Validate required fields
-        if not step.id or not step.name or not step.model or not step.provider or not step.call_type:
+
+        model_name = step.get_full_model_name()
+
+        if model_name == "/" and not allow_empty:
+            self.errors.append(
+                ValidationError(ValidationErrorType.STEP, "Model name and provider cannot be empty", step.id)
+            )
+            return False
+
+        # Check if the model names are allowed
+        if self.allowed_model_names and model_name not in self.allowed_model_names:
+            self.errors.append(
+                ValidationError(ValidationErrorType.STEP, f"Model name '{model_name}' is not allowed", step.id)
+            )
+            return False
+
+        if not step.id or not step.call_type:
             self.errors.append(ValidationError(ValidationErrorType.STEP, "Step missing required fields", step.id))
             return False
 
@@ -328,7 +380,7 @@ class WorkflowValidator:
         # Validate input fields
         input_names = set()
         for field in step.input_fields:
-            if not self._validate_input_field(field):
+            if not self._validate_input_field(field, allow_empty):
                 return False
             if field.name in input_names:
                 self.errors.append(
@@ -342,7 +394,7 @@ class WorkflowValidator:
         # Validate output fields
         output_names = set()
         for field in step.output_fields:
-            if not self._validate_output_field(field):
+            if not self._validate_output_field(field, allow_empty):
                 return False
             if field.name in output_names:
                 self.errors.append(
@@ -355,7 +407,7 @@ class WorkflowValidator:
 
         return True
 
-    def _validate_input_field(self, field: InputField) -> bool:
+    def _validate_input_field(self, field: InputField, allow_empty: bool = False) -> bool:
         """Validates an input field"""
         # Validate required fields
         if not field.name or not field.description or not field.variable:
@@ -365,7 +417,7 @@ class WorkflowValidator:
             return False
 
         # Validate field name
-        if not self._is_valid_identifier(field.name):
+        if not self._is_valid_identifier(field.name, allow_empty):
             self.errors.append(
                 ValidationError(
                     ValidationErrorType.NAMING,
@@ -410,7 +462,7 @@ class WorkflowValidator:
 
         return True
 
-    def _validate_output_field(self, field: OutputField) -> bool:
+    def _validate_output_field(self, field: OutputField, allow_empty: bool = False) -> bool:
         """Validates an output field"""
         # Validate required fields
         if not field.name or not field.description:
@@ -422,7 +474,7 @@ class WorkflowValidator:
             return False
 
         # Validate field name
-        if not self._is_valid_identifier(field.name):
+        if not self._is_valid_identifier(field.name, allow_empty):
             self.errors.append(
                 ValidationError(
                     ValidationErrorType.NAMING,
@@ -528,10 +580,12 @@ class WorkflowValidator:
 
         return True
 
-    def _is_valid_variable_reference(self, var: str) -> bool:
+    def _is_valid_variable_reference(self, var: str | None, allow_empty: bool = True) -> bool:
         """Validates if a variable reference is properly formatted"""
         if not self.workflow:
             return False
+        if var is None:
+            return allow_empty
         parts = var.split(".")
         if len(parts) == 1:
             return True  # External input
@@ -554,8 +608,8 @@ class WorkflowValidator:
             return False
         return True
 
-    def _is_valid_identifier(self, name: str) -> bool:
+    def _is_valid_identifier(self, name: str, allow_empty: bool = False) -> bool:
         """Validates if a string is a valid Python identifier"""
         if name and name.strip():
             return bool(re.match(r"^[a-zA-Z_][a-zA-Z0-9_]*$", name))
-        return False
+        return allow_empty