Zaturn

Instructions

Implementation Reference

• zaturn/tools/visualizations.py:270-308 (handler)

  The main handler function for the 'density_heatmap' MCP tool. It takes data source ID, SQL query, column names for X/Y axes, optional bin counts, executes the query, generates a density heatmap using plotly.express.density_heatmap, converts the figure to a base64-encoded PNG ImageContent object for MCP response, or returns an error string.

  def density_heatmap(self,
      source_id: Annotated[
          str, Field(description='The data source to run the query on')
      ],  
      query: Annotated[
          str, Field(description='SQL query to run on the data source')
      ],
      column_x: Annotated[
          str, Field(description='Column name from SQL result to use on the X-axis')
      ],
      column_y: Annotated[
          str, Field(description='Column name from SQL result to use on the Y-axis')
      ],
      nbins_x: Annotated[
          int | None, Field(description='Optional; number of bins to use on the X-axis')
      ] = None,
      nbins_y: Annotated[
          int | None, Field(description='Optional; number of bins to use on the Y-axis')
      ] = None,
  ) -> str | ImageContent:
      """
      Run query against specified source and make a 2d-histogram aka. Density Heatmap using result
      For both csv and parquet sources, use DuckDB SQL syntax
      Use 'CSV' as the table name in the SQL query for csv sources.
      Use 'PARQUET' as the table name in the SQL query for parquet sources.
  
      This will return an image of the plot
      """
  
      try:
          df = self._get_df_from_source(source_id, query)
          fig = px.density_heatmap(df, x=column_x, y=column_y, nbinsx=nbins_x, nbinsy=nbins_y)
          fig.update_xaxes(autotickangles=[0, 45, 60, 90])
          fig.update_yaxes(autotickangles=[0, 45, 60, 90])
  
          return _fig_to_image(fig)
      except Exception as e:
          return str(e)

• zaturn/mcp/__init__.py:92-93 (registration)

  Registers all tools from ZaturnTools instance (which includes density_heatmap) to the FastMCP server by converting each tool function to a Tool object and calling add_tool.

  for tool_function in zaturn_tools.tools:
      zaturn_mcp.add_tool(Tool.from_function(tool_function))

• zaturn/tools/visualizations.py:29-40 (registration)

  Adds the density_heatmap method to the self.tools list in Visualizations class __init__, making it available for collection into higher-level tool sets like ZaturnTools.

  self.tools = [
      self.scatter_plot,
      self.line_plot,
      self.histogram,
      self.strip_plot,
      self.box_plot,
      self.bar_plot,
  
      self.density_heatmap,
      self.polar_scatter,
      self.polar_line,
  ]

• zaturn/tools/__init__.py:7-9 (registration)

  Includes all tools from Visualizations class (including density_heatmap) into the ZaturnTools.tools list by unpacking visualizations.Visualizations(data_sources).tools.

  self.tools = [
      *core.Core(data_sources).tools,
      *visualizations.Visualizations(data_sources).tools,

• zaturn/tools/visualizations.py:13-22 (helper)

  Utility function to convert a Plotly figure to an MCP-compatible ImageContent object (base64 PNG), used by density_heatmap and all other visualization tools.

  def _fig_to_image(fig):
      fig_encoded = b64encode(fig.to_image(format='png')).decode()
      img_b64 = "data:image/png;base64," + fig_encoded
      
      return ImageContent(
          type = 'image',
          data = fig_encoded,
          mimeType = 'image/png',
          annotations = None,
      )

• zaturn/tools/visualizations.py:43-49 (helper)

  Utility method to retrieve data source by ID, execute the provided SQL query using query_utils, and return Pandas DataFrame, shared by density_heatmap and other viz tools.

  def _get_df_from_source(self, source_id, query):
      source = self.data_sources.get(source_id)
      if not source:
          raise Exception(f"Source {source_id} Not Found")
              
      return query_utils.execute_query(source, query)