vizro-mcp 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

vizro_mcp/server.py

@@ -8,45 +8,39 @@ from typing import Any, Literal, Optional
 
 import vizro.models as vm
 from mcp.server.fastmcp import FastMCP
-from pydantic import ValidationError
+from pydantic import Field, ValidationError
 from vizro import Vizro
 
 from vizro_mcp._schemas import (
     AgGridEnhanced,
     ChartPlan,
-    ContainerSimplified,
-    DashboardSimplified,
-    FilterSimplified,
     GraphEnhanced,
-    PageSimplified,
-    ParameterSimplified,
-    TabsSimplified,
-    get_overview_vizro_models,
-    get_simple_dashboard_config,
 )
 from vizro_mcp._utils import (
+    CHART_INSTRUCTIONS,
     GAPMINDER,
     IRIS,
-    SAMPLE_DASHBOARD_CONFIG,
     STOCKS,
     TIPS,
     DFInfo,
     DFMetaData,
+    NoDefsGenerateJsonSchema,
     convert_github_url_to_raw,
     create_pycafe_url,
+    get_chart_prompt,
+    get_dashboard_instructions,
+    get_dashboard_prompt,
     get_dataframe_info,
     get_python_code_and_preview_link,
+    get_starter_dashboard_prompt,
     load_dataframe_by_format,
     path_or_url_check,
 )
 
-# PyCafe URL for Vizro snippets
-PYCAFE_URL = "https://py.cafe"
-
 
 @dataclass
-class ValidationResults:
-    """Results of the validation tool."""
+class ValidateResults:
+    """Results of validation tools."""
 
     valid: bool
     message: str
@@ -65,96 +59,50 @@ class DataAnalysisResults:
     df_metadata: Optional[DFMetaData]
 
 
-# TODO: what do I need to do here, as things are already set up?
-mcp = FastMCP(
-    "MCP server to help create Vizro dashboards and charts.",
-)
-
-
-@mcp.tool()
-def get_sample_data_info(data_name: Literal["iris", "tips", "stocks", "gapminder"]) -> DFMetaData:
-    """If user provides no data, use this tool to get sample data information.
+@dataclass
+class ModelJsonSchemaResults:
+    """Results of the get_model_json_schema tool."""
 
-    Use the following data for the below purposes:
-        - iris: mostly numerical with one categorical column, good for scatter, histogram, boxplot, etc.
-        - tips: contains mix of numerical and categorical columns, good for bar, pie, etc.
-        - stocks: stock prices, good for line, scatter, generally things that change over time
-        - gapminder: demographic data, good for line, scatter, generally things with maps or many categories
+    model_name: str
+    json_schema: dict[str, Any]
+    additional_info: str
 
-    Args:
-        data_name: Name of the dataset to get sample data for
 
-    Returns:
-        Data info object containing information about the dataset.
-    """
-    if data_name == "iris":
-        return IRIS
-    elif data_name == "tips":
-        return TIPS
-    elif data_name == "stocks":
-        return STOCKS
-    elif data_name == "gapminder":
-        return GAPMINDER
+# TODO: check on https://github.com/modelcontextprotocol/python-sdk what new things are possible to do here
+mcp = FastMCP(
+    "MCP server to help create Vizro dashboards and charts.",
+)
 
 
 @mcp.tool()
-def validate_model_config(
-    config: dict[str, Any],
-    data_infos: list[DFMetaData],  # Should be Optional[..]=None, but Cursor complains..
-    auto_open: bool = True,
-) -> ValidationResults:
-    """Validate Vizro model configuration. Run ALWAYS when you have a complete dashboard configuration.
+def get_vizro_chart_or_dashboard_plan(
+    user_plan: Literal["chart", "dashboard"],
+    user_host: Literal["generic_host", "ide"],
+    advanced_mode: bool = False,
+) -> str:
+    """Get instructions for creating a Vizro chart or dashboard. Call FIRST when asked to create Vizro things.
 
-    If successful, the tool will return the python code and, if it is a remote file, the py.cafe link to the chart.
-    The PyCafe link will be automatically opened in your default browser if auto_open is True.
+    Must be ALWAYS called FIRST with advanced_mode=False, then call again with advanced_mode=True
+    if the JSON config does not suffice anymore.
 
     Args:
