update progress bar

vizro-ai/pyproject.toml

@@ -19,7 +19,7 @@ dependencies = [
   "openai>=1.0.0",
   "langchain>=0.1.0, <0.3.0",  # TODO update all LLMChain class, update to pydantic v2 and remove upper bound
   "langchain-openai",
-  "langgraph",
+  "langgraph>=0.1.2",
   "python-dotenv>=1.0.0",  # TODO decide env var management to see if we need this
   "vizro>=0.1.4",  # TODO set upper bound later
   "ipython>=8.10.0",  # not directly required, pinned by Snyk to avoid a vulnerability: https://app.snyk.io/vuln/SNYK-PYTHON-IPYTHON-3318382

vizro-ai/src/vizro_ai/chains/_llm_models.py

@@ -28,6 +28,7 @@
     "Anthropic": [
         "claude-3-haiku-20240307",
         "claude-3-sonnet-20240229",
+        "claude-3-opus-20240229",
     ],
 }
 

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

@@ -10,14 +10,15 @@
 from langchain_core.runnables import RunnableConfig
 from langgraph.constants import END, Send
 from langgraph.graph import StateGraph
+from tqdm.auto import tqdm
 from vizro_ai.dashboard.nodes.build import PageBuilder
 from vizro_ai.dashboard.nodes.data_summary import DfInfo, _get_df_info, df_sum_prompt
 from vizro_ai.dashboard.nodes.plan import (
     DashboardPlanner,
     PagePlanner,
     _get_dashboard_plan,
 )
-from vizro_ai.dashboard.utils import DataFrameMetadata, DfMetadata
+from vizro_ai.dashboard.utils import DataFrameMetadata, DfMetadata, _execute_step
 
 try:
     from pydantic.v1 import BaseModel, validator
@@ -26,7 +27,6 @@
 
 
 logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
 
 
 Messages = List[BaseMessage]
@@ -71,37 +71,48 @@ def check_dataframes(cls, v):
 
 def _store_df_info(state: GraphState, config: RunnableConfig) -> Dict[str, DfMetadata]:
     """Store information about the dataframes."""
-    logger.info("*** _store_df_info ***")
     dfs = state.dfs
     df_metadata = state.df_metadata
     messages = state.messages
     current_df_names = []
-    for df in dfs:
-        df_schema, df_sample = _get_df_info(df)
+    with tqdm(total=len(dfs), desc="Store df info") as pbar:
+        for df in dfs:
+            df_schema, df_sample = _get_df_info(df)
 
-        llm = config["configurable"].get("model", None)
-        data_sum_chain = df_sum_prompt | llm.with_structured_output(DfInfo)
+            llm = config["configurable"].get("model", None)
+            data_sum_chain = df_sum_prompt | llm.with_structured_output(DfInfo)
 
-        df_name = data_sum_chain.invoke(
-            {"messages": messages, "df_schema": df_schema, "df_sample": df_sample, "current_df_names": current_df_names}
-        ).dataset_name
+            df_name = data_sum_chain.invoke(
+                {
+                    "messages": messages,
+                    "df_schema": df_schema,
+                    "df_sample": df_sample,
+                    "current_df_names": current_df_names,
+                }
+            ).dataset_name
 
-        current_df_names.append(df_name)
+            current_df_names.append(df_name)
 
-        logger.info(f"df_name: {df_name}")
-        df_metadata.metadata[df_name] = DataFrameMetadata(df_schema=df_schema, df=df, df_sample=df_sample)
+            pbar.write(f"df_name: {df_name}")
+            pbar.update(1)
+            df_metadata.metadata[df_name] = DataFrameMetadata(df_schema=df_schema, df=df, df_sample=df_sample)
 
     return {"df_metadata": df_metadata}
 
 
 def _dashboard_plan(state: GraphState, config: RunnableConfig) -> Dict[str, DashboardPlanner]:
     """Generate a dashboard plan."""
-    logger.info("*** _dashboard_plan ***")
+    node_desc = "Generate dashboard plan"
+    pbar = tqdm(total=2, desc=node_desc)
     query = state.messages[0].content
     df_metadata = state.df_metadata
 
     llm = config["configurable"].get("model", None)
+
+    _execute_step(pbar, node_desc + " --> in progress", None)
     dashboard_plan = _get_dashboard_plan(query=query, model=llm, df_metadata=df_metadata)
+    _execute_step(pbar, node_desc + " --> done", None)
+    pbar.close()
 
     return {"dashboard_plan": dashboard_plan}
 
@@ -136,7 +147,6 @@ def _build_page(state: BuildPageState, config: RunnableConfig) -> Dict[str, List
 
 def _continue_to_pages(state: GraphState) -> List[Send]:
     """Continue to build pages."""
-    logger.info("*** build_page ***")
     df_metadata = state.df_metadata
     return [
         Send(node="_build_page", arg={"page_plan": v, "df_metadata": df_metadata}) for v in state.dashboard_plan.pages
@@ -145,7 +155,6 @@ def _continue_to_pages(state: GraphState) -> List[Send]:
 
 def _build_dashboard(state: GraphState) -> Dict[str, vm.Dashboard]:
     """Build a dashboard."""
-    logger.info("*** build_dashboard ***")
     dashboard_plan = state.dashboard_plan
     pages = state.pages
 

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

@@ -26,7 +26,7 @@ def _get_df_info(df: pd.DataFrame) -> Tuple[Dict[str, str], str]:
             Inspect the provided data and give a short unique name to the dataset. \n
             Here is the dataframe sample:  \n ------- \n  {df_sample} \n ------- \n
             Here is the schema:  \n ------- \n  {df_schema} \n ------- \n
-            Avoid the following names currently in use: \n ------- \n {current_df_names} \n ------- \n
+            AVOID the following names: \n ------- \n {current_df_names} \n ------- \n
             \n ------- \n
             """,
         ),

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

@@ -58,8 +58,8 @@ class DashboardOutputs:
 
 
 def _execute_step(pbar: tqdm.std.tqdm, description: str, value: Any) -> Any:
-    pbar.set_description_str(description)
     pbar.update(1)
+    pbar.set_description_str(description)
     return value