Move get_all_env_var_names() to _model_provider

Author: mliu-cloudera

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

@@ -36,6 +36,7 @@
 #  DATA.
 #
 import abc
+import itertools
 import os
 
 from llama_index.core.base.embeddings.base import BaseEmbedding
@@ -128,3 +129,12 @@ def get_reranking_model(name: str, top_n: int) -> BaseNodePostprocessor:
     @abc.abstractmethod
     def get_model_source() -> ModelSource:
         raise NotImplementedError
+
+
+def get_all_env_var_names() -> set[str]:
+    """Return the names of all the env vars required by all model providers."""
+    return set(
+        itertools.chain.from_iterable(
+            subcls.get_env_var_names() for subcls in ModelProvider.__subclasses__()
+        )
+    )

llm-service/app/tests/provider_mocks/bedrock.py

@@ -49,6 +49,7 @@
 from app.services.caii.types import ModelResponse
 from app.services.models import ModelProvider
 from app.services.models.providers import BedrockModelProvider
+from app.services.models.providers._model_provider import get_all_env_var_names
 
 TEXT_MODELS = [
     ("test.unavailable-text-model-v1", "NOT_AVAILABLE"),
@@ -175,15 +176,6 @@ def list_reranking_models() -> list[ModelResponse]:
         yield
 
 
-def get_all_env_var_names() -> set[str]:
-    """Return the names of all the env vars required by all model providers."""
-    return set(
-        itertools.chain.from_iterable(
-            subcls.get_env_var_names() for subcls in ModelProvider.__subclasses__()
-        )
-    )
-
-
 # TODO: move this test function to a discoverable place
 def test_bedrock(mock_bedrock, client) -> None:
     response = client.get("/llm-service/models/model_source")

llm-service/app/tests/services/test_models.py

@@ -43,16 +43,10 @@
 from app.services.caii import caii
 from app.services.caii.types import ListEndpointEntry
 from app.services.models.providers import BedrockModelProvider
-from app.services.models.providers._model_provider import ModelProvider
-
-
-def get_all_env_var_names() -> set[str]:
-    """Return the names of all the env vars required by all model providers."""
-    return set(
-        itertools.chain.from_iterable(
-            subcls.get_env_var_names() for subcls in ModelProvider.__subclasses__()
-        )
-    )
+from app.services.models.providers._model_provider import (
+    ModelProvider,
+    get_all_env_var_names,
+)
 
 
 @pytest.fixture(params=ModelProvider.__subclasses__())