-        config: Either a JSON string or a dictionary representing a Vizro model configuration
-        data_infos: List of DFMetaData objects containing information about the data files
-        auto_open: Whether to automatically open the PyCafe link in a browser
+        user_plan: The type of Vizro thing the user wants to create
+        user_host: The host the user is using, if "ide" you can use the IDE/editor to run python code
+        advanced_mode: Only call if you need to use custom CSS, custom components or custom actions.
+            No need to call this with advanced_mode=True if you need advanced charts, use `custom_charts` in
+            the `validate_dashboard_config` tool instead.
 
     Returns:
-        ValidationResults object with status and dashboard details
+        Instructions for creating a Vizro chart or dashboard
     """
-    Vizro._reset()
-
-    try:
-        dashboard = vm.Dashboard.model_validate(config)
-    except ValidationError as e:
-        return ValidationResults(
-            valid=False,
-            message=f"Validation Error: {e!s}",
-            python_code="",
-            pycafe_url=None,
-            browser_opened=False,
-        )
-
-    else:
-        result = get_python_code_and_preview_link(dashboard, data_infos)
-
-        pycafe_url = result.pycafe_url if all(info.file_location_type == "remote" for info in data_infos) else None
-        browser_opened = False
-
-        if pycafe_url and auto_open:
-            try:
-                browser_opened = webbrowser.open(pycafe_url)
-            except Exception:
-                browser_opened = False
-
-        return ValidationResults(
-            valid=True,
-            message="Configuration is valid for Dashboard!",
-            python_code=result.python_code,
-            pycafe_url=pycafe_url,
-            browser_opened=browser_opened,
-        )
-
-    finally:
-        Vizro._reset()
+    if user_plan == "chart":
+        return CHART_INSTRUCTIONS
+    elif user_plan == "dashboard":
+        return f"{get_dashboard_instructions(advanced_mode, user_host)}"
 
 
 @mcp.tool()
-def get_model_json_schema(model_name: str) -> dict[str, Any]:
+def get_model_json_schema(model_name: str) -> ModelJsonSchemaResults:
     """Get the JSON schema for the specified Vizro model.
 
     Args:
