Clean up logic for determining the active model provider (#309)

* Parametrize ModelProvider fixture directly

* Simplify clearing env vars

* Move get_model_source() to the tops of classes

* Factor out get_provider_class()

* Rename is_enabled() to env_vars_are_set()

* Privatize ModelProvider base class

* Rework get_provider_class()

* Remove redundant ModelProviderType

* Clean up logs

Author: mliu-cloudera

llm-service/app/ai/indexing/summary_indexer.py

@@ -73,10 +73,9 @@
 from .base import BaseTextIndexer
 from .readers.base_reader import ReaderConfig, ChunksResult
 from ..vector_stores.vector_store_factory import VectorStoreFactory
-from ...config import settings
+from ...config import settings, ModelSource
 from ...services.metadata_apis import data_sources_metadata_api
-from ...services.models.providers import ModelProvider
-from ...services.models import ModelSource
+from ...services.models.providers import get_provider_class
 
 logger = logging.getLogger(__name__)
 
@@ -133,9 +132,7 @@ def __index_configuration(
         embed_summaries: bool = True,
     ) -> Dict[str, Any]:
         prompt_helper: Optional[PromptHelper] = None
-        model_source: ModelSource = (
-            ModelProvider.get_provider_class().get_model_source()
-        )
+        model_source: ModelSource = get_provider_class().get_model_source()
         if model_source == "CAII":
             # if we're using CAII, let's be conservative and use a small context window to account for mistral's small context
             prompt_helper = PromptHelper(context_window=3000)

llm-service/app/config.py

@@ -46,14 +46,21 @@
 
 import logging
 import os.path
+from enum import Enum
 from typing import cast, Optional, Literal
 
 
 SummaryStorageProviderType = Literal["Local", "S3"]
 ChatStoreProviderType = Literal["Local", "S3"]
 VectorDbProviderType = Literal["QDRANT", "OPENSEARCH"]
 MetadataDbProviderType = Literal["H2", "PostgreSQL"]
-ModelProviderType = Literal["Azure", "CAII", "OpenAI", "Bedrock"]
+
+
+class ModelSource(str, Enum):
+    AZURE = "Azure"
+    OPENAI = "OpenAI"
+    BEDROCK = "Bedrock"
+    CAII = "CAII"
 
 
 class _Settings:
@@ -185,14 +192,15 @@ def openai_api_base(self) -> Optional[str]:
         return os.environ.get("OPENAI_API_BASE")
 
     @property
-    def model_provider(self) -> Optional[ModelProviderType]:
+    def model_provider(self) -> Optional[ModelSource]:
         """The preferred model provider to use.
         Options: 'AZURE', 'CAII', 'OPENAI', 'BEDROCK'
         If not set, will use the first available provider in priority order."""
         provider = os.environ.get("MODEL_PROVIDER")
-        if provider and provider in ["Azure", "CAII", "OpenAI", "Bedrock"]:
-            return cast(ModelProviderType, provider)
-        return None
+        try:
+            return ModelSource(provider)
+        except ValueError:
+            return None
 
 
 settings = _Settings()

llm-service/app/routers/index/models/__init__.py

@@ -40,7 +40,7 @@
 from fastapi import APIRouter
 
 import app.services.models
-import app.services.models._model_source
+from app.config import ModelSource
 from .... import exceptions
 from ....services import models
 from ....services.caii.caii import describe_endpoint, build_model_response
@@ -71,7 +71,7 @@ def get_reranking_models() -> List[ModelResponse]:
     "/model_source", summary="Model source enabled - Bedrock, CAII, OpenAI or Azure"
 )
 @exceptions.propagates
-def get_model() -> app.services.models._model_source.ModelSource:
+def get_model() -> ModelSource:
     return app.services.models.get_model_source()
 
 

llm-service/app/services/amp_metadata/__init__.py

@@ -50,7 +50,7 @@
     ChatStoreProviderType,
     VectorDbProviderType,
     MetadataDbProviderType,
-    ModelProviderType,
+    ModelSource,
 )
 from app.services.models.providers import (
     CAIIModelProvider,
@@ -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
     aws_config: AwsConfig
     azure_config: AzureConfig
     caii_config: CaiiConfig
@@ -216,13 +216,13 @@ def validate_model_config(environ: dict[str, str]) -> ValidationResult:
             f"Preferred provider {preferred_provider} is properly configured. \n"
         )
         if preferred_provider == "Bedrock":
