diff --git a/optuna_dashboard/_app.py b/optuna_dashboard/_app.py index 2ad7874e..d7ddf72b 100644 --- a/optuna_dashboard/_app.py +++ b/optuna_dashboard/_app.py @@ -316,7 +316,7 @@ def create_app( except ValueError: response.status = 400 return {"reason": "invalid request."} - if component_type not in ["Note", "Artifact"]: + if component_type not in ["note", "artifact"]: response.status = 400 return {"reason": "component_type must be either 'Note' or 'Artifact'."} diff --git a/optuna_dashboard/_preference_setting.py b/optuna_dashboard/_preference_setting.py index d9d7c1c2..2310aa03 100644 --- a/optuna_dashboard/_preference_setting.py +++ b/optuna_dashboard/_preference_setting.py @@ -1,6 +1,7 @@ from __future__ import annotations from typing import TYPE_CHECKING +from typing import Any from optuna.storages import BaseStorage @@ -10,7 +11,7 @@ from .preferential._study import PreferentialStudy if TYPE_CHECKING: from typing import Literal - OUTPUT_COMPONENT_TYPE = Literal["Note", "Artifact"] + OUTPUT_COMPONENT_TYPE = Literal["note", "artifact"] _SYSTEM_ATTR_FEEDBACK_COMPONENT = "preference:component" @@ -19,22 +20,22 @@ def _register_preference_feedback_component_type( study_id: int, storage: BaseStorage, component_type: OUTPUT_COMPONENT_TYPE, - artifact_key: str = "", + artifact_key: str | None = None, ) -> None: + value: dict[str, Any] = {"type": component_type} + if artifact_key is not None: + value["artifact_key"] = artifact_key storage.set_study_system_attr( study_id=study_id, key=_SYSTEM_ATTR_FEEDBACK_COMPONENT, - value={ - "type": component_type, - "artifact_key": artifact_key, - }, + value=value, ) def register_preference_feedback_component_type( study: PreferentialStudy, component_type: OUTPUT_COMPONENT_TYPE, - artifact_key: str = "", + artifact_key: str | None = None, ) -> None: """Register output component to the study. @@ -48,6 +49,11 @@ def register_preference_feedback_component_type( this argument is used as the attribute key of the artifact. Each trial displays the artifact whose id is the value of the attribute. """ + if component_type == "artifact": + assert ( + artifact_key is not None + ), "artifact_key must be specified when component_type is Artifact" + _register_preference_feedback_component_type( study_id=study._study._study_id, storage=study._study._storage, diff --git a/python_tests/test_api.py b/python_tests/test_api.py index 1908c188..c478dc26 100644 --- a/python_tests/test_api.py +++ b/python_tests/test_api.py @@ -183,7 +183,7 @@ class APITestCase(TestCase): def test_change_component(self) -> None: storage = optuna.storages.InMemoryStorage() study = create_study(storage=storage, n_generate=3) - register_preference_feedback_component_type(study, "Note") + register_preference_feedback_component_type(study, "note") for _ in range(3): study.ask() @@ -193,7 +193,7 @@ class APITestCase(TestCase): app, f"/api/studies/{study_id}/preference_feedback_component_type", "PUT", - body=json.dumps({"type": "Artifact", "artifact_key": "image"}), + body=json.dumps({"type": "artifact", "artifact_key": "image"}), content_type="application/json", ) self.assertEqual(status, 204) @@ -207,7 +207,7 @@ class APITestCase(TestCase): self.assertEqual(status, 200) study_detail = json.loads(body) - assert study_detail["feedback_component_type"]["type"] == "Artifact" + assert study_detail["feedback_component_type"]["type"] == "artifact" assert study_detail["feedback_component_type"]["artifact_key"] == "image" def test_skip_trial(self) -> None: diff --git a/python_tests/test_preference_setting.py b/python_tests/test_preference_setting.py index 8a601af7..bcd53e42 100644 --- a/python_tests/test_preference_setting.py +++ b/python_tests/test_preference_setting.py @@ -11,10 +11,10 @@ from optuna_dashboard.preferential._study import PreferentialStudy class FeedbackSettingTestCase(TestCase): def test_widget_to_dict_from_dict(self) -> None: study = PreferentialStudy(optuna.create_study()) - register_preference_feedback_component_type(study, "Artifact", "image_key") + register_preference_feedback_component_type(study, "artifact", "image_key") system_attrs = study._study.system_attrs feedback_type = system_attrs.get(_SYSTEM_ATTR_FEEDBACK_COMPONENT, {}) assert "type" in feedback_type - assert feedback_type["type"] == "Artifact" + assert feedback_type["type"] == "artifact" assert "artifact_key" in feedback_type assert feedback_type["artifact_key"] == "image_key"