@@ -163,88 +111,65 @@ def get_model_json_schema(model_name: str) -> dict[str, Any]:
     Returns:
         JSON schema of the requested Vizro model
     """
-    # Dictionary mapping model names to their simplified versions
     modified_models = {
-        "Page": PageSimplified,
-        "Dashboard": DashboardSimplified,
         "Graph": GraphEnhanced,
         "AgGrid": AgGridEnhanced,
         "Table": AgGridEnhanced,
-        "Tabs": TabsSimplified,
-        "Container": ContainerSimplified,
-        "Filter": FilterSimplified,
-        "Parameter": ParameterSimplified,
     }
 
-    # Check if model_name is in the simplified models dictionary
     if model_name in modified_models:
-        return modified_models[model_name].model_json_schema()
+        return ModelJsonSchemaResults(
+            model_name=model_name,
+            json_schema=modified_models[model_name].model_json_schema(schema_generator=NoDefsGenerateJsonSchema),
+            additional_info="""LLM must remember to replace `$ref` with the actual config. Request the schema of
+that model if necessary. Do NOT forget to call `validate_dashboard_config` after each iteration.""",
+        )
 
-    # Check if model exists in vizro.models
     if not hasattr(vm, model_name):
-        return {"error": f"Model '{model_name}' not found in vizro.models"}
+        return ModelJsonSchemaResults(
+            model_name=model_name,
+            json_schema={},
+            additional_info=f"Model '{model_name}' not found in vizro.models",
+        )
 
-    # Get schema for standard model
     model_class = getattr(vm, model_name)
-    return model_class.model_json_schema()
+    return ModelJsonSchemaResults(
+        model_name=model_name,
+        json_schema=model_class.model_json_schema(schema_generator=NoDefsGenerateJsonSchema),
+        additional_info="""LLM must remember to replace `$ref` with the actual config. Request the schema of
+that model if necessary. Do NOT forget to call `validate_dashboard_config` after each iteration.""",
+    )
 
 
 @mcp.tool()
-def get_vizro_chart_or_dashboard_plan(user_plan: Literal["chart", "dashboard"]) -> str:
-    """Get instructions for creating a Vizro chart or dashboard. Call FIRST when asked to create Vizro things."""
-    if user_plan == "chart":
-        return """
-IMPORTANT:
-    - KEEP IT SIMPLE: rather than iterating yourself, ask the user for more instructions
-    - ALWAYS VALIDATE:if you iterate over a valid produced solution, make sure to ALWAYS call the
-        validate_chart_code tool to validate the chart code, display the figure code to the user
-    - DO NOT modify the background (with plot_bgcolor) or color sequences unless explicitly asked for
-
-Instructions for creating a Vizro chart:
-    - analyze the datasets needed for the chart using the load_and_analyze_data tool - the most important
-        information here are the column names and column types
-    - if the user provides no data, but you need to display a chart or table, use the get_sample_data_info
-        tool to get sample data information
-    - create a chart using plotly express and/or plotly graph objects, and call the function `custom_chart`
-    - call the validate_chart_code tool to validate the chart code, display the figure code to the user (as artifact)
-    - do NOT call any other tool after, especially do NOT create a dashboard
-            """
-    elif user_plan == "dashboard":
-        return f"""
-IMPORTANT:
-    - KEEP IT SIMPLE: rather than iterating yourself, ask the user for more instructions
-    - ALWAYS VALIDATE:if you iterate over a valid produced solution, make sure to ALWAYS call the
-        validate_model_config tool again to ensure the solution is still valid
-    - DO NOT show any code or config to the user until you have validated the solution, do not say you are preparing
-        a solution, just do it and validate it
-    - IF STUCK: try enquiring the schema of the component in question
-
-
-Instructions for creating a Vizro dashboard:
-    - IF the user has no plan (ie no components or pages), use the config at the bottom of this prompt
-        and validate that solution without any additions, OTHERWISE:
-    - analyze the datasets needed for the dashboard using the load_and_analyze_data tool - the most
-        important information here are the column names and column types
-    - if the user provides no data, but you need to display a chart or table, use the get_sample_data_info
-        tool to get sample data information
-    - make a plan of what components you would like to use, then request all necessary schemas
-        using the get_model_json_schema tool
-    - assemble your components into a page, then add the page or pages to a dashboard, DO NOT show config or code
-        to the user until you have validated the solution
-    - ALWAYS validate the dashboard configuration using the validate_model_config tool
-    - if you display any code artifact, you must use the above created code, do not add new config to it
-
-Models you can use:
-{get_overview_vizro_models()}
-
-Very simple dashboard config:
-{get_simple_dashboard_config()}
+def get_sample_data_info(data_name: Literal["iris", "tips", "stocks", "gapminder"]) -> DFMetaData:
+    """If user provides no data, use this tool to get sample data information.
+
+    Use the following data for the below purposes:
+        - iris: mostly numerical with one categorical column, good for scatter, histogram, boxplot, etc.
+        - tips: contains mix of numerical and categorical columns, good for bar, pie, etc.
+        - stocks: stock prices, good for line, scatter, generally things that change over time
+        - gapminder: demographic data, good for line, scatter, generally things with maps or many categories
+
+    Args:
+        data_name: Name of the dataset to get sample data for
+
+    Returns:
+        Data info object containing information about the dataset.
     """
+    if data_name == "iris":
+        return IRIS
+    elif data_name == "tips":
+        return TIPS
+    elif data_name == "stocks":
+        return STOCKS
+    elif data_name == "gapminder":
+        return GAPMINDER
 
 
 @mcp.tool()
 def load_and_analyze_data(path_or_url: str) -> DataAnalysisResults:
-    """Load data from various file formats into a pandas DataFrame and analyze its structure.
+    """Use to understand local or remote data files. Must be called with absolute paths or URLs.
 
     Supported formats:
     - CSV (.csv)
@@ -255,7 +180,7 @@ def load_and_analyze_data(path_or_url: str) -> DataAnalysisResults:
     - Parquet (.parquet)
 
     Args:
