Skip to content

Commit 50d15a9

Browse files
sawa3030sawa3030
committed
use mcperror
1 parent 95e79b1 commit 50d15a9

File tree

2 files changed

+118
-143
lines changed

2 files changed

+118
-143
lines changed

‎optuna_mcp/server.py‎

Lines changed: 99 additions & 124 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,9 @@
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
10+
from mcp.shared.exceptions import McpError
11+
from mcp.types import ErrorData
12+
from mcp.types import INTERNAL_ERROR
1213
import optuna
1314
import optuna_dashboard
1415
import plotly
@@ -112,23 +113,21 @@ def create_study(
112113
return StudyResponse(study_name=study_name)
113114

114115
@mcp.tool(structured_output=True)
115-
def get_all_study_names() -> list[StudyResponse] | CallToolResult:
116+
def get_all_study_names() -> list[StudyResponse]:
116117
"""Get all study names from the storage."""
117118
storage: str | optuna.storages.BaseStorage | None = None
118119
if mcp.study is not None:
119120
storage = mcp.study._storage
120121
elif mcp.storage is not None:
121122
storage = mcp.storage
122123
else:
123-
return CallToolResult(
124-
content=[TextContent(type="text", text="No storage specified.")], isError=True
125-
)
124+
raise McpError(ErrorData(code=INTERNAL_ERROR, message="No storage specified."))
126125

127126
study_names = optuna.get_all_study_names(storage)
128127
return [StudyResponse(study_name=name) for name in study_names]
129128

130129
@mcp.tool(structured_output=True)
131-
def ask(search_space: dict) -> TrialResponse | CallToolResult:
130+
def ask(search_space: dict) -> TrialResponse:
132131
"""Suggest new parameters using Optuna
133132
134133
search_space must be a string that can be evaluated to a dictionary to specify Optuna's distributions.
@@ -142,18 +141,14 @@ def ask(search_space: dict) -> TrialResponse | CallToolResult:
142141
for name, dist in search_space.items()
143142
}
144143
except Exception as e:
145-
return CallToolResult(
146-
content=[TextContent(type="text", text=f"Error: {e}")], isError=True
147-
)
144+
raise McpError(ErrorData(code=INTERNAL_ERROR, message=f"Error: {e}")) from e
148145

149146
if mcp.study is None:
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,
147+
raise McpError(
148+
ErrorData(
149+
code=INTERNAL_ERROR,
150+
message="No study has been created. Please create a study first.",
151+
)
157152
)
158153

159154
trial = mcp.study.ask(fixed_distributions=distributions)
@@ -164,16 +159,14 @@ def ask(search_space: dict) -> TrialResponse | CallToolResult:
164159
)
165160

166161
@mcp.tool(structured_output=True)
167-
def tell(trial_number: int, values: float | list[float]) -> TrialResponse | CallToolResult:
162+
def tell(trial_number: int, values: float | list[float]) -> TrialResponse:
168163
"""Report the result of a trial"""
169164
if mcp.study is None:
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,
165+
raise McpError(
166+
ErrorData(
167+
code=INTERNAL_ERROR,
168+
message="No study has been created. Please create a study first.",
169+
)
177170
)
178171

179172
mcp.study.tell(
@@ -190,7 +183,7 @@ def tell(trial_number: int, values: float | list[float]) -> TrialResponse | Call
190183
@mcp.tool(structured_output=True)
191184
def set_sampler(
192185
name: typing.Literal["TPESampler", "NSGAIISampler", "RandomSampler", "GPSampler"],
193-
) -> StudyResponse | CallToolResult:
186+
) -> StudyResponse:
194187
"""Set the sampler for the study.
195188
The sampler must be one of the following:
196189
- TPESampler
@@ -204,33 +197,28 @@ def set_sampler(
204197
"""
205198
sampler = getattr(optuna.samplers, name)()
206199
if mcp.study is None:
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,
200+
raise McpError(
201+
ErrorData(
202+
code=INTERNAL_ERROR,
203+
message="No study has been created. Please create a study first.",
204+
)
214205
)
206+
215207
mcp.study.sampler = sampler
216208
return StudyResponse(
217209
study_name=mcp.study.study_name,
218210
sampler_name=name,
219211
)
220212

221213
@mcp.tool(structured_output=True)
222-
def set_trial_user_attr(
223-
trial_number: int, key: str, value: typing.Any
224-
) -> str | CallToolResult:
214+
def set_trial_user_attr(trial_number: int, key: str, value: typing.Any) -> str:
225215
"""Set user attributes for a trial"""
226216
if mcp.study is None:
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,
217+
raise McpError(
218+
ErrorData(
219+
code=INTERNAL_ERROR,
220+
message="No study has been created. Please create a study first.",
221+
)
234222
)
235223

236224
storage = mcp.study._storage
@@ -241,17 +229,16 @@ def set_trial_user_attr(
241229
return f"User attribute {key} set to {json.dumps(value)} for trial {trial_number}"
242230

243231
@mcp.tool(structured_output=True)
244-
def get_trial_user_attrs(trial_number: int) -> TrialResponse | CallToolResult:
232+
def get_trial_user_attrs(trial_number: int) -> TrialResponse:
245233
"""Get user attributes in a trial"""
246234
if mcp.study is None:
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,
235+
raise McpError(
236+
ErrorData(
237+
code=INTERNAL_ERROR,
238+
message="No study has been created. Please create a study first.",
239+
)
254240
)
241+
255242
storage = mcp.study._storage
256243
trial_id = storage.get_trial_id_from_study_id_trial_number(
257244
mcp.study._study_id, trial_number
@@ -263,7 +250,7 @@ def get_trial_user_attrs(trial_number: int) -> TrialResponse | CallToolResult:
263250
)
264251

265252
@mcp.tool(structured_output=True)
266-
def set_metric_names(metric_names: list[str]) -> str | CallToolResult:
253+
def set_metric_names(metric_names: list[str]) -> str:
267254
"""Set metric_names. metric_names are labels used to distinguish what each objective value is.
268255
269256
Args:
@@ -272,78 +259,72 @@ def set_metric_names(metric_names: list[str]) -> str | CallToolResult:
272259
The length of metric_names list must be the same with the number of objectives.
273260
"""
274261
if mcp.study is None:
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,
262+
raise McpError(
263+
ErrorData(
264+
code=INTERNAL_ERROR,
265+
message="No study has been created. Please create a study first.",
266+
)
282267
)
268+
283269
mcp.study.set_metric_names(metric_names)
284270
return f"metric_names set to {json.dumps(metric_names)}"
285271

286272
@mcp.tool(structured_output=True)
287-
def get_metric_names() -> str | CallToolResult:
273+
def get_metric_names() -> str:
288274
"""Get metric_names"""
289275
if mcp.study is None:
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,
276+
raise McpError(
277+
ErrorData(
278+
code=INTERNAL_ERROR,
279+
message="No study has been created. Please create a study first.",
280+
)
297281
)
282+
298283
return f"Metric names: {json.dumps(mcp.study.metric_names)}"
299284

300285
@mcp.tool(structured_output=True)
301-
def get_directions() -> StudyResponse | CallToolResult:
286+
def get_directions() -> StudyResponse:
302287
"""Get the directions of the study."""
303288
if mcp.study is None:
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,
289+
raise McpError(
290+
ErrorData(
291+
code=INTERNAL_ERROR,
292+
message="No study has been created. Please create a study first.",
293+
)
311294
)
295+
312296
directions = [d.name.lower() for d in mcp.study.directions]
313297
return StudyResponse(
314298
study_name=mcp.study.study_name,
315299
directions=directions,
316300
)
317301

318302
@mcp.tool(structured_output=True)
319-
def get_trials() -> str | CallToolResult:
303+
def get_trials() -> str:
320304
"""Get all trials in a CSV format"""
321305
if mcp.study is None:
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,
306+
raise McpError(
307+
ErrorData(
308+
code=INTERNAL_ERROR,
309+
message="No study has been created. Please create a study first.",
310+
)
329311
)
312+
330313
csv_string = mcp.study.trials_dataframe().to_csv()
331314
return f"Trials: \n{csv_string}"
332315

333316
@mcp.tool(structured_output=True)
334-
def best_trial() -> TrialResponse | CallToolResult:
317+
def best_trial() -> TrialResponse:
335318
"""Get the best trial
336319
337320
This feature can only be used for single-objective optimization. If your study is multi-objective, use best_trials instead.
338321
"""
339322
if mcp.study is None:
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,
323+
raise McpError(
324+
ErrorData(
325+
code=INTERNAL_ERROR,
326+
message="No study has been created. Please create a study first.",
327+
)
347328
)
348329

349330
trial = mcp.study.best_trial
@@ -356,17 +337,16 @@ def best_trial() -> TrialResponse | CallToolResult:
356337
)
357338

358339
@mcp.tool(structured_output=True)
359-
def best_trials() -> list[TrialResponse] | CallToolResult:
340+
def best_trials() -> list[TrialResponse]:
360341
"""Return trials located at the Pareto front in the study."""
361342
if mcp.study is None:
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,
343+
raise McpError(
344+
ErrorData(
345+
code=INTERNAL_ERROR,
346+
message="No study has been created. Please create a study first.",
347+
)
369348
)
349+
370350
return [
371351
TrialResponse(
372352
trial_number=trial.number,
@@ -393,33 +373,31 @@ def _create_trial(trial: TrialToAdd) -> optuna.trial.FrozenTrial:
393373
)
394374

395375
@mcp.tool(structured_output=True)
396-
def add_trial(trial: TrialToAdd) -> str | CallToolResult:
376+
def add_trial(trial: TrialToAdd) -> str:
397377
"""Add a trial to the study."""
398378
if mcp.study is None:
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,
379+
raise McpError(
380+
ErrorData(
381+
code=INTERNAL_ERROR,
382+
message="No study has been created. Please create a study first.",
383+
)
406384
)
385+
407386
mcp.study.add_trial(_create_trial(trial))
408387
return "Trial was added."
409388

410389
@mcp.tool(structured_output=True)
411-
def add_trials(trials: list[TrialToAdd]) -> str | CallToolResult:
390+
def add_trials(trials: list[TrialToAdd]) -> str:
412391
"""Add multiple trials to the study."""
413392
frozen_trials = [_create_trial(trial) for trial in trials]
414393
if mcp.study is None:
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,
394+
raise McpError(
395+
ErrorData(
396+
code=INTERNAL_ERROR,
397+
message="No study has been created. Please create a study first.",
398+
)
422399
)
400+
423401
mcp.study.add_trials(frozen_trials)
424402
return f"{len(trials)} trials were added."
425403

@@ -640,7 +618,7 @@ def plot_rank(
640618
return Image(data=plotly.io.to_image(fig), format="png")
641619

642620
@mcp.tool(structured_output=True)
643-
def launch_optuna_dashboard(port: int = 58080) -> str | CallToolResult:
621+
def launch_optuna_dashboard(port: int = 58080) -> str:
644622
"""Launch the Optuna dashboard"""
645623
storage: str | optuna.storages.BaseStorage | None = None
646624
if mcp.dashboard_thread_port is not None:
@@ -651,14 +629,11 @@ def launch_optuna_dashboard(port: int = 58080) -> str | CallToolResult:
651629
elif mcp.storage is not None:
652630
storage = mcp.storage
653631
else:
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,
632+
raise McpError(
633+
ErrorData(
634+
code=INTERNAL_ERROR,
635+
message="No study has been created. Please create a study first.",
636+
)
662637
)
663638

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

0 commit comments

Comments
 (0)