vizro-mcp 0.0.1.dev0__py3-none-any.whl

vizro_mcp/__init__.py

@@ -0,0 +1,19 @@
+import logging
+import sys
+
+from .server import mcp
+
+__version__ = "0.0.1.dev0"
+
+
+def main():
+    """Run the Vizro MCP server - makes charts and dashboards available to AI assistants."""
+    # Configure logging to show warnings by default
+    logging.basicConfig(level=logging.WARNING, stream=sys.stderr)
+
+    # Run the MCP server
+    mcp.run()
+
+
+if __name__ == "__main__":
+    main()

vizro_mcp/_schemas/__init__.py

@@ -0,0 +1,29 @@
+from .schemas import (
+    MODEL_GROUPS,
+    AgGridEnhanced,
+    ChartPlan,
+    ContainerSimplified,
+    DashboardSimplified,
+    FilterSimplified,
+    GraphEnhanced,
+    PageSimplified,
+    ParameterSimplified,
+    TabsSimplified,
+    get_overview_vizro_models,
+    get_simple_dashboard_config,
+)
+
+__all__ = [
+    "MODEL_GROUPS",
+    "AgGridEnhanced",
+    "ChartPlan",
+    "ContainerSimplified",
+    "DashboardSimplified",
+    "FilterSimplified",
+    "GraphEnhanced",
+    "PageSimplified",
+    "ParameterSimplified",
+    "TabsSimplified",
+    "get_overview_vizro_models",
+    "get_simple_dashboard_config",
+]

vizro_mcp/_schemas/schemas.py