-        path_or_url: Local file path or URL to a data file
+        path_or_url: Absolute (important!) local file path or URL to a data file
 
     Returns:
         DataAnalysisResults object containing DataFrame information and metadata
@@ -276,7 +201,14 @@ def load_and_analyze_data(path_or_url: str) -> DataAnalysisResults:
         df, read_fn = load_dataframe_by_format(processed_path_or_url, mime_type)
 
     except Exception as e:
-        return DataAnalysisResults(valid=False, message=f"Failed to load data: {e!s}", df_info=None, df_metadata=None)
+        return DataAnalysisResults(
+            valid=False,
+            message=f"""Failed to load data: {e!s}. Remember to use the ABSOLUTE path or URL!
+Alternatively, you can use any data analysis means available to you. Most important information are the column names and
+column types for passing along to the `validate_dashboard_config` or `validate_chart_code` tools.""",
+            df_info=None,
+            df_metadata=None,
+        )
 
     df_info = get_dataframe_info(df)
     df_metadata = DFMetaData(
@@ -289,58 +221,97 @@ def load_and_analyze_data(path_or_url: str) -> DataAnalysisResults:
     return DataAnalysisResults(valid=True, message="Data loaded successfully", df_info=df_info, df_metadata=df_metadata)
 
 
+# TODO: Additional things we could validate:
+# - data_infos: check we are referring to the correct dataframe, or at least A DF
+@mcp.tool()
+def validate_dashboard_config(
+    dashboard_config: dict[str, Any],
+    data_infos: list[DFMetaData],
+    custom_charts: list[ChartPlan],
+    auto_open: bool = True,
+) -> ValidateResults:
+    """Validate Vizro model configuration. Run ALWAYS when you have a complete dashboard configuration.
+
+    If successful, the tool will return the python code and, if it is a remote file, the py.cafe link to the chart.
+    The PyCafe link will be automatically opened in your default browser if auto_open is True.
+
+    Args:
+        dashboard_config: Either a JSON string or a dictionary representing a Vizro dashboard model configuration
+        data_infos: List of DFMetaData objects containing information about the data files
+        custom_charts: List of ChartPlan objects containing information about the custom charts in the dashboard
+        auto_open: Whether to automatically open the PyCafe link in a browser
+
+    Returns:
+        ValidationResults object with status and dashboard details
+    """
+    Vizro._reset()
+
+    try:
+        dashboard = vm.Dashboard.model_validate(
+            dashboard_config,
+            context={"allow_undefined_captured_callable": [custom_chart.chart_name for custom_chart in custom_charts]},
+        )
+    except ValidationError as e:
+        return ValidateResults(
+            valid=False,
+            message=f"""Validation Error: {e!s}. Fix the error and call this tool again.
+Calling `get_model_json_schema` may help.""",
+            python_code="",
+            pycafe_url=None,
+            browser_opened=False,
+        )
+
+    else:
+        code_link = get_python_code_and_preview_link(dashboard, data_infos, custom_charts)
+
+        pycafe_url = code_link.pycafe_url if all(info.file_location_type == "remote" for info in data_infos) else None
+        browser_opened = False
+
+        if pycafe_url and auto_open:
+            try:
+                browser_opened = webbrowser.open(pycafe_url)
+            except Exception:
+                browser_opened = False
+
+        return ValidateResults(
+            valid=True,
+            message="""Configuration is valid for Dashboard! Do not forget to call this tool again after each iteration.
+If you are creating an `app.py` file, you MUST use the code from the validation tool, do not modify it, watch out for
+differences to previous `app.py`""",
+            python_code=code_link.python_code,
+            pycafe_url=pycafe_url,
+            browser_opened=browser_opened,
+        )
+
+    finally:
+        Vizro._reset()
+
+
 @mcp.prompt()
 def create_starter_dashboard():
     """Prompt template for getting started with Vizro."""
-    content = f"""
-Create a super simple Vizro dashboard with one page and one chart and one filter:
-- No need to call any tools except for validate_model_config
-- Call this tool with the precise config as shown below
-- The PyCafe link will be automatically opened in your default browser
-- THEN show the python code after validation, but do not show the PyCafe link
-- Be concise, do not explain anything else, just create the dashboard
-- Finally ask the user what they would like to do next, then you can call other tools to get more information,
-    you should then start with the get_chart_or_dashboard_plan tool
-
-{SAMPLE_DASHBOARD_CONFIG}
-"""
-    return content
+    return get_starter_dashboard_prompt()
 
 
 @mcp.prompt()
-def create_eda_dashboard(
-    file_path_or_url: str,
+def create_dashboard(
+    file_path_or_url: str = Field(description="The absolute path or URL to the data file you want to use."),
+    context: Optional[str] = Field(default=None, description="(Optional) Describe the dashboard you want to create."),
 ) -> str:
     """Prompt template for creating an EDA dashboard based on one dataset."""
-    content = f"""
-Create an EDA dashboard based on the following dataset:{file_path_or_url}. Proceed as follows:
-1. Analyze the data using the load_and_analyze_data tool first, passing the file path or github url {file_path_or_url}
-    to the tool.
-2. Create a dashboard with 4 pages:
-    - Page 1: Summary of the dataset using the Card component and the dataset itself using the plain AgGrid component.
-    - Page 2: Visualizing the distribution of all numeric columns using the Graph component with a histogram.
-        - use a Parameter that targets the Graph component and the x argument, and you can select the column to
-            be displayed
-        - IMPORTANT:remember that you target the chart like: <graph_id>.x and NOT <graph_id>.figure.x
-        - do not use any color schemes etc.
-        - add filters for all categorical columns
-    - Page 3: Visualizing the correlation between all numeric columns using the Graph component with a scatter plot.
-    - Page 4: Two interesting charts side by side, use the Graph component for this. Make sure they look good
-        but do not try something beyond the scope of plotly express
-"""
-    return content
+    return get_dashboard_prompt(file_path_or_url, context)
 
 
 @mcp.tool()
 def validate_chart_code(
-    config: ChartPlan,
+    chart_config: ChartPlan,
     data_info: DFMetaData,
     auto_open: bool = True,
-) -> ValidationResults:
+) -> ValidateResults:
     """Validate the chart code created by the user and optionally open the PyCafe link in a browser.
 
     Args:
-        config: A ChartPlan object with the chart configuration
+        chart_config: A ChartPlan object with the chart configuration
         data_info: Metadata for the dataset to be used in the chart
         auto_open: Whether to automatically open the PyCafe link in a browser
 
@@ -350,9 +321,9 @@ def validate_chart_code(
     Vizro._reset()
 
     try:
-        chart_plan_obj = ChartPlan.model_validate(config)
+        chart_plan_obj = ChartPlan.model_validate(chart_config)
     except ValidationError as e:
-        return ValidationResults(
+        return ValidateResults(
             valid=False,
             message=f"Validation Error: {e!s}",
             python_code="",
@@ -372,7 +343,7 @@ def validate_chart_code(
             except Exception:
                 browser_opened = False
 
-        return ValidationResults(
+        return ValidateResults(
             valid=True,
             message="Chart only dashboard created successfully!",
             python_code=chart_plan_obj.get_chart_code(vizro=True),
@@ -386,15 +357,8 @@ def validate_chart_code(
 
 @mcp.prompt()
 def create_vizro_chart(
-    chart_type: str,
-    file_path_or_url: Optional[str] = None,
+    file_path_or_url: str = Field(description="The absolute path or URL to the data file you want to use."),
+    context: Optional[str] = Field(default=None, description="(Optional) Describe the chart you want to create."),
 ) -> str:
     """Prompt template for creating a Vizro chart."""
-    content = f"""
- - Create a chart using the following chart type: {chart_type}.
- - You MUST name the function containing the fig `custom_chart`
- - Make sure to analyze the data using the load_and_analyze_data tool first, passing the file path or github url
- {file_path_or_url} OR choose the most appropriate sample data using the get_sample_data_info tool.
- Then you MUST use the validate_chart_code tool to validate the chart code.
-            """
-    return content
+    return get_chart_prompt(file_path_or_url, context)