vizro-mcp 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

vizro_mcp/__init__.py

@@ -3,7 +3,7 @@ import sys
 
 from .server import mcp
 
-__version__ = "0.1.1"
+__version__ = "0.1.3"
 
 
 def main():

vizro_mcp/_schemas/__init__.py

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

vizro_mcp/_schemas/schemas.py

@@ -1,15 +1,14 @@
 """Schema defining pydantic models for usage in the MCP server."""
 
-from typing import Annotated, Any, Literal, Optional
+from typing import Annotated, Any, Optional
 
+import vizro.figures as vf
 import vizro.models as vm
-from pydantic import AfterValidator, BaseModel, Field, PrivateAttr, conlist
+from pydantic import AfterValidator, BaseModel, Field, PrivateAttr, ValidationInfo
 
-from vizro_mcp._utils import SAMPLE_DASHBOARD_CONFIG, DFMetaData
+from vizro_mcp._utils import DFMetaData
 
-# Constants used in chart validation
-CUSTOM_CHART_NAME = "custom_chart"
-ADDITIONAL_IMPORTS = [
+BASE_IMPORTS = [
     "import vizro.plotly.express as px",
     "import plotly.graph_objects as go",
     "import pandas as pd",
@@ -17,91 +16,6 @@ ADDITIONAL_IMPORTS = [
     "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.
@@ -110,15 +24,23 @@ class GraphEnhanced(vm.Graph):
 
     figure: dict[str, Any] = Field(
         description="""
+For simpler charts and without need for data manipulation, use this field:
 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).
-
+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".
+- you must refer to the dataframe by name, check file_name in the data_infos field ("data_frame": "<file_name>")
 - do not use a title if your Graph model already has a title.
+
+For more complex charts and those that require data manipulation, use the `custom_charts` field:
+- create the suitable number of custom charts and add them to the `custom_charts` field
+- refer here to the function signature you created
+- you must use the key: "_target_: "<custom_chart_name>"
+- you must refer to the dataframe by name, check file_name in the data_infos field ("data_frame": "<file_name>")
+- in general, DO NOT modify the background (with plot_bgcolor) or color sequences unless explicitly asked for
+- when creating hover templates, EXPLICITLY style them to work on light and dark mode
 """
     )
 
@@ -137,7 +59,22 @@ The only difference to the dash version is that:
     )
 
 
-###### Chart functionality - not sure if I should include this in the MCP server
+FIGURE_NAMESPACE_FUNCTION_DOCS = {func: vf.__dict__[func].__doc__ for func in vf.__all__}
+
+
+class FigureEnhanced(vm.Figure):
+    """Figure model that allows to use dynamic figure functions."""
+
+    figure: dict[str, Any] = Field(
+        description=f"""This is the figure function to be displayed.
+
+Only use arguments from the below mapping of _target_ to figure function documentation:
+
+{FIGURE_NAMESPACE_FUNCTION_DOCS}"""
+    )
+
+
+###### Chart functionality ######
 def _strip_markdown(code_string: str) -> str:
     """Remove any code block wrappers (markdown or triple quotes)."""
     wrappers = [("```python\n", "```"), ("```py\n", "```"), ("```\n", "```"), ('"""', '"""'), ("'''", "'''")]
@@ -150,22 +87,19 @@ def _strip_markdown(code_string: str) -> str:
     return code_string.strip()
 
 
-def _check_chart_code(v: str) -> str:
+def _check_chart_code(v: str, info: ValidationInfo) -> 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}("
+    func_def = f"def {info.data['chart_name']}("
     if func_def not in v:
-        raise ValueError(f"The chart code must be wrapped in a function named `{CUSTOM_CHART_NAME}`")
+        raise ValueError(f"The chart code must be wrapped in a function named `{info.data['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."""
-        )
+    if "(data_frame" not in first_line:
+        raise ValueError("""The chart code must accept as first argument `data_frame` which is a pandas DataFrame.""")
     return v
 
 
@@ -177,6 +111,12 @@ class ChartPlan(BaseModel):
         Describes the chart type that best reflects the user request.
         """,
     )
+    chart_name: str = Field(
+        description="""
+        The name of the chart function. Should be unique, concise and in snake_case.
+        """,
+        pattern=r"^[a-z][a-z0-9_]*$",
+    )
     imports: list[str] = Field(
         description="""
         List of import statements required to render the chart defined by the `chart_code` field. Ensure that every
@@ -192,20 +132,22 @@ class ChartPlan(BaseModel):
         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
+        1. Must be wrapped in a function that is named `chart_name`
+        2. Must accept as first argument 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.
+        5. DO NOT modify the background (with plot_bgcolor) or color sequences unless explicitly asked for
+        6. When creating hover templates, explicitly ensure that it works on light and dark mode
         """,
         ),
     ]
 
-    _additional_vizro_imports: list[str] = PrivateAttr(ADDITIONAL_IMPORTS)
+    _base_chart_imports: list[str] = PrivateAttr(BASE_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 = list(dict.fromkeys(self.imports + self._base_chart_imports))  # remove duplicates
+        if vizro:
             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]
@@ -214,9 +156,9 @@ class ChartPlan(BaseModel):
     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}")
+            chart_code = chart_code.replace(f"def {self.chart_name}", f"@capture('graph')\ndef {self.chart_name}")
         if chart_name is not None:
-            chart_code = chart_code.replace(f"def {CUSTOM_CHART_NAME}", f"def {chart_name}")
+            chart_code = chart_code.replace(f"def {self.chart_name}", f"def {chart_name}")
         return chart_code
 
     def get_dashboard_template(self, data_info: DFMetaData) -> str:
@@ -232,14 +174,14 @@ class ChartPlan(BaseModel):
         imports = self.get_imports(vizro=True)
 
         # Add the Vizro-specific imports if not present
-        additional_imports = [
+        additional_dashboard_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")))
+        all_imports = list(dict.fromkeys(additional_dashboard_imports + imports.split("\n")))
 
         dashboard_template = f"""
 {chr(10).join(imp for imp in all_imports if imp)}
@@ -259,7 +201,7 @@ dashboard = vm.Dashboard(
             components=[
                 vm.Graph(
                     id="{self.chart_type}_graph",
-                    figure={CUSTOM_CHART_NAME}("{data_info.file_name}"),
+                    figure={self.chart_name}("{data_info.file_name}"),
                 )
             ],
         )
@@ -274,24 +216,15 @@ 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
-
+if __name__ == "__main__":
+    plan = ChartPlan(
+        chart_type="scatter",
+        chart_name="scatter",
+        imports=["import pandas as pd", "import plotly.express as px"],
+        chart_code="""
+def scatter(data_frame):
+    return px.scatter(data_frame, x="sepal_length", y="sepal_width")
+        """,
+    )
 
-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
+    # print(plan.get_chart_code(chart_name="poo", vizro=True))

vizro_mcp/_utils/__init__.py

@@ -1,11 +1,22 @@
-from .utils import (
+from .configs import (
     GAPMINDER,
     IRIS,
     SAMPLE_DASHBOARD_CONFIG,
     STOCKS,
     TIPS,
+)
+from .prompts import (
+    CHART_INSTRUCTIONS,
+    LAYOUT_INSTRUCTIONS,
+    get_chart_prompt,
+    get_dashboard_instructions,
+    get_dashboard_prompt,
+    get_starter_dashboard_prompt,
+)
+from .utils import (
     DFInfo,
     DFMetaData,
+    NoDefsGenerateJsonSchema,
     VizroCodeAndPreviewLink,
     convert_github_url_to_raw,
     create_pycafe_url,
@@ -15,19 +26,29 @@ from .utils import (
     path_or_url_check,
 )
 
-__all__ = [
+__all__ = [  # noqa: RUF022
+    # Classes
+    "DFInfo",
+    "DFMetaData",
+    "NoDefsGenerateJsonSchema",
+    "VizroCodeAndPreviewLink",
+    # Constants
+    "CHART_INSTRUCTIONS",
+    "LAYOUT_INSTRUCTIONS",
     "GAPMINDER",
     "IRIS",
     "SAMPLE_DASHBOARD_CONFIG",
     "STOCKS",
     "TIPS",
-    "DFInfo",
-    "DFMetaData",
-    "VizroCodeAndPreviewLink",
+    # Functions
     "convert_github_url_to_raw",
     "create_pycafe_url",
     "get_dataframe_info",
+    "get_dashboard_instructions",
+    "get_dashboard_prompt",
+    "get_chart_prompt",
     "get_python_code_and_preview_link",
+    "get_starter_dashboard_prompt",
     "load_dataframe_by_format",
     "path_or_url_check",
 ]

vizro_mcp/_utils/configs.py

@@ -0,0 +1,142 @@
+"""Pre-set configs for the Vizro MCP."""
+
+from dataclasses import dataclass
+from typing import Any, Literal, Optional
+
+
+@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 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