diff --git a/src/graphon/dsl/slim/package_loader.py b/src/graphon/dsl/slim/package_loader.py index 268ff3e..fc9359c 100644 --- a/src/graphon/dsl/slim/package_loader.py +++ b/src/graphon/dsl/slim/package_loader.py @@ -522,6 +522,6 @@ def _convert_model_type(value: str | None) -> ModelType | None: if value is None: return None try: - return ModelType.value_of(str(value)) + return ModelType(str(value)) except ValueError: return None diff --git a/src/graphon/model_runtime/entities/model_entities.py b/src/graphon/model_runtime/entities/model_entities.py index b0136f3..e7a043a 100644 --- a/src/graphon/model_runtime/entities/model_entities.py +++ b/src/graphon/model_runtime/entities/model_entities.py @@ -9,6 +9,20 @@ from graphon.model_runtime.entities.common_entities import I18nObject from graphon.model_runtime.entities.message_entities import PromptMessageContentType +_ORIGIN_MODEL_TYPE_BY_MODEL_TYPE: dict[str, str] = { + "llm": "text-generation", + "text-embedding": "embeddings", + "rerank": "reranking", + "speech2text": "speech2text", + "moderation": "moderation", + "tts": "tts", +} + +_MODEL_TYPE_BY_ORIGIN_MODEL_TYPE: dict[str, str] = { + origin_model_type: model_type + for model_type, origin_model_type in _ORIGIN_MODEL_TYPE_BY_MODEL_TYPE.items() +} + class ModelType(StrEnum): """Enum class for model type.""" @@ -21,52 +35,22 @@ class ModelType(StrEnum): TTS = auto() @classmethod - def value_of(cls, origin_model_type: str) -> ModelType: - """Map a provider-native model type string to a `ModelType`.""" - return _normalize_origin_model_type(origin_model_type) + def _missing_(cls, value: object) -> ModelType | None: + if isinstance(value, str): + model_type = _MODEL_TYPE_BY_ORIGIN_MODEL_TYPE.get(value) + if model_type is not None: + return cls(model_type) + return None def to_origin_model_type(self) -> str: """Map `ModelType` back to the provider-native model type string.""" - origin_model_type = _ORIGIN_MODEL_TYPE_BY_MODEL_TYPE.get(self) + origin_model_type = _ORIGIN_MODEL_TYPE_BY_MODEL_TYPE.get(self.value) if origin_model_type is None: msg = f"invalid model type {self}" raise ValueError(msg) return origin_model_type -_ORIGIN_MODEL_TYPE_BY_MODEL_TYPE: dict[ModelType, str] = { - ModelType.LLM: "text-generation", - ModelType.TEXT_EMBEDDING: "embeddings", - ModelType.RERANK: "reranking", - ModelType.SPEECH2TEXT: "speech2text", - ModelType.MODERATION: "moderation", - ModelType.TTS: "tts", -} - - -def _normalize_origin_model_type(origin_model_type: str) -> ModelType: - match origin_model_type: - case "text-generation": - normalized_model_type = ModelType.LLM - case "embeddings": - normalized_model_type = ModelType.TEXT_EMBEDDING - case "reranking": - normalized_model_type = ModelType.RERANK - case "speech2text": - normalized_model_type = ModelType.SPEECH2TEXT - case "moderation": - normalized_model_type = ModelType.MODERATION - case "tts": - normalized_model_type = ModelType.TTS - case _: - try: - normalized_model_type = ModelType(origin_model_type) - except ValueError as error: - msg = f"invalid origin model type {origin_model_type}" - raise ValueError(msg) from error - return normalized_model_type - - class FetchFrom(StrEnum): """Enum class for fetch from.""" diff --git a/tests/model_runtime/test_model_dispatch.py b/tests/model_runtime/test_model_dispatch.py index c88cb96..9562c3c 100644 --- a/tests/model_runtime/test_model_dispatch.py +++ b/tests/model_runtime/test_model_dispatch.py @@ -347,7 +347,7 @@ def invoke_speech_to_text( @pytest.mark.parametrize( - ("origin_model_type", "expected_model_type"), + ("raw_model_type", "expected_model_type"), [ ("text-generation", ModelType.LLM), (ModelType.LLM.value, ModelType.LLM), @@ -359,11 +359,11 @@ def invoke_speech_to_text( ("tts", ModelType.TTS), ], ) -def test_model_type_value_of_uses_model_map( - origin_model_type: str, +def test_model_type_accepts_origin_model_type_aliases( + raw_model_type: str, expected_model_type: ModelType, ) -> None: - assert ModelType.value_of(origin_model_type) == expected_model_type + assert ModelType(raw_model_type) == expected_model_type @pytest.mark.parametrize(