diff --git a/vizro-ai/src/vizro_ai/dashboard/response_models/components.py b/vizro-ai/src/vizro_ai/dashboard/response_models/components.py index 58afe2b56..2c246901f 100644 --- a/vizro-ai/src/vizro_ai/dashboard/response_models/components.py +++ b/vizro-ai/src/vizro_ai/dashboard/response_models/components.py @@ -1,7 +1,6 @@ """Component plan model.""" import logging -from typing import Union import vizro.models as vm from vizro.models.types import ComponentType @@ -51,11 +50,14 @@ def create(self, model, df_metadata) -> ComponentType: figure=vizro_ai.plot(df=df_metadata.get_df(self.df_name), 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.df_name)) + return vm.AgGrid( + id=self.component_id + "_" + self.page_id, figure=dash_ag_grid(data_frame=self.df_name) + ) elif self.component_type == "Card": return _get_pydantic_output(query=self.component_description, llm_model=model, result_model=vm.Card) except DebugFailure as e: logger.warning( f"Failed to build component: {self.component_id}.\n ------- \n " - f"Reason: {e} \n ------- \n Relevant prompt: `{self.component_description}`") + 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}") diff --git a/vizro-ai/src/vizro_ai/dashboard/response_models/controls.py b/vizro-ai/src/vizro_ai/dashboard/response_models/controls.py index 53e6668bd..e10dd312f 100644 --- a/vizro-ai/src/vizro_ai/dashboard/response_models/controls.py +++ b/vizro-ai/src/vizro_ai/dashboard/response_models/controls.py @@ -98,9 +98,11 @@ def create(self, model, available_components, df_metadata): ) except ValidationError as e: - logger.warning(f"Build failed for `Control`, returning default values. Try rephrase the prompt or " - f"select a different model. \n ------- \n Error details: {e} \n ------- \n " - f"Relevant prompt: `{self.control_description}`") + logger.warning( + f"Build failed for `Control`, returning default values. Try rephrase the prompt or " + f"select a different model. \n ------- \n Error details: {e} \n ------- \n " + f"Relevant prompt: `{self.control_description}`" + ) return None return actual diff --git a/vizro-ai/src/vizro_ai/dashboard/response_models/df_info.py b/vizro-ai/src/vizro_ai/dashboard/response_models/df_info.py index 54a54a52c..e47b8b7da 100644 --- a/vizro-ai/src/vizro_ai/dashboard/response_models/df_info.py +++ b/vizro-ai/src/vizro_ai/dashboard/response_models/df_info.py @@ -1,9 +1,8 @@ """Data Summary Node.""" -from typing import Any, Dict, Tuple, List +from typing import Any, Dict, Tuple import pandas as pd -from langchain_core.language_models.chat_models import BaseChatModel try: from pydantic.v1 import BaseModel, Field diff --git a/vizro-ai/src/vizro_ai/dashboard/response_models/layout.py b/vizro-ai/src/vizro_ai/dashboard/response_models/layout.py index ae5a89a6d..fa362f0d5 100644 --- a/vizro-ai/src/vizro_ai/dashboard/response_models/layout.py +++ b/vizro-ai/src/vizro_ai/dashboard/response_models/layout.py @@ -8,7 +8,7 @@ try: from pydantic.v1 import BaseModel, Field, ValidationError, validator except ImportError: # pragma: no cov - from pydantic import BaseModel, Field, ValidationError, validator + from pydantic import BaseModel, Field, validator import numpy as np from vizro.models._layout import _get_grid_lines, _get_unique_grid_component_ids, _validate_grid_areas from vizro_ai.dashboard._pydantic_output import _get_pydantic_output @@ -71,9 +71,11 @@ def create(self, model) -> Union[vm.Layout, None]: proxy = _get_pydantic_output(query=layout_prompt, llm_model=model, result_model=LayoutProxyModel) actual = vm.Layout.parse_obj(proxy.dict(exclude={})) except DebugFailure as e: - logger.warning(f"Build failed for `Layout`, returning default values. Try rephrase the prompt or " - f"select a different model. \n ------- \n Error details: {e} \n ------- \n " - f"Relevant prompt: `{self.layout_description}`") + logger.warning( + f"Build failed for `Layout`, returning default values. Try rephrase the prompt or " + f"select a different model. \n ------- \n Error details: {e} \n ------- \n " + f"Relevant prompt: `{self.layout_description}`" + ) actual = None return actual diff --git a/vizro-ai/src/vizro_ai/dashboard/response_models/page.py b/vizro-ai/src/vizro_ai/dashboard/response_models/page.py index 16400fb04..b2fed6918 100644 --- a/vizro-ai/src/vizro_ai/dashboard/response_models/page.py +++ b/vizro-ai/src/vizro_ai/dashboard/response_models/page.py @@ -4,16 +4,15 @@ from typing import List, Union try: - from pydantic.v1 import BaseModel, Field, validator, PrivateAttr + from pydantic.v1 import BaseModel, Field, PrivateAttr, validator except ImportError: # pragma: no cov - from pydantic import BaseModel, Field, validator, PrivateAttr + from pydantic import BaseModel, Field, PrivateAttr, validator +import vizro.models as vm +from tqdm.auto import tqdm from vizro_ai.dashboard.response_models.components import ComponentPlan from vizro_ai.dashboard.response_models.controls import ControlPlan from vizro_ai.dashboard.response_models.layout import LayoutPlan -import vizro.models as vm -from tqdm.auto import tqdm from vizro_ai.dashboard.utils import _execute_step -from vizro_ai.utils.helper import DebugFailure logger = logging.getLogger(__name__) @@ -41,13 +40,13 @@ def _check_components_plan(cls, v): if len(v) == 0: raise ValueError("A page must contain at least one component.") return v - + def __init__(self, **data): super().__init__(**data) self._components = None self._controls = None self._layout = None - + def _get_components(self, df_metadata, model): if self._components is None: self._components = self._build_components(df_metadata, model) @@ -62,14 +61,12 @@ def _build_components(self, df_metadata, model): leave=False, ) as pbar: for component_plan in self.components_plan: - component_log.set_description_str( - f"[Page] <{self.title}>: [Component] {component_plan.component_id}" - ) + component_log.set_description_str(f"[Page] <{self.title}>: [Component] {component_plan.component_id}") pbar.update(1) components.append(component_plan.create(df_metadata=df_metadata, model=model)) component_log.close() return components - + def _get_layout(self, model): if self._layout is None: self._layout = self._build_layout(model) @@ -79,14 +76,18 @@ def _build_layout(self, model): if self.layout_plan is None: return None return self.layout_plan.create(model=model) - + def _get_controls(self, df_metadata, model): if self._controls is None: self._controls = self._build_controls(df_metadata, model) return self._controls def _available_components(self, df_metadata, model): - return [comp.id for comp in self._get_components(df_metadata=df_metadata, model=model) if isinstance(comp, (vm.Graph, vm.AgGrid))] + return [ + comp.id + for comp in self._get_components(df_metadata=df_metadata, model=model) + if isinstance(comp, (vm.Graph, vm.AgGrid)) + ] def _build_controls(self, df_metadata, model): controls = [] @@ -98,25 +99,27 @@ def _build_controls(self, df_metadata, model): for control_plan in self.controls_plan: pbar.update(1) control = control_plan.create( - model=model, available_components=self._available_components(df_metadata, model), df_metadata=df_metadata + model=model, + available_components=self._available_components(df_metadata, model), + df_metadata=df_metadata, ) if control: controls.append(control) return controls - - def create(self, model, df_metadata): page_desc = f"Building page: {self.title}" logger.info(page_desc) pbar = tqdm(total=5, desc=page_desc) title = _execute_step(pbar, page_desc + " --> add title", self.title) - components = _execute_step(pbar, page_desc + " --> add components", self._get_components(df_metadata=df_metadata, model=model)) + components = _execute_step( + pbar, page_desc + " --> add components", self._get_components(df_metadata=df_metadata, model=model) + ) controls = _execute_step(pbar, page_desc + " --> add controls", self._get_controls(df_metadata, model)) layout = _execute_step(pbar, page_desc + " --> add layout", self._get_layout(model)) - + page = vm.Page(title=title, components=components, controls=controls, layout=layout) _execute_step(pbar, page_desc + " --> done", None) pbar.close()