Skip to content

Commit 95e79b1

Browse files
sawa3030sawa3030
committed
add error
1 parent 94eb06b commit 95e79b1

File tree

2 files changed

+149
-45
lines changed

2 files changed

+149
-45
lines changed

‎optuna_mcp/server.py‎

Lines changed: 137 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@
77

88
from mcp.server.fastmcp import FastMCP
99
from mcp.server.fastmcp import Image
10+
from mcp.types import CallToolResult
11+
from mcp.types import TextContent
1012
import optuna
1113
import optuna_dashboard
1214
import plotly
@@ -110,21 +112,23 @@ def create_study(
110112
return StudyResponse(study_name=study_name)
111113

112114
@mcp.tool(structured_output=True)
113-
def get_all_study_names() -> list[StudyResponse] | str:
115+
def get_all_study_names() -> list[StudyResponse] | CallToolResult:
114116
"""Get all study names from the storage."""
115117
storage: str | optuna.storages.BaseStorage | None = None
116118
if mcp.study is not None:
117119
storage = mcp.study._storage
118120
elif mcp.storage is not None:
119121
storage = mcp.storage
120122
else:
121-
return "No storage specified."
123+
return CallToolResult(
124+
content=[TextContent(type="text", text="No storage specified.")], isError=True
125+
)
122126

123127
study_names = optuna.get_all_study_names(storage)
124128
return [StudyResponse(study_name=name) for name in study_names]
125129

126130
@mcp.tool(structured_output=True)
127-
def ask(search_space: dict) -> TrialResponse | str:
131+
def ask(search_space: dict) -> TrialResponse | CallToolResult:
128132
"""Suggest new parameters using Optuna
129133
130134
search_space must be a string that can be evaluated to a dictionary to specify Optuna's distributions.
@@ -138,10 +142,19 @@ def ask(search_space: dict) -> TrialResponse | str:
138142
for name, dist in search_space.items()
139143
}
140144
except Exception as e:
141-
return f"Error: {e}"
145+
return CallToolResult(
146+
content=[TextContent(type="text", text=f"Error: {e}")], isError=True
147+
)
142148

143149
if mcp.study is None:
144-
raise ValueError("No study has been created. Please create a study first.")
150+
return CallToolResult(
151+
content=[
152+
TextContent(
153+
type="text", text="No study has been created. Please create a study first."
154+
)
155+
],
156+
isError=True,
157+
)
145158

146159
trial = mcp.study.ask(fixed_distributions=distributions)
147160

@@ -151,10 +164,17 @@ def ask(search_space: dict) -> TrialResponse | str:
151164
)
152165

153166
@mcp.tool(structured_output=True)
154-
def tell(trial_number: int, values: float | list[float]) -> TrialResponse:
167+
def tell(trial_number: int, values: float | list[float]) -> TrialResponse | CallToolResult:
155168
"""Report the result of a trial"""
156169
if mcp.study is None:
157-
raise ValueError("No study has been created. Please create a study first.")
170+
return CallToolResult(
171+
content=[
172+
TextContent(
173+
type="text", text="No study has been created. Please create a study first."
174+
)
175+
],
176+
isError=True,
177+
)
158178

159179
mcp.study.tell(
160180
trial=trial_number,
@@ -170,7 +190,7 @@ def tell(trial_number: int, values: float | list[float]) -> TrialResponse:
170190
@mcp.tool(structured_output=True)
171191
def set_sampler(
172192
name: typing.Literal["TPESampler", "NSGAIISampler", "RandomSampler", "GPSampler"],
173-
) -> StudyResponse:
193+
) -> StudyResponse | CallToolResult:
174194
"""Set the sampler for the study.
175195
The sampler must be one of the following:
176196
- TPESampler
@@ -184,18 +204,34 @@ def set_sampler(
184204
"""
185205
sampler = getattr(optuna.samplers, name)()
186206
if mcp.study is None:
187-
raise ValueError("No study has been created. Please create a study first.")
207+
return CallToolResult(
208+
content=[
209+
TextContent(
210+
type="text", text="No study has been created. Please create a study first."
211+
)
212+
],
213+
isError=True,
214+
)
188215
mcp.study.sampler = sampler
189216
return StudyResponse(
190217
study_name=mcp.study.study_name,
191218
sampler_name=name,
192219
)
193220

194221
@mcp.tool(structured_output=True)
195-
def set_trial_user_attr(trial_number: int, key: str, value: typing.Any) -> str:
222+
def set_trial_user_attr(
223+
trial_number: int, key: str, value: typing.Any
224+
) -> str | CallToolResult:
196225
"""Set user attributes for a trial"""
197226
if mcp.study is None:
198-
raise ValueError("No study has been created. Please create a study first.")
227+
return CallToolResult(
228+
content=[
229+
TextContent(
230+
type="text", text="No study has been created. Please create a study first."
231+
)
232+
],
233+
isError=True,
234+
)
199235

200236
storage = mcp.study._storage
201237
trial_id = storage.get_trial_id_from_study_id_trial_number(
@@ -205,10 +241,17 @@ def set_trial_user_attr(trial_number: int, key: str, value: typing.Any) -> str:
205241
return f"User attribute {key} set to {json.dumps(value)} for trial {trial_number}"
206242

207243
@mcp.tool(structured_output=True)
208-
def get_trial_user_attrs(trial_number: int) -> TrialResponse:
244+
def get_trial_user_attrs(trial_number: int) -> TrialResponse | CallToolResult:
209245
"""Get user attributes in a trial"""
210246
if mcp.study is None:
211-
raise ValueError("No study has been created. Please create a study first.")
247+
return CallToolResult(
248+
content=[
249+
TextContent(
250+
type="text", text="No study has been created. Please create a study first."
251+
)
252+
],
253+
isError=True,
254+
)
212255
storage = mcp.study._storage
213256
trial_id = storage.get_trial_id_from_study_id_trial_number(
214257
mcp.study._study_id, trial_number
@@ -220,7 +263,7 @@ def get_trial_user_attrs(trial_number: int) -> TrialResponse:
220263
)
221264

222265
@mcp.tool(structured_output=True)
223-
def set_metric_names(metric_names: list[str]) -> str:
266+
def set_metric_names(metric_names: list[str]) -> str | CallToolResult:
224267
"""Set metric_names. metric_names are labels used to distinguish what each objective value is.
225268
226269
Args:
@@ -229,44 +272,79 @@ def set_metric_names(metric_names: list[str]) -> str:
229272
The length of metric_names list must be the same with the number of objectives.
230273
"""
231274
if mcp.study is None:
232-
raise ValueError("No study has been created. Please create a study first.")
275+
return CallToolResult(
276+
content=[
277+
TextContent(
278+
type="text", text="No study has been created. Please create a study first."
279+
)
280+
],
281+
isError=True,
282+
)
233283
mcp.study.set_metric_names(metric_names)
234284
return f"metric_names set to {json.dumps(metric_names)}"
235285

236286
@mcp.tool(structured_output=True)
237-
def get_metric_names() -> str:
287+
def get_metric_names() -> str | CallToolResult:
238288
"""Get metric_names"""
239289
if mcp.study is None:
240-
raise ValueError("No study has been created. Please create a study first.")
290+
return CallToolResult(
291+
content=[
292+
TextContent(
293+
type="text", text="No study has been created. Please create a study first."
294+
)
295+
],
296+
isError=True,
297+
)
241298
return f"Metric names: {json.dumps(mcp.study.metric_names)}"
242299

243300
@mcp.tool(structured_output=True)
244-
def get_directions() -> StudyResponse:
301+
def get_directions() -> StudyResponse | CallToolResult:
245302
"""Get the directions of the study."""
246303
if mcp.study is None:
247-
raise ValueError("No study has been created. Please create a study first.")
304+
return CallToolResult(
305+
content=[
306+
TextContent(
307+
type="text", text="No study has been created. Please create a study first."
308+
)
309+
],
310+
isError=True,
311+
)
248312
directions = [d.name.lower() for d in mcp.study.directions]
249313
return StudyResponse(
250314
study_name=mcp.study.study_name,
251315
directions=directions,
252316
)
253317

254318
@mcp.tool(structured_output=True)
255-
def get_trials() -> str:
319+
def get_trials() -> str | CallToolResult:
256320
"""Get all trials in a CSV format"""
257321
if mcp.study is None:
258-
raise ValueError("No study has been created. Please create a study first.")
322+
return CallToolResult(
323+
content=[
324+
TextContent(
325+
type="text", text="No study has been created. Please create a study first."
326+
)
327+
],
328+
isError=True,
329+
)
259330
csv_string = mcp.study.trials_dataframe().to_csv()
260331
return f"Trials: \n{csv_string}"
261332

262333
@mcp.tool(structured_output=True)
263-
def best_trial() -> TrialResponse:
334+
def best_trial() -> TrialResponse | CallToolResult:
264335
"""Get the best trial
265336
266337
This feature can only be used for single-objective optimization. If your study is multi-objective, use best_trials instead.
267338
"""
268339
if mcp.study is None:
269-
raise ValueError("No study has been created. Please create a study first.")
340+
return CallToolResult(
341+
content=[
342+
TextContent(
343+
type="text", text="No study has been created. Please create a study first."
344+
)
345+
],
346+
isError=True,
347+
)
270348

271349
trial = mcp.study.best_trial
272350
return TrialResponse(
@@ -278,10 +356,17 @@ def best_trial() -> TrialResponse:
278356
)
279357

280358
@mcp.tool(structured_output=True)
281-
def best_trials() -> list[TrialResponse]:
359+
def best_trials() -> list[TrialResponse] | CallToolResult:
282360
"""Return trials located at the Pareto front in the study."""
283361
if mcp.study is None:
284-
raise ValueError("No study has been created. Please create a study first.")
362+
return CallToolResult(
363+
content=[
364+
TextContent(
365+
type="text", text="No study has been created. Please create a study first."
366+
)
367+
],
368+
isError=True,
369+
)
285370
return [
286371
TrialResponse(
287372
trial_number=trial.number,
@@ -308,19 +393,33 @@ def _create_trial(trial: TrialToAdd) -> optuna.trial.FrozenTrial:
308393
)
309394

310395
@mcp.tool(structured_output=True)
311-
def add_trial(trial: TrialToAdd) -> str:
396+
def add_trial(trial: TrialToAdd) -> str | CallToolResult:
312397
"""Add a trial to the study."""
313398
if mcp.study is None:
314-
raise ValueError("No study has been created. Please create a study first.")
399+
return CallToolResult(
400+
content=[
401+
TextContent(
402+
type="text", text="No study has been created. Please create a study first."
403+
)
404+
],
405+
isError=True,
406+
)
315407
mcp.study.add_trial(_create_trial(trial))
316408
return "Trial was added."
317409

318410
@mcp.tool(structured_output=True)
319-
def add_trials(trials: list[TrialToAdd]) -> str:
411+
def add_trials(trials: list[TrialToAdd]) -> str | CallToolResult:
320412
"""Add multiple trials to the study."""
321413
frozen_trials = [_create_trial(trial) for trial in trials]
322414
if mcp.study is None:
323-
raise ValueError("No study has been created. Please create a study first.")
415+
return CallToolResult(
416+
content=[
417+
TextContent(
418+
type="text", text="No study has been created. Please create a study first."
419+
)
420+
],
421+
isError=True,
422+
)
324423
mcp.study.add_trials(frozen_trials)
325424
return f"{len(trials)} trials were added."
326425

@@ -541,7 +640,7 @@ def plot_rank(
541640
return Image(data=plotly.io.to_image(fig), format="png")
542641

543642
@mcp.tool(structured_output=True)
544-
def launch_optuna_dashboard(port: int = 58080) -> str:
643+
def launch_optuna_dashboard(port: int = 58080) -> str | CallToolResult:
545644
"""Launch the Optuna dashboard"""
546645
storage: str | optuna.storages.BaseStorage | None = None
547646
if mcp.dashboard_thread_port is not None:
@@ -552,9 +651,14 @@ def launch_optuna_dashboard(port: int = 58080) -> str:
552651
elif mcp.storage is not None:
553652
storage = mcp.storage
554653
else:
555-
raise ValueError(
556-
"No study has been created or no storage URL has been provided."
557-
"Please create a study or provide a storage URL directly."
654+
return CallToolResult(
655+
content=[
656+
TextContent(
657+
type="text",
658+
text="No study has been created or no storage URL has been provided. Please create a study or provide a storage URL directly.",
659+
)
660+
],
661+
isError=True,
558662
)
559663

560664
def runner(storage: optuna.storages.BaseStorage | str, port: int) -> None:

0 commit comments

Comments
 (0)