From 6b678a2d7f21328035b1ea19e7bcd737cd25b02a Mon Sep 17 00:00:00 2001 From: Lingyi Zhang Date: Wed, 17 Jul 2024 23:13:24 -0400 Subject: [PATCH] refactor page build --- .../vizro_ai/dashboard/_pydantic_output.py | 2 +- .../dashboard/graph/dashboard_creation.py | 34 +++--- .../dashboard/response_models/components.py | 27 +++-- .../dashboard/response_models/controls.py | 4 +- .../dashboard/response_models/df_info.py | 1 - .../dashboard/response_models/layout.py | 7 +- .../dashboard/response_models/page.py | 95 ++++++++++++++- .../dashboard/response_models/page_build.py | 112 ------------------ 8 files changed, 135 insertions(+), 147 deletions(-) delete mode 100644 vizro-ai/src/vizro_ai/dashboard/response_models/page_build.py diff --git a/vizro-ai/src/vizro_ai/dashboard/_pydantic_output.py b/vizro-ai/src/vizro_ai/dashboard/_pydantic_output.py index 33bd3e4e2..106fa2b7c 100644 --- a/vizro-ai/src/vizro_ai/dashboard/_pydantic_output.py +++ b/vizro-ai/src/vizro_ai/dashboard/_pydantic_output.py @@ -78,4 +78,4 @@ def _get_pydantic_output( return res except ValidationError as validation_error: last_validation_error = validation_error - return last_validation_error + raise last_validation_error diff --git a/vizro-ai/src/vizro_ai/dashboard/graph/dashboard_creation.py b/vizro-ai/src/vizro_ai/dashboard/graph/dashboard_creation.py index b11ed94c4..c6c5536da 100644 --- a/vizro-ai/src/vizro_ai/dashboard/graph/dashboard_creation.py +++ b/vizro-ai/src/vizro_ai/dashboard/graph/dashboard_creation.py @@ -15,8 +15,8 @@ from vizro_ai.dashboard.response_models.dashboard import DashboardPlanner from vizro_ai.dashboard.response_models.df_info import DfInfo, _create_df_info_content, _get_df_info from vizro_ai.dashboard.response_models.page import PagePlanner -from vizro_ai.dashboard.response_models.page_build import PageBuilder from vizro_ai.dashboard.utils import DfMetadata, MetadataContent, _execute_step +from vizro_ai.utils.helper import DebugFailure try: from pydantic.v1 import BaseModel, validator @@ -81,12 +81,16 @@ def _store_df_info(state: GraphState, config: RunnableConfig) -> Dict[str, DfMet ) llm = config["configurable"].get("model", None) - df_name = _get_pydantic_output( - query=query, - llm_model=llm, - result_model=DfInfo, - df_info=df_info, - ).dataset_name + try: + df_name = _get_pydantic_output( + query=query, + llm_model=llm, + result_model=DfInfo, + df_info=df_info, + ).dataset_name + except DebugFailure as e: + logger.warning(f"Failed in name generation {e}") + df_name = f"df_{len(current_df_names)}" current_df_names.append(df_name) @@ -111,9 +115,13 @@ def _dashboard_plan(state: GraphState, config: RunnableConfig) -> Dict[str, Dash node_desc + " --> in progress \n(this step could take longer " "when more complex requirements are given)", None, ) - dashboard_plan = _get_pydantic_output( - query=query, llm_model=llm, result_model=DashboardPlanner, df_info=df_metadata.get_schemas_and_samples() - ) + try: + dashboard_plan = _get_pydantic_output( + query=query, llm_model=llm, result_model=DashboardPlanner, df_info=df_metadata.get_schemas_and_samples() + ) + except DebugFailure as e: + logger.error(f"Error in dashboard plan generation: {e}", exc_info=True) + raise _execute_step(pbar, node_desc + " --> done", None) pbar.close() @@ -139,11 +147,7 @@ def _build_page(state: BuildPageState, config: RunnableConfig) -> Dict[str, List page_plan = state["page_plan"] llm = config["configurable"].get("model", None) - page = PageBuilder( - model=llm, - df_metadata=df_metadata, - page_plan=page_plan, - ).page + page = page_plan.create(model=llm, df_metadata=df_metadata) return {"pages": [page]} 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 606e7644f..58afe2b56 100644 --- a/vizro-ai/src/vizro_ai/dashboard/response_models/components.py +++ b/vizro-ai/src/vizro_ai/dashboard/response_models/components.py @@ -13,6 +13,7 @@ from vizro.tables import dash_ag_grid from vizro_ai.dashboard._constants import component_type from vizro_ai.dashboard._pydantic_output import _get_pydantic_output +from vizro_ai.utils.helper import DebugFailure logger = logging.getLogger(__name__) @@ -37,18 +38,24 @@ class ComponentPlan(BaseModel): "used, please specify that as N/A.", ) - def create(self, model, df_metadata) -> Union[ComponentType, None]: + def create(self, model, df_metadata) -> ComponentType: """Create the component.""" from vizro_ai import VizroAI vizro_ai = VizroAI(model=model) - if self.component_type == "Graph": - return vm.Graph( - id=self.component_id + "_" + self.page_id, - 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)) - elif self.component_type == "Card": - return _get_pydantic_output(query=self.component_description, llm_model=model, result_model=vm.Card) + try: + if self.component_type == "Graph": + return vm.Graph( + id=self.component_id + "_" + self.page_id, + 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)) + 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}`") + 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 4b71e36ea..53e6668bd 100644 --- a/vizro-ai/src/vizro_ai/dashboard/response_models/controls.py +++ b/vizro-ai/src/vizro_ai/dashboard/response_models/controls.py @@ -98,7 +98,9 @@ def create(self, model, available_components, df_metadata): ) except ValidationError as e: - logger.info(f"Build failed for `Control`, returning default values. Error details: {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}`") 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 bb0c71d5b..54a54a52c 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 @@ -10,7 +10,6 @@ except ImportError: # pragma: no cov from pydantic import BaseModel, Field -from vizro_ai.dashboard._pydantic_output import _get_pydantic_output DF_SUM_PROMPT = """ Inspect the provided data and give a short unique name to the dataset. \n 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 4abad9328..ae5a89a6d 100644 --- a/vizro-ai/src/vizro_ai/dashboard/response_models/layout.py +++ b/vizro-ai/src/vizro_ai/dashboard/response_models/layout.py @@ -12,6 +12,7 @@ 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 +from vizro_ai.utils.helper import DebugFailure logger = logging.getLogger(__name__) @@ -69,8 +70,10 @@ def create(self, model) -> Union[vm.Layout, None]: try: proxy = _get_pydantic_output(query=layout_prompt, llm_model=model, result_model=LayoutProxyModel) actual = vm.Layout.parse_obj(proxy.dict(exclude={})) - except (ValidationError, AttributeError) as e: - logger.info(f"Build failed for `Layout`, returning default values. Error details: {e}") + 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}`") 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 4ca51e8e3..16400fb04 100644 --- a/vizro-ai/src/vizro_ai/dashboard/response_models/page.py +++ b/vizro-ai/src/vizro_ai/dashboard/response_models/page.py @@ -1,15 +1,19 @@ """Page plan model.""" import logging -from typing import List +from typing import List, Union try: - from pydantic.v1 import BaseModel, Field, validator + from pydantic.v1 import BaseModel, Field, validator, PrivateAttr except ImportError: # pragma: no cov - from pydantic import BaseModel, Field, validator + from pydantic import BaseModel, Field, validator, PrivateAttr 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__) @@ -28,11 +32,92 @@ class PagePlanner(BaseModel): controls_plan: List[ControlPlan] = Field([], description="Controls of the page.") layout_plan: LayoutPlan = Field(None, description="Layout of the page.") + _components: List[Union[vm.Card, vm.AgGrid, vm.Figure]] = PrivateAttr() + _controls: List[vm.Filter] = PrivateAttr() + _layout: vm.Layout = PrivateAttr() + @validator("components_plan") 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) + return self._components + + def _build_components(self, df_metadata, model): + components = [] + component_log = tqdm(total=0, bar_format="{desc}", leave=False) + with tqdm( + total=len(self.components_plan), + desc=f"Currently Building ... [Page] <{self.title}> components", + 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}" + ) + 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) + return self._layout + + 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))] + + def _build_controls(self, df_metadata, model): + controls = [] + with tqdm( + total=len(self.controls_plan), + desc=f"Currently Building ... [Page] <{self.title}> controls", + leave=False, + ) as pbar: + 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 + ) + 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) - def create(): - pass + 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)) + 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() + return page diff --git a/vizro-ai/src/vizro_ai/dashboard/response_models/page_build.py b/vizro-ai/src/vizro_ai/dashboard/response_models/page_build.py deleted file mode 100644 index 73077bae6..000000000 --- a/vizro-ai/src/vizro_ai/dashboard/response_models/page_build.py +++ /dev/null @@ -1,112 +0,0 @@ -"""Module that contains the builder functionality.""" - -import logging - -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__) - - -class PageBuilder: - """Class to build a page.""" - - def __init__(self, model, df_metadata, page_plan): - """Initialize PageBuilder.""" - self._model = model - self._df_metadata = df_metadata - self._page_plan = page_plan - self._components = None - self._controls = None - self._page = None - self._layout = None - - @property - def components(self): - """Property to get components.""" - if self._components is None: - self._components = self._build_components() - return self._components - - def _build_components(self): - components = [] - component_log = tqdm(total=0, bar_format="{desc}", leave=False) - with tqdm( - total=len(self._page_plan.components_plan), - desc=f"Currently Building ... [Page] <{self._page_plan.title}> components", - leave=False, - ) as pbar: - for component_plan in self._page_plan.components_plan: - component_log.set_description_str( - f"[Page] <{self._page_plan.title}>: [Component] {component_plan.component_id}" - ) - pbar.update(1) - try: - components.append(component_plan.create(df_metadata=self._df_metadata, model=self._model)) - except DebugFailure as e: - components.append(vm.Card(id=component_plan.component_id, text=f"Failed to build component: {e}")) - component_log.close() - return components - - @property - def layout(self): - """Property to get layout.""" - if self._layout is None: - self._layout = self._build_layout() - return self._layout - - def _build_layout(self): - if self._page_plan.layout_plan is None: - return None - return self._page_plan.layout_plan.create(model=self._model) - - @property - def controls(self): - """Property to get controls.""" - if self._controls is None: - self._controls = self._build_controls() - return self._controls - - @property - def available_components(self): - """Property to get available components.""" - return [comp.id for comp in self.components if isinstance(comp, (vm.Graph, vm.AgGrid))] - - def _build_controls(self): - controls = [] - with tqdm( - total=len(self._page_plan.controls_plan), - desc=f"Currently Building ... [Page] <{self._page_plan.title}> controls", - leave=False, - ) as pbar: - for control_plan in self._page_plan.controls_plan: - pbar.update(1) - control = control_plan.create( - model=self._model, available_components=self.available_components, df_metadata=self._df_metadata - ) - if control: - controls.append(control) - - return controls - - @property - def page(self): - """Property to get page.""" - if self._page is None: - page_desc = f"Building page: {self._page_plan.title}" - logger.info(page_desc) - pbar = tqdm(total=5, desc=page_desc) - - title = _execute_step(pbar, page_desc + " --> add title", self._page_plan.title) - components = _execute_step(pbar, page_desc + " --> add components", self.components) - controls = _execute_step(pbar, page_desc + " --> add controls", self.controls) - layout = _execute_step(pbar, page_desc + " --> add layout", self.layout) - - # TODO: handle error `ValidationError: 1 validation error for Page` - # Number of page and grid components need to be the same. - self._page = vm.Page(title=title, components=components, controls=controls, layout=layout) - _execute_step(pbar, page_desc + " --> done", None) - pbar.close() - return self._page