Optuna MCP Server

Instructions

Implementation Reference

• optuna_mcp/server.py:281-295 (handler)

  The main handler function for the 'get_metric_names' tool. It checks if a study exists and returns a StudyResponse containing the study's metric names.

  @mcp.tool(structured_output=True)
  def get_metric_names() -> StudyResponse:
      """Get metric_names"""
      if mcp.study is None:
          raise McpError(
              ErrorData(
                  code=INTERNAL_ERROR,
                  message="No study has been created. Please create a study first.",
              )
          )
  
      return StudyResponse(
          study_name=mcp.study.study_name,
          metric_names=mcp.study.metric_names,
      )

• optuna_mcp/server.py:71-82 (schema)

  Pydantic schema for the structured output of get_metric_names, defining fields including metric_names.

  class StudyResponse(BaseModel):
      study_name: str
      sampler_name: SamplerName | None = Field(
          default=None, description="The name of the sampler used in the study."
      )
      directions: list[DirectionName] | None = Field(
          default=None, description="The optimization directions for each objective."
      )
      metric_names: list[str] | None = Field(
          default=None, description="The metric names for each objective."
      )

• optuna_mcp/server.py:100-669 (registration)

  Function that registers all MCP tools, including get_metric_names via its decorator.

  def register_tools(mcp: OptunaMCP) -> OptunaMCP:
      @mcp.tool(structured_output=True)
      def create_study(
          study_name: str,
          directions: list[DirectionName] | None = None,
      ) -> StudyResponse:
          """Create a new Optuna study with the given study_name and directions.
  
          If the study already exists, it will be simply loaded.
          """
  
          mcp.study = optuna.create_study(
              study_name=study_name,
              storage=mcp.storage,
              load_if_exists=True,
              directions=directions,
          )
          if mcp.storage is None:
              mcp.storage = mcp.study._storage
  
          return StudyResponse(study_name=study_name)
  
      @mcp.tool(structured_output=True)
      def get_all_study_names() -> list[StudyResponse]:
          """Get all study names from the storage."""
          storage: str | optuna.storages.BaseStorage | None = None
          if mcp.study is not None:
              storage = mcp.study._storage
          elif mcp.storage is not None:
              storage = mcp.storage
          else:
              raise McpError(ErrorData(code=INTERNAL_ERROR, message="No storage specified."))
  
          study_names = optuna.get_all_study_names(storage)
          return [StudyResponse(study_name=name) for name in study_names]
  
      @mcp.tool(structured_output=True)
      def ask(search_space: dict) -> TrialResponse:
          """Suggest new parameters using Optuna
  
          search_space must be a string that can be evaluated to a dictionary to specify Optuna's distributions.
  
          Example:
              {"x": {"name": "FloatDistribution", "attributes": {"step": null, "low": -10.0, "high": 10.0, "log": false}}}
          """
          try:
              distributions = {
                  name: optuna.distributions.json_to_distribution(json.dumps(dist))
                  for name, dist in search_space.items()
              }
          except Exception as e:
              raise McpError(ErrorData(code=INTERNAL_ERROR, message=f"Error: {e}")) from e
  
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          trial = mcp.study.ask(fixed_distributions=distributions)
  
          return TrialResponse(
              trial_number=trial.number,
              params=trial.params,
          )
  
      @mcp.tool(structured_output=True)
      def tell(trial_number: int, values: float | list[float]) -> TrialResponse:
          """Report the result of a trial"""
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          mcp.study.tell(
              trial=trial_number,
              values=values,
              state=optuna.trial.TrialState.COMPLETE,
              skip_if_finished=True,
          )
          return TrialResponse(
              trial_number=trial_number,
              values=[values] if isinstance(values, float) else values,
          )
  
      @mcp.tool(structured_output=True)
      def set_sampler(
          name: SamplerName,
      ) -> StudyResponse:
          """Set the sampler for the study.
          The sampler must be one of the following:
          - TPESampler
          - NSGAIISampler
          - RandomSampler
          - GPSampler
  
          The default sampler for single-objective optimization is TPESampler.
          The default sampler for multi-objective optimization is NSGAIISampler.
          GPSampler is a Gaussian process-based sampler suitable for low-dimensional numerical optimization problems.
          """
          sampler = getattr(optuna.samplers, name)()
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          mcp.study.sampler = sampler
          return StudyResponse(
              study_name=mcp.study.study_name,
              sampler_name=name,
          )
  
      @mcp.tool(structured_output=True)
      def set_trial_user_attr(trial_number: int, key: str, value: typing.Any) -> str:
          """Set user attributes for a trial"""
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          storage = mcp.study._storage
          trial_id = storage.get_trial_id_from_study_id_trial_number(
              mcp.study._study_id, trial_number
          )
          storage.set_trial_user_attr(trial_id, key, value)
          return f"User attribute {key} set to {json.dumps(value)} for trial {trial_number}"
  
      @mcp.tool(structured_output=True)
      def get_trial_user_attrs(trial_number: int) -> TrialResponse:
          """Get user attributes in a trial"""
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          storage = mcp.study._storage
          trial_id = storage.get_trial_id_from_study_id_trial_number(
              mcp.study._study_id, trial_number
          )
          trial = storage.get_trial(trial_id)
          return TrialResponse(
              trial_number=trial_number,
              user_attrs=trial.user_attrs,
          )
  
      @mcp.tool(structured_output=True)
      def set_metric_names(metric_names: list[str]) -> StudyResponse:
          """Set metric_names. metric_names are labels used to distinguish what each objective value is.
  
          Args:
              metric_names:
                  The list of metric name for each objective value.
                  The length of metric_names list must be the same with the number of objectives.
          """
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
          mcp.study.set_metric_names(metric_names)
          return StudyResponse(
              study_name=mcp.study.study_name,
              metric_names=metric_names,
          )
  
      @mcp.tool(structured_output=True)
      def get_metric_names() -> StudyResponse:
          """Get metric_names"""
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          return StudyResponse(
              study_name=mcp.study.study_name,
              metric_names=mcp.study.metric_names,
          )
  
      @mcp.tool(structured_output=True)
      def get_directions() -> StudyResponse:
          """Get the directions of the study."""
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          directions = [d.name.lower() for d in mcp.study.directions]
          return StudyResponse(
              study_name=mcp.study.study_name,
              directions=typing.cast(list[DirectionName], directions),
          )
  
      @mcp.tool(structured_output=False)
      def get_trials() -> str:
          """Get all trials in a CSV format"""
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          csv_string = mcp.study.trials_dataframe().to_csv()
          return f"Trials: \n{csv_string}"
  
      @mcp.tool(structured_output=True)
      def best_trial() -> TrialResponse:
          """Get the best trial
  
          This feature can only be used for single-objective optimization. If your study is multi-objective, use best_trials instead.
          """
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          trial = mcp.study.best_trial
          return TrialResponse(
              trial_number=trial.number,
              params=trial.params,
              values=trial.values,
              user_attrs=trial.user_attrs,
              system_attrs=trial.system_attrs,
          )
  
      @mcp.tool(structured_output=True)
      def best_trials() -> list[TrialResponse]:
          """Return trials located at the Pareto front in the study."""
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          return [
              TrialResponse(
                  trial_number=trial.number,
                  params=trial.params,
                  values=trial.values,
                  user_attrs=trial.user_attrs,
                  system_attrs=trial.system_attrs,
              )
              for trial in mcp.study.best_trials
          ]
  
      def _create_trial(trial: TrialToAdd) -> optuna.trial.FrozenTrial:
          """Create a trial from the given parameters."""
          return optuna.trial.create_trial(
              params=trial.params,
              distributions={
                  k: optuna.distributions.json_to_distribution(json.dumps(d))
                  for k, d in trial.distributions.items()
              },
              values=trial.values,
              state=optuna.trial.TrialState[trial.state],
              user_attrs=trial.user_attrs,
              system_attrs=trial.system_attrs,
          )
  
      @mcp.tool(structured_output=True)
      def add_trial(trial: TrialToAdd) -> str:
          """Add a trial to the study."""
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          mcp.study.add_trial(_create_trial(trial))
          return "Trial was added."
  
      @mcp.tool(structured_output=True)
      def add_trials(trials: list[TrialToAdd]) -> str:
          """Add multiple trials to the study."""
          frozen_trials = [_create_trial(trial) for trial in trials]
          if mcp.study is None:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          mcp.study.add_trials(frozen_trials)
          return f"{len(trials)} trials were added."
  
      @mcp.tool()
      def plot_optimization_history(
          target: int | None = None,
          target_name: str = "Objective Value",
      ) -> Image:
          """Return the optimization history plot as an image.
  
          Args:
              target:
                  An index to specify the value to display. To plot nth objective value, set this to n.
                  Note that this is 0-indexed, i.e., to plot the first objective value, set this to 0.
                  For single-objective optimization, None (auto) is recommended.
                  For multi-objective optimization, this must be specified.
              target_name:
                  Target's name to display on the axis label and the legend.
          """
          fig = optuna.visualization.plot_optimization_history(
              mcp.study,
              target=(lambda t: t.values[target]) if target is not None else None,
              target_name=target_name,
          )
          return Image(data=plotly.io.to_image(fig), format="png")
  
      @mcp.tool()
      def plot_hypervolume_history(
          reference_point: list[float],
      ) -> Image:
          """Return the hypervolume history plot as an image.
  
          Args:
              reference_point:
                  A list of reference points to calculate the hypervolume.
          """
          fig = optuna.visualization.plot_hypervolume_history(
              mcp.study,
              reference_point=reference_point,
          )
          return Image(data=plotly.io.to_image(fig), format="png")
  
      @mcp.tool()
      def plot_pareto_front(
          target_names: list[str] | None = None,
          include_dominated_trials: bool = True,
          targets: list[int] | None = None,
      ) -> Image:
          """Return the Pareto front plot as an image for multi-objective optimization.
  
          Args:
              target_names:
                  Objective name list used as the axis titles. If :obj:`None` is specified,
                  "Objective {objective_index}" is used instead. If ``targets`` is specified
                  for a study that does not contain any completed trial,
                  ``target_name`` must be specified.
              include_dominated_trials:
                  A flag to include all dominated trial's objective values.
              targets:
                  A list of indices to specify the objective values to display.
                  Note that this is 0-indexed, i.e., to plot the first and second objective value, set this to [0, 1].
                  If the number of objectives is neither 2 nor 3, ``targets`` must be specified.
                  By default, all objectives are displayed.
          """
          fig = optuna.visualization.plot_pareto_front(
              mcp.study,
              target_names=target_names,
              include_dominated_trials=include_dominated_trials,
              targets=targets,
          )
          return Image(data=plotly.io.to_image(fig), format="png")
  
      @mcp.tool()
      def plot_contour(
          params: list[str] | None = None,
          target: int = 0,
          target_name: str = "Objective Value",
      ) -> Image:
          """Return the contour plot as an image.
  
          Args:
              params:
                  Parameter list to visualize. The default is all parameters.
              target:
                  An index to specify the value to display. To plot nth objective value, set this to n.
                  Note that this is 0-indexed, i.e., to plot the first objective value, set this to 0.
              target_name:
                  Target’s name to display on the color bar.
          """
          fig = optuna.visualization.plot_contour(
              mcp.study, params=params, target=lambda t: t.values[target], target_name=target_name
          )
          return Image(data=plotly.io.to_image(fig), format="png")
  
      @mcp.tool()
      def plot_parallel_coordinate(
          params: list[str] | None = None,
          target: int = 0,
          target_name: str = "Objective Value",
      ) -> Image:
          """Return the parallel coordinate plot as an image.
  
          Args:
              params:
                  Parameter list to visualize. The default is all parameters.
              target:
                  An index to specify the value to display. To plot nth objective value, set this to n.
                  Note that this is 0-indexed, i.e., to plot the first objective value, set this to 0.
              target_name:
                  Target’s name to display on the axis label and the legend.
          """
          fig = optuna.visualization.plot_parallel_coordinate(
              mcp.study,
              params=params,
              target=lambda t: t.values[target],
              target_name=target_name,
          )
          return Image(data=plotly.io.to_image(fig), format="png")
  
      @mcp.tool()
      def plot_slice(
          params: list[str] | None = None,
          target: int = 0,
          target_name: str = "Objective Value",
      ) -> Image:
          """Return the slice plot as an image.
  
          Args:
              params:
                  Parameter list to visualize. The default is all parameters.
              target:
                  An index to specify the value to display. To plot nth objective value, set this to n.
                  Note that this is 0-indexed, i.e., to plot the first objective value, set this to 0.
              target_name:
                  Target’s name to display on the axis label.
          """
          fig = optuna.visualization.plot_slice(
              mcp.study,
              params=params,
              target=lambda t: t.values[target],
              target_name=target_name,
          )
          return Image(data=plotly.io.to_image(fig), format="png")
  
      @mcp.tool()
      def plot_param_importances(
          params: list[str] | None = None,
          target: int | None = None,
          target_name: str = "Objective Value",
      ) -> Image:
          """Return the parameter importances plot as an image.
  
          Args:
              params:
                  Parameter list to visualize. The default is all parameters.
              target:
                  An index to specify the value to display. To plot nth objective value, set this to n.
                  Note that this is 0-indexed, i.e., to plot the first objective value, set this to 0.
                  By default, all objective will be plotted by setting target to None.
              target_name:
                  Target’s name to display on the legend.
          """
          evaluator = optuna.importance.PedAnovaImportanceEvaluator()
          fig = optuna.visualization.plot_param_importances(
              mcp.study,
              evaluator=evaluator,
              params=params,
              target=(lambda t: t.values[target]) if target is not None else None,
              target_name=target_name,
          )
          return Image(data=plotly.io.to_image(fig), format="png")
  
      @mcp.tool()
      def plot_edf(
          target: int = 0,
          target_name: str = "Objective Value",
      ) -> Image:
          """Return the EDF plot as an image.
  
          Args:
              target:
                  An index to specify the value to display. To plot nth objective value, set this to n.
                  Note that this is 0-indexed, i.e., to plot the first objective value, set this to 0.
              target_name:
                  Target’s name to display on the axis label.
          """
          fig = optuna.visualization.plot_edf(
              mcp.study,
              target=lambda t: t.values[target],
              target_name=target_name,
          )
          return Image(data=plotly.io.to_image(fig), format="png")
  
      @mcp.tool()
      def plot_timeline() -> Image:
          """Return the timeline plot as an image."""
          fig = optuna.visualization.plot_timeline(mcp.study)
          return Image(data=plotly.io.to_image(fig), format="png")
  
      @mcp.tool()
      def plot_rank(
          params: list[str] | None = None,
          target: int = 0,
          target_name: str = "Objective Value",
      ) -> Image:
          """Return the rank plot as an image.
  
          Args:
              params:
                  Parameter list to visualize. The default is all parameters.
              target:
                  An index to specify the value to display. To plot nth objective value, set this to n.
                  Note that this is 0-indexed, i.e., to plot the first objective value, set this to 0.
              target_name:
                  Target’s name to display on the color bar.
          """
          fig = optuna.visualization.plot_rank(mcp.study)
          return Image(data=plotly.io.to_image(fig), format="png")
  
      @mcp.tool(structured_output=True)
      def launch_optuna_dashboard(port: int = 58080) -> str:
          """Launch the Optuna dashboard"""
          storage: str | optuna.storages.BaseStorage | None = None
          if mcp.dashboard_thread_port is not None:
              return f"Optuna dashboard is already running. Open http://127.0.0.1:{mcp.dashboard_thread_port[1]}."
  
          if mcp.study is not None:
              storage = mcp.study._storage
          elif mcp.storage is not None:
              storage = mcp.storage
          else:
              raise McpError(
                  ErrorData(
                      code=INTERNAL_ERROR,
                      message="No study has been created. Please create a study first.",
                  )
              )
  
          def runner(storage: optuna.storages.BaseStorage | str, port: int) -> None:
              try:
                  optuna_dashboard.run_server(storage=storage, host="127.0.0.1", port=port)
              except Exception as e:
                  print(f"Error starting the dashboard: {e}", file=sys.stderr)
                  sys.exit(1)
  
          # TODO(y0z): Consider better implementation
          thread = threading.Thread(
              target=runner,
              args=(storage, port),
              daemon=True,
          )
          thread.start()
          mcp.dashboard_thread_port = (thread, port)
  
          return f"Optuna dashboard is running at http://127.0.0.1:{port}"
  
      return mcp