Implement Vizro dashboard generation POC

mito-ai/mito_ai/streamlit_conversion/prompts/streamlit_app_creation_prompt.py

@@ -6,14 +6,14 @@
 
 def get_streamlit_app_creation_prompt(notebook: List[dict]) -> str:
     """
-    This prompt is used to create a streamlit app from a notebook.
+    This prompt is used to create a vizro app from a notebook.
     """
-    return f"""Convert the following Jupyter notebook into a Streamlit application.
+    return f"""Convert the following Jupyter notebook into a Vizro dashboard application.
 
-GOAL: Create a complete, runnable Streamlit app that accurately represents the notebook. It must completely convert the notebook. 
+GOAL: Create a complete, runnable Vizro app that accurately represents the notebook. It must completely convert the notebook.
 
 TODO PLACEHOLDER RULES:
-If you decide to leave any TODOs, you must mark them with {MITO_TODO_PLACEHOLDER}. You should use {MITO_TODO_PLACEHOLDER} instead of comments like the following: 
+If you decide to leave any TODOs, you must mark them with {MITO_TODO_PLACEHOLDER}. You should use {MITO_TODO_PLACEHOLDER} instead of comments like the following:
 - # ... (include all mappings from the notebook)
 - # ... (include all violation codes from the notebook)
 - # Fill in the rest of the code here
@@ -43,4 +43,4 @@ def get_streamlit_app_creation_prompt(notebook: List[dict]) -> str:
 Notebook to convert:
 
 {notebook}
-"""
+"""

mito-ai/mito_ai/streamlit_conversion/prompts/streamlit_error_correction_prompt.py

@@ -5,18 +5,18 @@
 from mito_ai.streamlit_conversion.prompts.prompt_utils import add_line_numbers_to_code
 
 def get_streamlit_error_correction_prompt(error: str, streamlit_app_code: str) -> str:
-    
+
     existing_streamlit_app_code_with_line_numbers = add_line_numbers_to_code(streamlit_app_code)
-    
-    return f"""You've created a Streamlit app, but it has an error in it when you try to run it.
+
+    return f"""You've created a Vizro dashboard app, but it has an error in it when you try to run it.
 
 Your job is to fix the error now. Only fix the specific error that you are instructed to fix now. Do not fix other error that that you anticipate. You will be asked to fix other errors later.
 
 {search_replace_instructions}
 
 ===============================================
 
-EXISTING STREAMLIT APP:
+EXISTING VIZRO APP:
 {existing_streamlit_app_code_with_line_numbers}
 
 ===============================================

mito-ai/mito_ai/streamlit_conversion/prompts/streamlit_finish_todo_prompt.py

@@ -6,10 +6,10 @@
 from mito_ai.streamlit_conversion.prompts.prompt_utils import add_line_numbers_to_code
 
 def get_finish_todo_prompt(notebook: List[dict], existing_streamlit_app_code: str, todo_placeholder: str) -> str:
-  
+
     existing_streamlit_app_code_with_line_numbers = add_line_numbers_to_code(existing_streamlit_app_code)
-    
-    return f"""You've already created the first draft of a Streamlit app representation of a Jupyter notebook, but you left yourself some TODOs marked as `{MITO_TODO_PLACEHOLDER}`. 
+
+    return f"""You've already created the first draft of a Vizro dashboard representation of a Jupyter notebook, but you left yourself some TODOs marked as `{MITO_TODO_PLACEHOLDER}`.
 
 **CRITICAL COMPLETION REQUIREMENT:**
 You have ONE and ONLY ONE opportunity to complete this TODO. If you do not finish the entire task completely, the application will be broken and unusable. This is your final chance to get it right.
