tidy

Author: lingyielia

vizro-ai/src/vizro_ai/_vizro_ai.py

@@ -48,8 +48,6 @@ def __init__(self, model: Optional[Union[BaseChatModel, str]] = None):
 
         """
         self.model = _get_llm_model(model=model)
-        # self.components_instances = {} # TODO: This line is not used anywhere?
-
         self.chart_type_enum = create_chart_type_enum()
         self.chart_type_examples = [member.value for member in self.chart_type_enum]
         print(f"Available chart types: {self.chart_type_examples}")

vizro-ai/src/vizro_ai/plot/_response_models.py

@@ -231,8 +231,17 @@ def __new__(
         validate_code: bool = True,
         chart_type_examples: Optional[list[str]] = None,
     ) -> type[BaseChartPlan]:
-        """Creates a chart plan model with additional validation."""
-        # Use provided examples or fallback to defaults
+        """Creates a chart plan model with additional enrichments.
+
+        Args:
+            data_frame: DataFrame to use for validation.
+            chart_plan: Chart plan model to run extended validation against. Defaults to ChartPlan.
+            validate_code: Whether to validate the code.
+            chart_type_examples: List of chart type examples.
+
+        Returns:
+            Chart plan model with additional enrichments.
+        """
         examples = chart_type_examples or []
         description = f"""
         {chart_plan.model_fields["chart_type"].description}
@@ -242,15 +251,12 @@ def __new__(
         you may choose a different chart type if it better suits the user's needs.
         """
 
-        # Create field with examples but keep original description
         chart_type_field = Field(description=description, examples=examples)
 
-        # Set up validators
         validators = {}
         if validate_code:
             validators = {"validator1": field_validator("chart_code")(_test_execute_chart_code(data_frame))}
 
-        # Create and return the dynamic model
         return create_model(
             "ChartPlanDynamic",
             __base__=chart_plan,

vizro-ai/src/vizro_ai/utils/helper.py

@@ -23,20 +23,6 @@
 }
 
 
-# def _get_code_templates() -> dict[str, str]:
-#     """Get code templates using the vizro_visual_vocabulary package.
-
-#     Returns:
-#         Dictionary mapping chart types to their example code.
-#     """
-#     return vv_get_code_templates()
-
-
-# def create_chart_type_enum() -> type[Enum]:
-#     """Create a dynamic Enum from the available chart types using vizro_visual_vocabulary."""
-#     return vv_create_chart_type_enum()
-
-
 def _get_df_info(df: pd.DataFrame, n_sample: int = 5) -> tuple[str, str]:
     """Get the dataframe schema and head info as string."""
     formatted_pairs = [f"{col_name}: {dtype}" for col_name, dtype in df.dtypes.items()]