@@ -0,0 +1,297 @@
+"""Schema defining pydantic models for usage in the MCP server."""
+
+from typing import Annotated, Any, Literal, Optional
+
+import vizro.models as vm
+from pydantic import AfterValidator, BaseModel, Field, PrivateAttr, conlist
+
+from vizro_mcp._utils import SAMPLE_DASHBOARD_CONFIG, DFMetaData
+
+# Constants used in chart validation
+CUSTOM_CHART_NAME = "custom_chart"
+ADDITIONAL_IMPORTS = [
+    "import vizro.plotly.express as px",
+    "import plotly.graph_objects as go",
+    "import pandas as pd",
+    "import numpy as np",
+    "from vizro.models.types import capture",
+]
+
+# These types are used to simplify the schema for the LLM.
+SimplifiedComponentType = Literal["Card", "Button", "Text", "Container", "Tabs", "Graph", "AgGrid"]
+SimplifiedSelectorType = Literal[
+    "Dropdown", "RadioItems", "Checklist", "DatePicker", "Slider", "RangeSlider", "DatePicker"
+]
+SimplifiedControlType = Literal["Filter", "Parameter"]
+SimplifiedLayoutType = Literal["Grid", "Flex"]
+
+# This dict is used to give the model and overview of what is available in the vizro.models namespace.
+# It helps it to narrow down the choices when asking for a model.
+MODEL_GROUPS: dict[str, list[type[vm.VizroBaseModel]]] = {
+    "main": [vm.Dashboard, vm.Page],
+    "components": [vm.Card, vm.Button, vm.Text, vm.Container, vm.Tabs, vm.Graph, vm.AgGrid],  #'Figure', 'Table'
+    "layouts": [vm.Grid, vm.Flex],
+    "controls": [vm.Filter, vm.Parameter],
+    "selectors": [
+        vm.Dropdown,
+        vm.RadioItems,
+        vm.Checklist,
+        vm.DatePicker,
+        vm.Slider,
+        vm.RangeSlider,
+        vm.DatePicker,
+    ],
+    "navigation": [vm.Navigation, vm.NavBar, vm.NavLink],
+}
+
+
+# These simplified page, container, tabs and dashboard models are used to return a flatter schema to the LLM in order to
+# reduce the context size. Especially the dashboard model schema is huge as it contains all other models.
+
+
+class FilterSimplified(vm.Filter):
+    """Simplified Filter model for reduced schema. LLM should remember to insert actual components."""
+
+    selector: Optional[SimplifiedSelectorType] = Field(
+        default=None, description="Selector to be displayed. Only provide if asked for!"
+    )
+
+
+class ParameterSimplified(vm.Parameter):
+    """Simplified Parameter model for reduced schema. LLM should remember to insert actual components."""
+
+    selector: SimplifiedSelectorType = Field(description="Selector to be displayed.")
+
+
+class ContainerSimplified(vm.Container):
+    """Simplified Container model for reduced schema. LLM should remember to insert actual components."""
+
+    components: list[SimplifiedComponentType] = Field(description="List of component names to be displayed.")
+    layout: Optional[SimplifiedLayoutType] = Field(
+        default=None, description="Layout to place components in. Only provide if asked for!"
+    )
+
+
+class TabsSimplified(vm.Tabs):
+    """Simplified Tabs model for reduced schema. LLM should remember to insert actual components."""
+
+    tabs: conlist(ContainerSimplified, min_length=1)
+
+
+class PageSimplified(BaseModel):
+    """Simplified Page modes for reduced schema. LLM should remember to insert actual components."""
+
+    components: list[SimplifiedComponentType] = Field(description="List of component names to be displayed.")
+    title: str = Field(description="Title to be displayed.")
+    description: str = Field(default="", description="Description for meta tags.")
+    layout: Optional[SimplifiedLayoutType] = Field(
+        default=None, description="Layout to place components in. Only provide if asked for!"
+    )
+    controls: list[SimplifiedControlType] = Field(default=[], description="Controls to be displayed.")
+
+
+class DashboardSimplified(BaseModel):
+    """Simplified Dashboard model for reduced schema. LLM should remember to insert actual components."""
+
+    pages: list[Literal["Page"]] = Field(description="List of page names to be included in the dashboard.")
+    theme: Literal["vizro_dark", "vizro_light"] = Field(
+        default="vizro_dark", description="Theme to be applied across dashboard. Defaults to `vizro_dark`."
+    )
+    navigation: Optional[Literal["Navigation"]] = Field(
+        default=None, description="Navigation component for the dashboard. Only provide if asked for!"
+    )
+    title: str = Field(default="", description="Dashboard title to appear on every page on top left-side.")
+
+
+# These enhanced models are used to return a more complete schema to the LLM. Although we do not have actual schemas for
+# the figure fields, we can prompt the model via the description to produce something likely correct.
+class GraphEnhanced(vm.Graph):
+    """A Graph model that uses Plotly Express to create the figure."""
+
+    figure: dict[str, Any] = Field(
+        description="""
+This is the plotly express figure to be displayed. Only use valid plotly express functions to create the figure.
+Only use the arguments that are supported by the function you are using and where no extra modules such as statsmodels
+are needed (e.g. trendline).
+
+- Configure a dictionary as if this would be added as **kwargs to the function you are using.
+- You must use the key: "_target_: "<function_name>" to specify the function you are using. Do NOT precede by
+    namespace (like px.line)
+- you must refer to the dataframe by name, for now it is one of "gapminder", "iris", "tips".
+- do not use a title if your Graph model already has a title.
+"""
+    )
+
+
+class AgGridEnhanced(vm.AgGrid):
+    """AgGrid model that uses dash-ag-grid to create the figure."""
+
+    figure: dict[str, Any] = Field(
+        description="""
+This is the ag-grid figure to be displayed. Only use arguments from the [`dash-ag-grid` function](https://dash.plotly.com/dash-ag-grid/reference).
+
+The only difference to the dash version is that:
+    - you must use the key: "_target_: "dash_ag_grid"
+    - you must refer to data via "data_frame": <data_frame_name> and NOT via columnDefs and rowData (do NOT set)
+        """
+    )
+
+
+###### Chart functionality - not sure if I should include this in the MCP server
+def _strip_markdown(code_string: str) -> str:
+    """Remove any code block wrappers (markdown or triple quotes)."""
+    wrappers = [("```python\n", "```"), ("```py\n", "```"), ("```\n", "```"), ('"""', '"""'), ("'''", "'''")]
+
+    for start, end in wrappers:
+        if code_string.startswith(start) and code_string.endswith(end):
+            code_string = code_string[len(start) : -len(end)]
+            break
+
+    return code_string.strip()
+
+
+def _check_chart_code(v: str) -> str:
+    v = _strip_markdown(v)
+
+    # TODO: add more checks: ends with return, has return, no second function def, only one indented line
+    func_def = f"def {CUSTOM_CHART_NAME}("
+    if func_def not in v:
+        raise ValueError(f"The chart code must be wrapped in a function named `{CUSTOM_CHART_NAME}`")
+
+    v = v[v.index(func_def) :].strip()
+
+    first_line = v.split("\n")[0].strip()
+    if "data_frame" not in first_line:
+        raise ValueError(
+            """The chart code must accept a single argument `data_frame`,
+and it should be the first argument of the chart."""
+        )
+    return v
+
+
+class ChartPlan(BaseModel):
+    """Base chart plan used to generate chart code based on user visualization requirements."""
+
+    chart_type: str = Field(
+        description="""
+        Describes the chart type that best reflects the user request.
+        """,
+    )
+    imports: list[str] = Field(
+        description="""
+        List of import statements required to render the chart defined by the `chart_code` field. Ensure that every
+        import statement is a separate list/array entry: An example of valid list of import statements would be:
+
+        ["import pandas as pd",
+        "import plotly.express as px"]
+        """,
+    )
+    chart_code: Annotated[
+        str,
+        AfterValidator(_check_chart_code),
+        Field(
+            description="""
+        Python code that generates a generates a plotly go.Figure object. It must fulfill the following criteria:
+        1. Must be wrapped in a function name
+        2. Must accept a single argument `data_frame` which is a pandas DataFrame
+        3. Must return a plotly go.Figure object
+        4. All data used in the chart must be derived from the data_frame argument, all data manipulations
+        must be done within the function.
+        """,
+        ),
+    ]
+
+    _additional_vizro_imports: list[str] = PrivateAttr(ADDITIONAL_IMPORTS)
+
+    def get_imports(self, vizro: bool = False):
+        imports = list(dict.fromkeys(self.imports + self._additional_vizro_imports))  # remove duplicates
+        if vizro:  # TODO: improve code of below
+            imports = [imp for imp in imports if "import plotly.express as px" not in imp]
+        else:
+            imports = [imp for imp in imports if "vizro" not in imp]
+        return "\n".join(imports) + "\n"
+
+    def get_chart_code(self, chart_name: Optional[str] = None, vizro: bool = False):
+        chart_code = self.chart_code
+        if vizro:
+            chart_code = chart_code.replace(f"def {CUSTOM_CHART_NAME}", f"@capture('graph')\ndef {CUSTOM_CHART_NAME}")
+        if chart_name is not None:
+            chart_code = chart_code.replace(f"def {CUSTOM_CHART_NAME}", f"def {chart_name}")
+        return chart_code
+
+    def get_dashboard_template(self, data_info: DFMetaData) -> str:
+        """Create a simple dashboard template for displaying the chart.
+
+        Args:
+            data_info: The metadata of the dataset to use.
+
+        Returns:
+            Complete Python code for a Vizro dashboard displaying the chart.
+        """
+        chart_code = self.get_chart_code(vizro=True)
+        imports = self.get_imports(vizro=True)
+
+        # Add the Vizro-specific imports if not present
+        additional_imports = [
+            "import vizro.models as vm",
+            "from vizro import Vizro",
+            "from vizro.managers import data_manager",
+        ]
+
+        # Combine imports without duplicates
+        all_imports = list(dict.fromkeys(additional_imports + imports.split("\n")))
+
+        dashboard_template = f"""
+{chr(10).join(imp for imp in all_imports if imp)}
+
+# Load the data
+data_manager["{data_info.file_name}"] = {data_info.read_function_string}("{data_info.file_path_or_url}")
+
+
+# Custom chart code
+{chart_code}
+
+# Create a dashboard to display the chart
+dashboard = vm.Dashboard(
+    pages=[
+        vm.Page(
+            title="{self.chart_type.capitalize()} Chart",
+            components=[
+                vm.Graph(
+                    id="{self.chart_type}_graph",
+                    figure={CUSTOM_CHART_NAME}("{data_info.file_name}"),
+                )
+            ],
+        )
+    ],
+    title="{self.chart_type.capitalize()} Dashboard",
+)
+
+# Run the dashboard
+Vizro().build(dashboard).run()
+"""
+
+        return dashboard_template
+
+
+def get_overview_vizro_models() -> dict[str, list[dict[str, str]]]:
+    """Get all available models in the vizro.models namespace.
+
+    Returns:
+        Dictionary with categories of models and their descriptions
+    """
+    result: dict[str, list[dict[str, str]]] = {}
+    for category, models_list in MODEL_GROUPS.items():
+        result[category] = [
+            {
+                "name": model_class.__name__,
+                "description": (model_class.__doc__ or "No description available").split("\n")[0],
+            }
+            for model_class in models_list
+        ]
+    return result
+
+
+def get_simple_dashboard_config() -> str:
+    """Very simple Vizro dashboard configuration. Use this config as a starter when no other config is provided."""
+    return SAMPLE_DASHBOARD_CONFIG

