add code sample and remove unit tests

vizro-ai/src/vizro_ai/dashboard/_pydantic_output.py

@@ -79,3 +79,13 @@ def _get_pydantic_output(
         except ValidationError as validation_error:
             last_validation_error = validation_error
     raise last_validation_error
+
+
+if __name__ == "__main__":
+    import vizro.models as vm
+    from vizro_ai.chains._llm_models import _get_llm_model
+
+    model = _get_llm_model()
+    component_description = "Create a card with the following content: 'Hello, world!'"
+    res = _get_pydantic_output(query=component_description, llm_model=model, result_model=vm.Card)
+    print(res)  # noqa: T201

vizro-ai/src/vizro_ai/dashboard/response_models/components.py

@@ -61,3 +61,21 @@ def create(self, model, df_metadata) -> ComponentType:
                 f"Reason: {e} \n ------- \n Relevant prompt: `{self.component_description}`"
             )
             return vm.Card(id=self.component_id, text=f"Failed to build component: {self.component_id}")
+
+
+if __name__ == "__main__":
+    from vizro_ai.chains._llm_models import _get_llm_model
+    from vizro_ai.dashboard.utils import DfMetadata
+
+    model = _get_llm_model()
+
+    df_metadata = DfMetadata({})
+    component_plan = ComponentPlan(
+        component_type="Card",
+        component_description="Create a card says 'this is worldwide GDP'.",
+        component_id="gdp_card",
+        page_id="1",
+        df_name="N/A",
+    )
+    component = component_plan.create(model, df_metadata)
+    print(component)  # noqa: T201

vizro-ai/src/vizro_ai/dashboard/response_models/controls.py

@@ -80,7 +80,7 @@ def create(self, model, available_components, df_metadata):
             _df_cols = list(_df_schema.keys())
         # when wrong dataframe name is given
         except KeyError:
-            logger.info(f"Dataframe {self.df_name} not found in metadata, returning default values.")
+            logger.warning(f"Dataframe {self.df_name} not found in metadata, returning default values.")
             return None
 
         try:
@@ -106,3 +106,24 @@ def create(self, model, available_components, df_metadata):
             return None
 
         return actual
+
+
+if __name__ == "__main__":
+    import pandas as pd
+    from vizro_ai.chains._llm_models import _get_llm_model
+    from vizro_ai.dashboard.utils import DfMetadata, MetadataContent
+
+    model = _get_llm_model()
+
+    df_metadata = DfMetadata({})
+    df_metadata.metadata["gdp_chart"] = MetadataContent(
+        df_schema={"a": "int64", "b": "int64"},
+        df=pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [4, 5, 6, 7, 8]}),
+        df_sample=pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [4, 5, 6, 7, 8]}),
+    )
+    control_plan = ControlPlan(
+        control_type="Filter",
+        control_description="Create a filter that filters the data based on the column 'a'.",
+        df_name="gdp_chart",
+    )
+    control = control_plan.create(model, ["gdp_chart"], df_metadata)

vizro-ai/src/vizro_ai/dashboard/response_models/df_info.py

@@ -36,3 +36,11 @@ def _get_df_info(df: pd.DataFrame) -> Tuple[Dict[str, str], pd.DataFrame]:
 def _create_df_info_content(df_schema: Any, df_sample: Any, current_df_names: list) -> dict:
     """Create the message content for the dataframe summarization."""
     return DF_SUM_PROMPT.format(df_sample=df_sample, df_schema=df_schema, current_df_names=current_df_names)
+
+
+if __name__ == "__main__":
+    df = pd.DataFrame({"a": [1, 2, 3, 4, 5], "b": [4, 5, 6, 7, 8]})
+    df_schema, df_sample = _get_df_info(df)
+    current_df_names = ["df1", "df2"]
+    print(_create_df_info_content(df_schema, df_sample, current_df_names))  # noqa: T201
+    print(DfInfo(dataset_name="test").dict())  # noqa: T201

vizro-ai/src/vizro_ai/dashboard/response_models/layout.py

@@ -79,3 +79,16 @@ def create(self, model) -> Union[vm.Layout, None]:
             actual = None
 
         return actual
+
+
+if __name__ == "__main__":
+    from vizro_ai.chains._llm_models import _get_llm_model
+
+    model = _get_llm_model()
+    layout_plan = LayoutPlan(
+        layout_description="Create a layout with a graph on the left and a card on the right.",
+        layout_grid_template_areas=["graph card"],
+    )
+    layout = layout_plan.create(model)
+    print(layout)  # noqa: T201
+    print(layout.dict())  # noqa: T201