@@ -29,17 +29,17 @@ def get_finish_todo_prompt(notebook: List[dict], existing_streamlit_app_code: st
 
 ===============================================
 
-Input Notebook that you are converting into the Streamlit app:
+Input Notebook that you are converting into the Vizro dashboard:
 {notebook}
 
 ===============================================
 
-EXISTING STREAMLIT APP:
+EXISTING VIZRO APP:
 {existing_streamlit_app_code_with_line_numbers}
 
 ===============================================
 
 Please make the changes for this TODO. Only focus on this one TODO right now. You will be asked to fix others later:
 {todo_placeholder}
 
-"""
+"""

mito-ai/mito_ai/streamlit_conversion/prompts/streamlit_system_prompt.py

@@ -1,56 +1,177 @@
 # Copyright (c) Saga Inc.
 # Distributed under the terms of the GNU Affero General Public License v3.0 License.
 
-streamlit_system_prompt = """You are a code conversion specialist who converts Jupyter notebooks into Streamlit applications with ABSOLUTE FIDELITY.
+streamlit_system_prompt = """You are a code conversion specialist who converts Jupyter notebooks into Vizro dashboard applications with ABSOLUTE FIDELITY.
 
 ROLE AND EXPERTISE:
-- Expert in Python, Jupyter notebooks, Streamlit, and data visualization
+- Expert in Python, Jupyter notebooks, Vizro, and data visualization
 - Experienced in creating executive-ready dashboards for business stakeholders
 - Skilled in translating technical analysis into clear, interactive presentations
 
 TASK REQUIREMENTS:
-1. Convert Jupyter notebook content into a complete Streamlit application (app.py)
+1. Convert Jupyter notebook content into a complete Vizro dashboard application (app.py)
 2. Preserve ALL outputs from code cells and markdown cells as they appear in the notebook
 3. Maintain the logical flow and structure of the original analysis
 4. Create an executive-friendly dashboard suitable for company leadership
 
-STREAMLIT IMPLEMENTATION GUIDELINES:
-- Use appropriate Streamlit components (st.title, st.header, st.subheader, st.markdown, etc.)
-- Display all visualizations using st.pyplot(), st.plotly_chart(), or st.altair_chart() as appropriate
-- Do not convert database connections into Streamlit's secret.toml format. If the user inlined their database credentials, are importing from an environment variable, or reading from a connections file, assume that same approach will work in the streamlit app.
-- Show dataframes and tables using st.dataframe() or st.table()
-- Include all text explanations and insights from markdown cells
-- Add interactive elements where beneficial (filters, selectors, etc.)
+VIZRO IMPLEMENTATION GUIDELINES:
+- Use Vizro's declarative structure: vm.Page, vm.Dashboard, vm.Graph, vm.Table, vm.Card, etc.
+- Import vizro.plotly.express as px for simple charts
+- Display visualizations using vm.Graph with @capture("graph") decorated functions
+- Show dataframes and tables using vm.Table with @capture("table") decorated functions
+- Include text content using vm.Card(text="...")
+- Add interactive elements using vm.Filter() and vm.Parameter() where beneficial
+- Do not convert database connections. If the user inlined their database credentials, are importing from an environment variable, or reading from a connections file, assume that same approach will work in the vizro app.
 - Ensure professional styling and layout suitable for executives
-- Just create the streamlit app code, do not include a _main_ function block. The file will be run directly using `streamlit run app.py`.
+- Structure the app as: load data → define @capture functions → create page(s) with components → build dashboard → run
+
+CRITICAL VIZRO CONSTRAINTS:
+
+**1. vm.Graph ONLY accepts @capture("graph") decorated functions, NOT raw Plotly figures!**
+
+WRONG - This will NOT work:
+```python
+import plotly.graph_objects as go
+fig = go.Figure(go.Bar(x=[1,2,3], y=[4,5,6]))
+vm.Graph(figure=fig)  # ERROR! Cannot use go.Figure directly
+```
+
+CORRECT - Use @capture("graph") decorator:
+```python
+from vizro.models.types import capture
+
+@capture("graph")
+def my_bar_chart(data_frame):
+    import plotly.graph_objects as go
+    fig = go.Figure(go.Bar(
+        x=data_frame['category'],
+        y=data_frame['value']
+    ))
+    fig.update_layout(title='My Chart')
+    return fig
+
+# Then use it:
+vm.Graph(figure=my_bar_chart(df))
+```
+
+**2. vm.Table ONLY accepts @capture("table") decorated functions, NOT raw DataFrames!**
+
+WRONG - This will NOT work:
+```python
+vm.Table(figure=df)  # ERROR! Cannot use DataFrame directly
+vm.Table(figure=df.head(100))  # ERROR! Still wrong
+```
+
+CORRECT - Use @capture("table") decorator:
+```python
+from vizro.models.types import capture
+
+@capture("table")
+def my_table(data_frame):
+    from dash import dash_table
+    return dash_table.DataTable(
+        data=data_frame.to_dict('records'),
+        columns=[{"name": i, "id": i} for i in data_frame.columns],
+        page_size=20,
+        style_table={'overflowX': 'auto'},
+        style_cell={'textAlign': 'left', 'padding': '10px'},
+        style_header={'backgroundColor': 'lightgrey', 'fontWeight': 'bold'}
+    )
+
+# Then use it:
+vm.Table(figure=my_table(df))
+```
+
+VIZRO APP STRUCTURE:
+```python
+import pandas as pd
+import vizro.plotly.express as px
+from vizro import Vizro
+import vizro.models as vm
+from vizro.models.types import capture
+
+# Data loading and processing code here
+df = pd.read_csv('data.csv')
+
+# Define graph functions with @capture("graph")
+@capture("graph")
+def my_bar_chart(data_frame):
+    import plotly.graph_objects as go
+    fig = go.Figure(go.Bar(x=data_frame['x'], y=data_frame['y']))
+    fig.update_layout(title='My Chart', height=500)
+    return fig
+
+@capture("graph")
+def my_scatter_chart(data_frame):
+    import plotly.graph_objects as go
+    fig = go.Figure(go.Scatter(x=data_frame['x'], y=data_frame['y'], mode='markers'))
+    return fig
+
+# Define table function with @capture("table")
+@capture("table")
+def data_table(data_frame):
+    from dash import dash_table
+    return dash_table.DataTable(
+        data=data_frame.to_dict('records'),
+        columns=[{"name": i, "id": i} for i in data_frame.columns],
+        page_size=20,
+        style_table={'overflowX': 'auto'},
+        style_cell={'textAlign': 'left', 'padding': '10px'},
+        style_header={'backgroundColor': 'lightgrey', 'fontWeight': 'bold'}
+    )
+
+# Create page with components
+page = vm.Page(
+    title="Dashboard Title",
+    components=[
+        vm.Card(text="# Dashboard Overview\\n\\nThis dashboard shows..."),
+        vm.Graph(figure=my_bar_chart(df)),
+        vm.Graph(figure=my_scatter_chart(df)),
+        vm.Table(figure=data_table(df.head(100)))
+    ],
+    controls=[
+        vm.Filter(column="category"),
+    ],
+)
+
+# Build and run dashboard
+dashboard = vm.Dashboard(pages=[page])
+Vizro().build(dashboard).run()
+```
 
 CRITICAL REQUIREMENTS:
 1. **PRESERVE ALL CODE EXACTLY**: Every line of code, every data structure, every import must be included in full
 2. **NO PLACEHOLDERS**: Never use comments like "# Add more data here" or "# Fill in the rest"
 3. **NO SIMPLIFICATION**: Do not replace actual data with sample data or hardcoded examples
 4. **COMPLETE DATA STRUCTURES**: If a notebook has a 1000-line dictionary, include all 1000 lines
-5. **PRESERVE DATA LOADING**: If the notebook reads from files, the Streamlit app must read from the same files
-6. **NO IMPROVIZAITION**: Do not provide your own interpretations of the analysis. Just convert the existing analysis into a streamlit app.
+5. **PRESERVE DATA LOADING**: If the notebook reads from files, the Vizro app must read from the same files
+6. **NO IMPROVISATION**: Do not provide your own interpretations of the analysis. Just convert the existing analysis into a vizro dashboard.
+7. **USE @capture DECORATORS**: Always use @capture("graph") for vm.Graph and @capture("table") for vm.Table - NEVER pass raw figures or DataFrames
 
-STYLE GUIDELINES: 
+COMPONENT MAPPING:
+- Code cells that generate plots → @capture("graph") function + vm.Graph(figure=func(df))
+- DataFrames to display → @capture("table") function + vm.Table(figure=func(df))
+- Markdown cells with text → vm.Card(text="...")
+- Variables that users might want to filter → vm.Filter(column="...")
+- Maintain the order of outputs as they appear in the notebook
+
+CHART TYPE CONVERSIONS (all must use @capture("graph")):
+- matplotlib/seaborn plots → convert to go.Figure inside @capture("graph") function
+- go.Bar → @capture("graph") function returning go.Figure(go.Bar(...))
+- go.Scatter → @capture("graph") function returning go.Figure(go.Scatter(...))
+- go.Histogram → @capture("graph") function returning go.Figure(go.Histogram(...))
+- px.scatter_mapbox → @capture("graph") function returning go.Figure with go.Scattermapbox
+
+STYLE GUIDELINES:
 - Create a professional, executive-friendly dashboard
-- If there are variables in the notebook that the streamlit app viewer would likely want to configure, then use the appropriate streamlit component to allow them to do so. For examples, if the notebook has a variable called "start_date" and "end_date", then use the st.date_input component to allow the user to select the start and end dates.
+- If there are variables in the notebook that viewers would likely want to configure, use vm.Filter or vm.Parameter
 - Do not use emojis unless they are in the notebook already
-- Do not modify the graphs or analysis. If the notebook has a graph, use the same graph in the streamlit app.
-- Always include the following code at the top of the file so the user does not use the wrong deploy button
-```python
-st.markdown(\"\"\"
-    <style>
-        #MainMenu {visibility: hidden;}
-        .stAppDeployButton {display:none;}
-        footer {visibility: hidden;}
-        .stMainBlockContainer {padding: 2rem 1rem 2rem 1rem;}
-    </style>
-\"\"\", unsafe_allow_html=True)
-```
+- Do not modify the graphs or analysis. If the notebook has a graph, convert it to Vizro-compatible format using @capture
+- Always end with: Vizro().build(dashboard).run()
 
 OUTPUT FORMAT:
-- Output the complete, runnable app.py file.
-- Do not output any extra text, just give the python code. 
+- Output the complete, runnable app.py file
+- Do not output any extra text, just give the python code
+- The file will be run directly using Python (not streamlit run)
 
-"""
+"""

mito-ai/mito_ai/streamlit_conversion/prompts/update_existing_app_prompt.py

@@ -7,14 +7,14 @@
 
 def get_update_existing_app_prompt(notebook: List[dict], streamlit_app_code: str, edit_prompt: str) -> str:
     """
-    This prompt is used to update an existing streamlit app.
+    This prompt is used to update an existing vizro app.
     """
-    
+
     existing_streamlit_app_code_with_line_numbers = add_line_numbers_to_code(streamlit_app_code)
-    
+
     return f"""
 
-GOAL: You've previously created a first draft of the Streamlit app. Now the user reviewed it and provided feedback.Update the existing streamlit app according to the feedback provided by the user. Use the input notebook to help you understand what code needs to be added, changed, or modified to fulfill the user's edit request.
+GOAL: You've previously created a first draft of the Vizro dashboard app. Now the user reviewed it and provided feedback. Update the existing vizro app according to the feedback provided by the user. Use the input notebook to help you understand what code needs to be added, changed, or modified to fulfill the user's edit request.
 
 **CRITICAL COMPLETION REQUIREMENT:**
 You have ONE and ONLY ONE opportunity to complete this edit request. If you do not finish the entire task completely, the application will be broken and unusable. This is your final chance to get it right.
@@ -34,17 +34,17 @@ def get_update_existing_app_prompt(notebook: List[dict], streamlit_app_code: str
 
 ===============================================
 
-INPUT NOTEBOOK: 
+INPUT NOTEBOOK:
 {notebook}
 
 ===============================================
 
-EXISTING STREAMLIT APP:
+EXISTING VIZRO APP:
 {existing_streamlit_app_code_with_line_numbers}
 
 ===============================================
 
 USER EDIT REQUEST:
 {edit_prompt}
 
-"""
+"""

mito-ai/mito_ai/streamlit_conversion/streamlit_agent_handler.py

@@ -120,24 +120,24 @@ async def streamlit_handler(notebook_path: AbsoluteNotebookPath, app_file_name:
         # Otherwise generate a new streamlit app
         streamlit_code = await generate_new_streamlit_code(notebook_code)
 
-    # Then, after creating/updating the app, validate that the new code runs 
-    errors = validate_app(streamlit_code, notebook_path)
-    tries = 0
-    while len(errors)>0 and tries < 5:
-        for error in errors:
-            streamlit_code = await correct_error_in_generation(error, streamlit_code)
-        
-        errors = validate_app(streamlit_code, notebook_path)
-        
-        if len(errors)>0:
-            # TODO: We can't easily get the key type here, so for the beta release
-            # we are just defaulting to the mito server key since that is by far the most common.
-            log_streamlit_app_validation_retry('mito_server_key', MessageType.STREAMLIT_CONVERSION, errors)
-        tries+=1
+    # Validation disabled for Vizro POC
+    # errors = validate_app(streamlit_code, notebook_path)
+    # tries = 0
+    # while len(errors)>0 and tries < 5:
+    #     for error in errors:
+    #         streamlit_code = await correct_error_in_generation(error, streamlit_code)
+    #
+    #     errors = validate_app(streamlit_code, notebook_path)
+    #
+    #     if len(errors)>0:
+    #         # TODO: We can't easily get the key type here, so for the beta release
+    #         # we are just defaulting to the mito server key since that is by far the most common.
+    #         log_streamlit_app_validation_retry('mito_server_key', MessageType.STREAMLIT_CONVERSION, errors)
+    #     tries+=1
 
-    if len(errors)>0:
-        final_errors = ', '.join(errors)
-        raise StreamlitConversionError(f"Streamlit agent failed generating code after max retries. Errors: {final_errors}", 500)
+    # if len(errors)>0:
+    #     final_errors = ', '.join(errors)
+    #     raise StreamlitConversionError(f"Streamlit agent failed generating code after max retries. Errors: {final_errors}", 500)
 
     # Finally, update the app.py file with the new code
     create_app_file(app_path, streamlit_code)