lint

vizro-ai/src/vizro_ai/dashboard/graph/code_generation.py

@@ -2,8 +2,7 @@
 
 import logging
 import operator
-import re
-from typing import Annotated, Any, Dict, List, Union
+from typing import Annotated, Dict, List
 
 import pandas as pd
 import vizro.models as vm
@@ -30,7 +29,6 @@
 logger.setLevel(logging.INFO)
 
 
-
 Messages = List[BaseMessage]
 """List of messages."""
 

vizro-ai/src/vizro_ai/dashboard/nodes/_model.py

@@ -5,7 +5,6 @@
     from pydantic.v1 import BaseModel, ValidationError
 except ImportError:  # pragma: no cov
     from pydantic import BaseModel, ValidationError
-from typing import Dict
 
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import HumanMessage

vizro-ai/src/vizro_ai/dashboard/nodes/build.py

@@ -3,11 +3,11 @@
 import logging
 
 import vizro.models as vm
+from tqdm.auto import tqdm, trange
 from vizro_ai.dashboard.utils import _execute_step
 from vizro_ai.utils.helper import DebugFailure
 
 logger = logging.getLogger(__name__)
-from tqdm.auto import tqdm, trange
 
 
 class PageBuilder:

vizro-ai/src/vizro_ai/dashboard/nodes/data_summary.py

@@ -38,6 +38,4 @@ def _get_df_info(df: pd.DataFrame) -> Tuple[Dict[str, str], str]:
 class DfInfo(BaseModel):
     """Data Info output."""
 
-    dataset_name: str = Field(
-        pattern=r"^[a-z]+(_[a-z]+)?$", description="Small snake case name of the dataset."
-    )
+    dataset_name: str = Field(pattern=r"^[a-z]+(_[a-z]+)?$", description="Small snake case name of the dataset.")

vizro-ai/src/vizro_ai/dashboard/nodes/plan.py

@@ -1,7 +1,7 @@
 """Module containing the planner functionality."""
 
 import logging
-from typing import Dict, List, Literal, Union
+from typing import List, Literal, Union
 
 import vizro.models as vm
 from langchain_openai import ChatOpenAI
@@ -45,10 +45,7 @@ class Component(BaseModel):
     component_id: str = Field(
         pattern=r"^[a-z]+(_[a-z]+)?$", description="Small snake case description of this component."
     )
-    page_id: str = Field(
-        ...,
-        description="The page id where this component will be placed."
-    )
+    page_id: str = Field(..., description="The page id where this component will be placed.")
     data_frame: str = Field(
         ...,
         description="The name of the dataframe that this component will use. If the dataframe is "
@@ -63,13 +60,13 @@ def create(self, model, df_metadata) -> Union[ComponentType, None]:
 
         if self.component_type == "Graph":
             return vm.Graph(
-                id=self.component_id+"_"+self.page_id,
-                figure=vizro_ai.plot(df=df_metadata.metadata[self.data_frame].df, user_input=self.component_description)
+                id=self.component_id + "_" + self.page_id,
+                figure=vizro_ai.plot(
+                    df=df_metadata.metadata[self.data_frame].df, user_input=self.component_description
+                ),
             )
         elif self.component_type == "AgGrid":
-            return vm.AgGrid(
-                id=self.component_id+"_"+self.page_id,
-                figure=dash_ag_grid(data_frame=self.data_frame))
+            return vm.AgGrid(id=self.component_id + "_" + self.page_id, figure=dash_ag_grid(data_frame=self.data_frame))
         elif self.component_type == "Card":
             return _get_proxy_model(
                 query=self.component_description, llm_model=model, result_model=vm.Card, df_metadata=df_metadata

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

@@ -1,12 +1,11 @@
 """Helper Functions For Vizro AI dashboard."""
 
-from typing import Any, Dict
 from dataclasses import dataclass, field
 
-import pandas as pd
-
 # import black
-from typing import Any
+from typing import Any, Dict
+
+import pandas as pd
 import tqdm.std
 import vizro.models as vm
 
@@ -18,19 +17,22 @@
     "import vizro.models as vm\n"
 )
 
+
 @dataclass
 class DataFrameMetadata:
     """Dataclass containing metadata for a dataframe."""
 
     df_schema: Dict[str, str]
     df: pd.DataFrame
 
+
 @dataclass
 class DfMetadata:
     """Dataclass containing metadata for all dataframes."""
 
     metadata: Dict[str, DataFrameMetadata] = field(default_factory=dict)
 
+
 @dataclass
 class DashboardOutputs:
     """Dataclass containing all possible `VizroAI.dashboard()` output."""