-            valid_model_config_exists = BedrockModelProvider.is_enabled()
+            valid_model_config_exists = BedrockModelProvider.env_vars_are_set()
         elif preferred_provider == "Azure":
-            valid_model_config_exists = AzureModelProvider.is_enabled()
+            valid_model_config_exists = AzureModelProvider.env_vars_are_set()
         elif preferred_provider == "OpenAI":
-            valid_model_config_exists = OpenAiModelProvider.is_enabled()
+            valid_model_config_exists = OpenAiModelProvider.env_vars_are_set()
         elif preferred_provider == "CAII":
-            valid_model_config_exists = CAIIModelProvider.is_enabled()
+            valid_model_config_exists = CAIIModelProvider.env_vars_are_set()
         return ValidationResult(
             valid=valid_model_config_exists,
             message=valid_message if valid_model_config_exists else message,
@@ -276,7 +276,7 @@ def validate_model_config(environ: dict[str, str]) -> ValidationResult:
 
     if message == "":
         # check to see if CAII models are available via discovery
-        if CAIIModelProvider.is_enabled():
+        if CAIIModelProvider.env_vars_are_set():
             message = "CAII models are available."
             valid_model_config_exists = True
         else:
@@ -388,7 +388,7 @@ def build_configuration(
     validate_config = validate(frozenset(env.items()))
 
     model_provider = (
-        TypeAdapter(ModelProviderType).validate_python(env.get("MODEL_PROVIDER"))
+        TypeAdapter(ModelSource).validate_python(env.get("MODEL_PROVIDER"))
         if env.get("MODEL_PROVIDER")
         else None
     )

llm-service/app/services/models/__init__.py

@@ -37,12 +37,12 @@
 #
 from .embedding import Embedding
 from .llm import LLM
-from .providers import ModelProvider
+from .providers import get_provider_class
 from .reranking import Reranking
-from ._model_source import ModelSource
+from ...config import ModelSource
 
-__all__ = ["Embedding", "LLM", "Reranking", "ModelSource"]
+__all__ = ["Embedding", "LLM", "Reranking", "get_model_source"]
 
 
 def get_model_source() -> ModelSource:
-    return ModelProvider.get_provider_class().get_model_source()
+    return get_provider_class().get_model_source()

llm-service/app/services/models/_model_source.py

@@ -41,7 +41,7 @@
 from llama_index.core.base.embeddings.base import BaseEmbedding
 
 from . import _model_type, _noop
-from .providers._model_provider import ModelProvider
+from .providers import get_provider_class
 from ..caii.types import ModelResponse
 
 
@@ -51,15 +51,15 @@ def get(cls, model_name: Optional[str] = None) -> BaseEmbedding:
         if model_name is None:
             model_name = cls.list_available()[0].model_id
 
-        return ModelProvider.get_provider_class().get_embedding_model(model_name)
+        return get_provider_class().get_embedding_model(model_name)
 
     @staticmethod
     def get_noop() -> BaseEmbedding:
         return _noop.DummyEmbeddingModel()
 
     @staticmethod
     def list_available() -> list[ModelResponse]:
-        return ModelProvider.get_provider_class().list_embedding_models()
+        return get_provider_class().list_embedding_models()
 
     @classmethod
     def test(cls, model_name: str) -> str:

llm-service/app/services/models/embedding.py

@@ -41,7 +41,7 @@
 from llama_index.core.base.llms.types import ChatMessage, MessageRole
 
 from . import _model_type, _noop
-from .providers._model_provider import ModelProvider
+from .providers import get_provider_class
 from ..caii.types import ModelResponse
 
 
@@ -51,15 +51,15 @@ def get(cls, model_name: Optional[str] = None) -> llms.LLM:
         if not model_name:
             model_name = cls.list_available()[0].model_id
 
-        return ModelProvider.get_provider_class().get_llm_model(model_name)
+        return get_provider_class().get_llm_model(model_name)
 
     @staticmethod
     def get_noop() -> llms.LLM:
         return _noop.DummyLlm()
 
     @staticmethod
     def list_available() -> list[ModelResponse]:
-        return ModelProvider.get_provider_class().list_llm_models()
+        return get_provider_class().list_llm_models()
 
     @classmethod
     def test(cls, model_name: str) -> Literal["ok"]:

llm-service/app/services/models/llm.py

@@ -35,16 +35,51 @@
 #  BUSINESS ADVANTAGE OR UNAVAILABILITY, OR LOSS OR CORRUPTION OF
 #  DATA.
 #
+import logging
+
+from app.config import settings
 from .azure import AzureModelProvider
 from .bedrock import BedrockModelProvider
 from .caii import CAIIModelProvider
 from .openai import OpenAiModelProvider
-from ._model_provider import ModelProvider
+from ._model_provider import _ModelProvider
+
+logger = logging.getLogger(__name__)
 
 __all__ = [
     "AzureModelProvider",
     "BedrockModelProvider",
     "CAIIModelProvider",
     "OpenAiModelProvider",
-    "ModelProvider",
+    "get_provider_class",
 ]
+
+
+def get_provider_class() -> type[_ModelProvider]:
+    """Return the ModelProvider subclass for the given provider name."""
+    model_providers: list[type[_ModelProvider]] = sorted(
+        _ModelProvider.__subclasses__(),
+        key=lambda ModelProviderSubcls: ModelProviderSubcls.get_priority(),
+    )
+
+    model_provider = settings.model_provider
+    for ModelProviderSubcls in model_providers:
+        if model_provider == ModelProviderSubcls.get_model_source():
+            logger.info(
+                'using model provider "%s" based on `MODEL_PROVIDER` env var',
+                ModelProviderSubcls.get_model_source().value,
+            )
+            return ModelProviderSubcls
+
+    # Fallback if no specific provider is set
+    for ModelProviderSubcls in model_providers:
+        if ModelProviderSubcls.env_vars_are_set():
+            logger.info(
+                'falling back to model provider "%s" based on env vars %s',
+                ModelProviderSubcls.get_model_source().value,
+                ModelProviderSubcls.get_env_var_names(),
+            )
+            return ModelProviderSubcls
+
+    logger.info('falling back to model provider "CAII"')
+    return CAIIModelProvider

llm-service/app/services/models/providers/__init__.py

@@ -42,50 +42,36 @@
 from llama_index.core.llms import LLM
 from llama_index.core.postprocessor.types import BaseNodePostprocessor
 
-from app.config import settings
-from .._model_source import ModelSource
+from app.config import ModelSource
 from ...caii.types import ModelResponse
 
 
-class ModelProvider(abc.ABC):
+class _ModelProvider(abc.ABC):
     @classmethod
-    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."""
         return all(map(os.environ.get, cls.get_env_var_names()))
 
     @staticmethod
-    def get_provider_class() -> type["ModelProvider"]:
-        """Return the ModelProvider subclass for the given provider name."""
-        from . import (
-            AzureModelProvider,
-            CAIIModelProvider,
-            OpenAiModelProvider,
-            BedrockModelProvider,
-        )
-
-        model_provider = settings.model_provider
-        if model_provider == "Azure":
-            return AzureModelProvider
-        elif model_provider == "CAII":
-            return CAIIModelProvider
-        elif model_provider == "OpenAI":
-            return OpenAiModelProvider
-        elif model_provider == "Bedrock":
-            return BedrockModelProvider
+    @abc.abstractmethod
+    def get_env_var_names() -> set[str]:
+        """Return the names of the env vars required by this model provider."""
+        raise NotImplementedError
 
-        # 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
+    @staticmethod
+    @abc.abstractmethod
+    def get_model_source() -> ModelSource:
+        """Return the name of this model provider"""
+        raise NotImplementedError
 
     @staticmethod
     @abc.abstractmethod
-    def get_env_var_names() -> set[str]:
-        """Return the names of the env vars required by this model provider."""
+    def get_priority() -> int:
+        """Return the priority of this model provider relative to the others.
+
+        1 is the highest priority.
+
+        """
         raise NotImplementedError
 
     @staticmethod
@@ -123,8 +109,3 @@ def get_embedding_model(name: str) -> BaseEmbedding:
     def get_reranking_model(name: str, top_n: int) -> BaseNodePostprocessor:
         """Return reranking model with `name`."""
         raise NotImplementedError
-
-    @staticmethod
-    @abc.abstractmethod
-    def get_model_source() -> ModelSource:
-        raise NotImplementedError