vizro_mcp/_utils/__init__.py

@@ -0,0 +1,33 @@
+from .utils import (
+    GAPMINDER,
+    IRIS,
+    SAMPLE_DASHBOARD_CONFIG,
+    STOCKS,
+    TIPS,
+    DFInfo,
+    DFMetaData,
+    VizroCodeAndPreviewLink,
+    convert_github_url_to_raw,
+    create_pycafe_url,
+    get_dataframe_info,
+    get_python_code_and_preview_link,
+    load_dataframe_by_format,
+    path_or_url_check,
+)
+
+__all__ = [
+    "GAPMINDER",
+    "IRIS",
+    "SAMPLE_DASHBOARD_CONFIG",
+    "STOCKS",
+    "TIPS",
+    "DFInfo",
+    "DFMetaData",
+    "VizroCodeAndPreviewLink",
+    "convert_github_url_to_raw",
+    "create_pycafe_url",
+    "get_dataframe_info",
+    "get_python_code_and_preview_link",
+    "load_dataframe_by_format",
+    "path_or_url_check",
+]

vizro_mcp/_utils/utils.py

@@ -0,0 +1,300 @@
+"""Utility functions for the Vizro MCP."""
+
+import base64
+import gzip
+import io
+import json
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Literal, Optional, Union
+from urllib.parse import quote, urlencode
+
+import pandas as pd
+import vizro.models as vm
+
+# PyCafe URL for Vizro snippets
+PYCAFE_URL = "https://py.cafe"
+
+
+@dataclass
+class VizroCodeAndPreviewLink:
+    python_code: str
+    pycafe_url: str
+
+
+@dataclass
+class DFMetaData:
+    file_name: str
+    file_path_or_url: str
+    file_location_type: Literal["local", "remote"]
+    read_function_string: Literal["pd.read_csv", "pd.read_json", "pd.read_html", "pd.read_parquet", "pd.read_excel"]
+    column_names_types: Optional[dict[str, str]] = None
+
+
+@dataclass
+class DFInfo:
+    general_info: str
+    sample: dict[str, Any]
+
+
+IRIS = DFMetaData(
+    file_name="iris_data",
+    file_path_or_url="https://raw.githubusercontent.com/plotly/datasets/master/iris-id.csv",
+    file_location_type="remote",
+    read_function_string="pd.read_csv",
+    column_names_types={
+        "sepal_length": "float",
+        "sepal_width": "float",
+        "petal_length": "float",
+        "petal_width": "float",
+        "species": "str",
+    },
+)
+
+TIPS = DFMetaData(
+    file_name="tips_data",
+    file_path_or_url="https://raw.githubusercontent.com/plotly/datasets/master/tips.csv",
+    file_location_type="remote",
+    read_function_string="pd.read_csv",
+    column_names_types={
+        "total_bill": "float",
+        "tip": "float",
+        "sex": "str",
+        "smoker": "str",
+        "day": "str",
+        "time": "str",
+        "size": "int",
+    },
+)
+
+STOCKS = DFMetaData(
+    file_name="stocks_data",
+    file_path_or_url="https://raw.githubusercontent.com/plotly/datasets/master/stockdata.csv",
+    file_location_type="remote",
+    read_function_string="pd.read_csv",
+    column_names_types={
+        "Date": "str",
+        "IBM": "float",
+        "MSFT": "float",
+        "SBUX": "float",
+        "AAPL": "float",
+        "GSPC": "float",
+    },
+)
+
+GAPMINDER = DFMetaData(
+    file_name="gapminder_data",
+    file_path_or_url="https://raw.githubusercontent.com/plotly/datasets/master/gapminder_unfiltered.csv",
+    file_location_type="remote",
+    read_function_string="pd.read_csv",
+    column_names_types={
+        "country": "str",
+        "continent": "str",
+        "year": "int",
+        "lifeExp": "float",
+        "pop": "int",
+        "gdpPercap": "float",
+    },
+)
+
+SAMPLE_DASHBOARD_CONFIG = """
+{
+  `config`: {
+    `pages`: [
+      {
+        `title`: `Iris Data Analysis`,
+        `controls`: [
+          {
+            `id`: `species_filter`,
+            `type`: `filter`,
+            `column`: `species`,
+            `targets`: [
+              `scatter_plot`
+            ],
+            `selector`: {
+              `type`: `dropdown`,
+              `multi`: true
+            }
+          }
+        ],
+        `components`: [
+          {
+            `id`: `scatter_plot`,
+            `type`: `graph`,
+            `title`: `Sepal Dimensions by Species`,
+            `figure`: {
+              `x`: `sepal_length`,
+              `y`: `sepal_width`,
+              `color`: `species`,
+              `_target_`: `scatter`,
+              `data_frame`: `iris_data`,
+              `hover_data`: [
+                `petal_length`,
+                `petal_width`
+              ]
+            }
+          }
+        ]
+      }
+    ],
+    `theme`: `vizro_dark`,
+    `title`: `Iris Dashboard`
+  },
+  `data_infos`: `
+[
+    {
+        \"file_name\": \"iris_data\",
+        \"file_path_or_url\": \"https://raw.githubusercontent.com/plotly/datasets/master/iris-id.csv\",
+        \"file_location_type\": \"remote\",
+        \"read_function_string\": \"pd.read_csv\",
+    }
+]
+`
+}
+
+"""
+
+
+def convert_github_url_to_raw(path_or_url: str) -> str:
+    """Convert a GitHub URL to a raw URL if it's a GitHub URL, otherwise return the original path or URL."""
+    github_pattern = r"https?://(?:www\.)?github\.com/([^/]+)/([^/]+)/(?:blob|raw)/([^/]+)/(.+)"
+    github_match = re.match(github_pattern, path_or_url)
+
+    if github_match:
+        user, repo, branch, file_path = github_match.groups()
+        return f"https://raw.githubusercontent.com/{user}/{repo}/{branch}/{file_path}"
+
+    return path_or_url
+
+
+def load_dataframe_by_format(
+    path_or_url: Union[str, Path], mime_type: Optional[str] = None
+) -> tuple[pd.DataFrame, Literal["pd.read_csv", "pd.read_json", "pd.read_html", "pd.read_excel", "pd.read_parquet"]]:
+    """Load a dataframe based on file format determined by MIME type or file extension."""
+    file_path_str_lower = str(path_or_url).lower()
+
+    # Determine format
+    if mime_type == "text/csv" or file_path_str_lower.endswith(".csv"):
+        df = pd.read_csv(
+            path_or_url,
+            on_bad_lines="warn",
+            low_memory=False,
+        )
+        read_fn = "pd.read_csv"
+    elif mime_type == "application/json" or file_path_str_lower.endswith(".json"):
+        df = pd.read_json(path_or_url)
+        read_fn = "pd.read_json"
+    elif mime_type == "text/html" or file_path_str_lower.endswith((".html", ".htm")):
+        tables = pd.read_html(path_or_url)
+        if not tables:
+            raise ValueError("No HTML tables found in the provided file or URL")
+        df = tables[0]  # Get the first table by default
+        read_fn = "pd.read_html"
+    elif mime_type in [
+        "application/vnd.ms-excel",
+        "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+        "application/vnd.oasis.opendocument.spreadsheet",
+    ] or any(file_path_str_lower.endswith(ext) for ext in [".xls", ".xlsx", ".ods"]):
+        df = pd.read_excel(path_or_url)  # opens only sheet 0
+        read_fn = "pd.read_excel"
+    elif mime_type == "application/vnd.apache.parquet" or file_path_str_lower.endswith(
+        ".parquet"
+    ):  # mime type exists but I did not manage to ever extract it
+        df = pd.read_parquet(path_or_url)
+        read_fn = "pd.read_parquet"
+    else:
+        raise ValueError("Could not determine file format")
+
+    # Check if the result is a Series and convert to DataFrame if needed
+    if isinstance(df, pd.Series):
+        df = df.to_frame()
+
+    return df, read_fn
+
+
+def path_or_url_check(string: str) -> str:
+    """Check if a string is a link or a file path."""
+    if string.startswith(("http://", "https://", "www.")):
+        return "remote"
+
+    if Path(string).is_file():
+        return "local"
+
+    return "invalid"
+
+
+def get_dataframe_info(df: pd.DataFrame) -> DFInfo:
+    """Get the info of a DataFrame."""
+    buffer = io.StringIO()
+    df.info(buf=buffer)
+    info_string = buffer.getvalue()
+
+    # Sample only as many rows as exist in the dataframe, up to 5
+    sample_size = min(5, len(df)) if not df.empty else 0
+
+    return DFInfo(general_info=info_string, sample=df.sample(sample_size).to_dict() if sample_size > 0 else {})
+
+
+def create_pycafe_url(python_code: str) -> str:
+    """Create a PyCafe URL for a given Python code."""
+    # Create JSON object for py.cafe
+    json_object = {
+        "code": python_code,
+        "requirements": "vizro==0.1.38",
+        "files": [],
+    }
+
+    # Convert to compressed base64 URL
+    json_text = json.dumps(json_object)
+    compressed_json_text = gzip.compress(json_text.encode("utf8"))
+    base64_text = base64.b64encode(compressed_json_text).decode("utf8")
+    query = urlencode({"c": base64_text}, quote_via=quote)
+    pycafe_url = f"{PYCAFE_URL}/snippet/vizro/v1?{query}"
+
+    return pycafe_url
+
+
+def get_python_code_and_preview_link(
+    model_object: vm.VizroBaseModel, data_infos: list[DFMetaData]
+) -> VizroCodeAndPreviewLink:
+    """Get the Python code and preview link for a Vizro model object."""
+    # Get the Python code
+    python_code = model_object._to_python()
+
+    # Add imports after the first empty line
+    lines = python_code.splitlines()
+    for i, line in enumerate(lines):
+        if not line.strip():
+            # Found first empty line, insert imports here
+            imports_to_add = [
+                "from vizro import Vizro",
+                "import pandas as pd",
+                "from vizro.managers import data_manager",
+            ]
+            lines[i:i] = imports_to_add
+            break
+
+    python_code = "\n".join(lines)
+
+    # Prepare data loading code
+    data_loading_code = "\n".join(
+        f'data_manager["{info.file_name}"] = {info.read_function_string}("{info.file_path_or_url}")'
+        for info in data_infos
+    )
+
+    # Patterns to identify the data manager section
+    data_manager_start_marker = "####### Data Manager Settings #####"
+    data_manager_end_marker = "########### Model code ############"
+
+    # Replace everything between the markers with our data loading code
+    pattern = re.compile(f"{data_manager_start_marker}.*?{data_manager_end_marker}", re.DOTALL)
+    replacement = f"{data_manager_start_marker}\n{data_loading_code}\n\n{data_manager_end_marker}"
+    python_code = pattern.sub(replacement, python_code)
+
+    # Add final run line
+    python_code += "\n\nVizro().build(model).run()"
+
+    pycafe_url = create_pycafe_url(python_code)
+
+    return VizroCodeAndPreviewLink(python_code=python_code, pycafe_url=pycafe_url)