Clean up logic for determining the active model provider #309
Conversation
There was a problem hiding this comment.
Pull Request Overview
This PR refactors the test configuration by moving the parametrization from the test class decorator to the fixture definition itself, leveraging pytest's direct fixture parametrization feature.
- Removes the
@pytest.mark.parametrizedecorator from the test class - Adds
params=ModelProvider.__subclasses__()parameter directly to the@pytest.fixturedecorator
Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.
| def is_enabled(cls) -> bool: | ||
| """Return whether this model provider is enabled, based on the presence of required env vars.""" | ||
| def env_vars_are_set(cls) -> bool: | ||
| """Return whether this model provider's env vars have set values.""" |
There was a problem hiding this comment.
"is enabled" is a bit ambiguous between "is available for use" and "is the active model provider", since we now have the MODEL_PROVIDER env var that the method doesn't even account for.
| @staticmethod | ||
| @abc.abstractmethod | ||
| def get_model_source() -> ModelSource: | ||
| """Return the name of this model provider""" | ||
| raise NotImplementedError |
There was a problem hiding this comment.
Moved here from the bottom of the class definition to group it with the other "metadata" methods (as opposed to the "get model" methods).
| def get_priority() -> int: | ||
| """Return the priority of this model provider relative to the others. | ||
|
|
||
| 1 is the highest priority. | ||
|
|
||
| """ |
There was a problem hiding this comment.
I set the values for this based on the order we had in get_provider_class():
# Fallback to priority order if no specific provider is set
if AzureModelProvider.is_enabled():
return AzureModelProvider
elif OpenAiModelProvider.is_enabled():
return OpenAiModelProvider
elif BedrockModelProvider.is_enabled():
return BedrockModelProvider
return CAIIModelProvider| return all(map(os.environ.get, cls.get_env_var_names())) | ||
|
|
||
| @staticmethod | ||
| def get_provider_class() -> type["ModelProvider"]: |
There was a problem hiding this comment.
Moved to llm-service/app/services/models/providers/__init__.py since it doesn't make sense for it to be inherited.
| ] | ||
|
|
||
|
|
||
| def get_provider_class() -> type[_ModelProvider]: |
There was a problem hiding this comment.
Revamped to loop through ModelProvider's subclasses rather than have them hard-coded.
| class ModelProvider(abc.ABC): | ||
| class _ModelProvider(abc.ABC): |
There was a problem hiding this comment.
With get_provider_class() moved out, this class should no longer need to be invoked directly except for subclassing.
| ) | ||
| ) | ||
|
|
||
|
|
||
| @pytest.fixture() | ||
| @pytest.fixture(params=_ModelProvider.__subclasses__()) |
There was a problem hiding this comment.
I've discovered that we can do this instead of the
@pytest.mark.parametrize(
"EnabledModelProvider",
ModelProvider.__subclasses__(),
indirect=True,
)below! which is a lot simpler and easier on a test writer.
| for name in get_all_env_var_names(): | ||
| monkeypatch.delenv(name, raising=False) | ||
| for name in ModelProviderSubcls.get_env_var_names(): | ||
| monkeypatch.setenv(name, "test") | ||
| for name in get_all_env_var_names() - ModelProviderSubcls.get_env_var_names(): | ||
| monkeypatch.delenv(name, raising=False) |
There was a problem hiding this comment.
I've realized this is probably simpler?
| ModelProviderType = Literal["Azure", "CAII", "OpenAI", "Bedrock"] | ||
|
|
||
|
|
||
| class ModelSource(str, Enum): | ||
| AZURE = "Azure" | ||
| OPENAI = "OpenAI" | ||
| BEDROCK = "Bedrock" | ||
| CAII = "CAII" |
There was a problem hiding this comment.
This probably isn't the best place for ModelSource (I moved it here from llm-service/app/services/models/ to avoid a circular import), but I wanted to replace and remove the redundant ModelProviderType.
| @@ -136,7 +136,7 @@ class ProjectConfig(BaseModel): | |||
| chat_store_provider: ChatStoreProviderType | |||
| vector_db_provider: VectorDbProviderType | |||
| metadata_db_provider: MetadataDbProviderType | |||
| model_provider: Optional[ModelProviderType] = None | |||
| model_provider: Optional[ModelSource] = None | |||
There was a problem hiding this comment.
This still comes back as a string when the frontend calls /llm-service/amp/config 👍
mliu-cloudera
changed the title
mliu-cloudera
requested review from
jkwatson,
baasitsharief and
ewilliams-cloudera
| logger.info('falling back to model provider "CAII"') | ||
| return CAIIModelProvider |
There was a problem hiding this comment.
I think we should keep it for backwards compatibility.
Also helps lay groundwork for #307.