Optuna MCP Server

Instructions

Implementation Reference

• optuna_mcp/server.py:388-400 (handler)

  The main handler function for the 'add_trial' MCP tool. It is decorated with @mcp.tool(), registering it as a tool, and adds a single trial to the Optuna study by calling the _create_trial helper.

  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."

• optuna_mcp/server.py:35-69 (schema)

  Dataclass defining the structured input schema (TrialToAdd) for the 'add_trial' tool, specifying parameters, distributions, values, state, and attributes.

  @dataclass
  class TrialToAdd:
      """
      A trial to be added to an Optuna study.
  
      Attributes:
          params: The parameter values for the trial.
          distributions: The distributions used for the parameters.
              A key is the parameter name and a value is a distribution.
              The distribution is a dictionary that can be converted to a JSON string, e.g.,
              {
                  "name": "IntDistribution",
                  "attributes": {"step": null, "low": 1, "high": 9, "log": false}
              }.
              The name of the distribution must be one of the following:
              - FloatDistribution
              - IntDistribution
              - CategoricalDistribution
          values: The objective values for the trial, or None if not set.
              If the state is "COMPLETE", this must be set.
          state: The state of the trial.
              - "COMPLETE": The trial completed successfully.
              - "PRUNED": The trial was pruned.
              - "FAIL": The trial failed.
          user_attrs: User-defined attributes for the trial, or None if not set.
          system_attrs: System-defined attributes for the trial, or None if not set.
      """
  
      params: dict[str, typing.Any]
      distributions: dict[str, typing.Any]
      values: list[float] | None
      state: typing.Literal["COMPLETE", "PRUNED", "FAIL"]
      user_attrs: dict[str, typing.Any] | None
      system_attrs: dict[str, typing.Any] | None

• optuna_mcp/server.py:373-386 (helper)

  Helper function that converts a TrialToAdd dataclass instance into an Optuna FrozenTrial object, used by the add_trial handler.

  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,
      )