chore: turn OpenAIMixin into a pydantic.BaseModel (#3671)

# What does this PR do?

- implement get_api_key instead of relying on
LiteLLMOpenAIMixin.get_api_key
 - remove use of LiteLLMOpenAIMixin
 - add default initialize/shutdown methods to OpenAIMixin
 - remove __init__s to allow proper pydantic construction
- remove dead code from vllm adapter and associated / duplicate unit
tests
 - update vllm adapter to use openaimixin for model registration
 - remove ModelRegistryHelper from fireworks & together adapters
 - remove Inference from nvidia adapter
 - complete type hints on embedding_model_metadata
- allow extra fields on OpenAIMixin, for model_store, __provider_id__,
etc
 - new recordings for ollama
 - enhance the list models error handling
- update cerebras (remove cerebras-cloud-sdk) and anthropic (custom
model listing) inference adapters
 - parametrized test_inference_client_caching
- remove cerebras, databricks, fireworks, together from blanket mypy
exclude
 - removed unnecessary litellm deps

## Test Plan

ci

Author: mattf

docs/docs/providers/inference/remote_databricks.mdx

@@ -15,7 +15,7 @@ Databricks inference provider for running models on Databricks' unified analytic
 | Field | Type | Required | Default | Description |
 |-------|------|----------|---------|-------------|
 | `allowed_models` | `list[str \| None` | No |  | List of models that should be registered with the model registry. If None, all models are allowed. |
-| `url` | `<class 'str'>` | No |  | The URL for the Databricks model serving endpoint |
+| `url` | `str \| None` | No |  | The URL for the Databricks model serving endpoint |
 | `api_token` | `<class 'pydantic.types.SecretStr'>` | No |  | The Databricks API token |
 
 ## Sample Configuration

llama_stack/providers/registry/inference.py

@@ -52,9 +52,7 @@ def available_providers() -> list[ProviderSpec]:
             api=Api.inference,
             adapter_type="cerebras",
             provider_type="remote::cerebras",
-            pip_packages=[
-                "cerebras_cloud_sdk",
-            ],
+            pip_packages=[],
             module="llama_stack.providers.remote.inference.cerebras",
             config_class="llama_stack.providers.remote.inference.cerebras.CerebrasImplConfig",
             description="Cerebras inference provider for running models on Cerebras Cloud platform.",
@@ -169,7 +167,7 @@ def available_providers() -> list[ProviderSpec]:
             api=Api.inference,
             adapter_type="openai",
             provider_type="remote::openai",
-            pip_packages=["litellm"],
+            pip_packages=[],
             module="llama_stack.providers.remote.inference.openai",
             config_class="llama_stack.providers.remote.inference.openai.OpenAIConfig",
             provider_data_validator="llama_stack.providers.remote.inference.openai.config.OpenAIProviderDataValidator",
@@ -179,7 +177,7 @@ def available_providers() -> list[ProviderSpec]:
             api=Api.inference,
             adapter_type="anthropic",
             provider_type="remote::anthropic",
-            pip_packages=["litellm"],
+            pip_packages=["anthropic"],
             module="llama_stack.providers.remote.inference.anthropic",
             config_class="llama_stack.providers.remote.inference.anthropic.AnthropicConfig",
             provider_data_validator="llama_stack.providers.remote.inference.anthropic.config.AnthropicProviderDataValidator",
@@ -189,9 +187,7 @@ def available_providers() -> list[ProviderSpec]:
             api=Api.inference,
             adapter_type="gemini",
             provider_type="remote::gemini",
-            pip_packages=[
-                "litellm",
-            ],
+            pip_packages=[],
             module="llama_stack.providers.remote.inference.gemini",
             config_class="llama_stack.providers.remote.inference.gemini.GeminiConfig",
             provider_data_validator="llama_stack.providers.remote.inference.gemini.config.GeminiProviderDataValidator",
@@ -202,7 +198,6 @@ def available_providers() -> list[ProviderSpec]:
             adapter_type="vertexai",
             provider_type="remote::vertexai",
             pip_packages=[
-                "litellm",
                 "google-cloud-aiplatform",
             ],
             module="llama_stack.providers.remote.inference.vertexai",
@@ -233,9 +228,7 @@ def available_providers() -> list[ProviderSpec]:
             api=Api.inference,
             adapter_type="groq",
             provider_type="remote::groq",
-            pip_packages=[
-                "litellm",
-            ],
+            pip_packages=[],
             module="llama_stack.providers.remote.inference.groq",
             config_class="llama_stack.providers.remote.inference.groq.GroqConfig",
             provider_data_validator="llama_stack.providers.remote.inference.groq.config.GroqProviderDataValidator",
@@ -245,7 +238,7 @@ def available_providers() -> list[ProviderSpec]:
             api=Api.inference,
             adapter_type="llama-openai-compat",
             provider_type="remote::llama-openai-compat",
-            pip_packages=["litellm"],
+            pip_packages=[],
             module="llama_stack.providers.remote.inference.llama_openai_compat",
             config_class="llama_stack.providers.remote.inference.llama_openai_compat.config.LlamaCompatConfig",
             provider_data_validator="llama_stack.providers.remote.inference.llama_openai_compat.config.LlamaProviderDataValidator",
@@ -255,9 +248,7 @@ def available_providers() -> list[ProviderSpec]:
             api=Api.inference,
             adapter_type="sambanova",
             provider_type="remote::sambanova",
-            pip_packages=[
-                "litellm",
-            ],
+            pip_packages=[],
             module="llama_stack.providers.remote.inference.sambanova",
             config_class="llama_stack.providers.remote.inference.sambanova.SambaNovaImplConfig",
             provider_data_validator="llama_stack.providers.remote.inference.sambanova.config.SambaNovaProviderDataValidator",
@@ -287,7 +278,7 @@ def available_providers() -> list[ProviderSpec]:
             api=Api.inference,
             provider_type="remote::azure",
             adapter_type="azure",
-            pip_packages=["litellm"],
+            pip_packages=[],
             module="llama_stack.providers.remote.inference.azure",
             config_class="llama_stack.providers.remote.inference.azure.AzureConfig",
             provider_data_validator="llama_stack.providers.remote.inference.azure.config.AzureProviderDataValidator",

llama_stack/providers/remote/inference/anthropic/__init__.py

@@ -10,6 +10,6 @@
 async def get_adapter_impl(config: AnthropicConfig, _deps):
     from .anthropic import AnthropicInferenceAdapter
 
-    impl = AnthropicInferenceAdapter(config)
+    impl = AnthropicInferenceAdapter(config=config)
     await impl.initialize()
     return impl

llama_stack/providers/remote/inference/anthropic/anthropic.py

@@ -4,13 +4,19 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from llama_stack.providers.utils.inference.litellm_openai_mixin import LiteLLMOpenAIMixin
+from collections.abc import Iterable
+
+from anthropic import AsyncAnthropic
+
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 from .config import AnthropicConfig
 
 
-class AnthropicInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin):
+class AnthropicInferenceAdapter(OpenAIMixin):
+    config: AnthropicConfig
+
+    provider_data_api_key_field: str = "anthropic_api_key"
     # source: https://docs.claude.com/en/docs/build-with-claude/embeddings
     # TODO: add support for voyageai, which is where these models are hosted
     # embedding_model_metadata = {
@@ -23,22 +29,11 @@ class AnthropicInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin):
     #     "voyage-multimodal-3": {"embedding_dimension": 1024, "context_length": 32000},
     # }
 
-    def __init__(self, config: AnthropicConfig) -> None:
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            litellm_provider_name="anthropic",
-            api_key_from_config=config.api_key,
-            provider_data_api_key_field="anthropic_api_key",
-        )
-        self.config = config
-
-    async def initialize(self) -> None:
-        await super().initialize()
-
-    async def shutdown(self) -> None:
-        await super().shutdown()
-
-    get_api_key = LiteLLMOpenAIMixin.get_api_key
+    def get_api_key(self) -> str:
+        return self.config.api_key or ""
 
     def get_base_url(self):
         return "https://api.anthropic.com/v1"
+
+    async def list_provider_model_ids(self) -> Iterable[str]:
+        return [m.id async for m in AsyncAnthropic(api_key=self.get_api_key()).models.list()]

llama_stack/providers/remote/inference/azure/__init__.py

@@ -10,6 +10,6 @@
 async def get_adapter_impl(config: AzureConfig, _deps):
     from .azure import AzureInferenceAdapter
 
-    impl = AzureInferenceAdapter(config)
+    impl = AzureInferenceAdapter(config=config)
     await impl.initialize()
     return impl

llama_stack/providers/remote/inference/azure/azure.py

@@ -4,31 +4,20 @@
 # This source code is licensed under the terms described in the LICENSE file in
 # the root directory of this source tree.
 
-from typing import Any
 from urllib.parse import urljoin
 
-from llama_stack.apis.inference import ChatCompletionRequest
-from llama_stack.providers.utils.inference.litellm_openai_mixin import (
-    LiteLLMOpenAIMixin,
-)
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
 from .config import AzureConfig
 
 
-class AzureInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin):
-    def __init__(self, config: AzureConfig) -> None:
-        LiteLLMOpenAIMixin.__init__(
-            self,
-            litellm_provider_name="azure",
-            api_key_from_config=config.api_key.get_secret_value(),
-            provider_data_api_key_field="azure_api_key",
-            openai_compat_api_base=str(config.api_base),
-        )
-        self.config = config
+class AzureInferenceAdapter(OpenAIMixin):
+    config: AzureConfig
 
-    # Delegate the client data handling get_api_key method to LiteLLMOpenAIMixin
-    get_api_key = LiteLLMOpenAIMixin.get_api_key
+    provider_data_api_key_field: str = "azure_api_key"
+
+    def get_api_key(self) -> str:
+        return self.config.api_key.get_secret_value()
 
     def get_base_url(self) -> str:
         """
@@ -37,26 +26,3 @@ def get_base_url(self) -> str:
         Returns the Azure API base URL from the configuration.
         """
         return urljoin(str(self.config.api_base), "/openai/v1")
-
-    async def _get_params(self, request: ChatCompletionRequest) -> dict[str, Any]:
-        # Get base parameters from parent
-        params = await super()._get_params(request)
-
-        # Add Azure specific parameters
-        provider_data = self.get_request_provider_data()
-        if provider_data:
-            if getattr(provider_data, "azure_api_key", None):
-                params["api_key"] = provider_data.azure_api_key
-            if getattr(provider_data, "azure_api_base", None):
-                params["api_base"] = provider_data.azure_api_base
-            if getattr(provider_data, "azure_api_version", None):
-                params["api_version"] = provider_data.azure_api_version
-            if getattr(provider_data, "azure_api_type", None):
-                params["api_type"] = provider_data.azure_api_type
-        else:
-            params["api_key"] = self.config.api_key.get_secret_value()
-            params["api_base"] = str(self.config.api_base)
-            params["api_version"] = self.config.api_version
-            params["api_type"] = self.config.api_type
-
-        return params

llama_stack/providers/remote/inference/cerebras/__init__.py

@@ -12,7 +12,7 @@ async def get_adapter_impl(config: CerebrasImplConfig, _deps):
 
     assert isinstance(config, CerebrasImplConfig), f"Unexpected config type: {type(config)}"
 
-    impl = CerebrasInferenceAdapter(config)
+    impl = CerebrasInferenceAdapter(config=config)
 
     await impl.initialize()
 

llama_stack/providers/remote/inference/cerebras/cerebras.py

@@ -6,71 +6,21 @@
 
 from urllib.parse import urljoin
 
-from cerebras.cloud.sdk import AsyncCerebras
-
-from llama_stack.apis.inference import (
-    ChatCompletionRequest,
-    CompletionRequest,
-    Inference,
-    OpenAIEmbeddingsResponse,
-    TopKSamplingStrategy,
-)
-from llama_stack.providers.utils.inference.openai_compat import (
-    get_sampling_options,
-)
+from llama_stack.apis.inference import OpenAIEmbeddingsResponse
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
-from llama_stack.providers.utils.inference.prompt_adapter import (
-    chat_completion_request_to_prompt,
-    completion_request_to_prompt,
-)
 
 from .config import CerebrasImplConfig
 
 
-class CerebrasInferenceAdapter(
-    OpenAIMixin,
-    Inference,
-):
-    def __init__(self, config: CerebrasImplConfig) -> None:
-        self.config = config
-
-        # TODO: make this use provider data, etc. like other providers
-        self._cerebras_client = AsyncCerebras(
-            base_url=self.config.base_url,
-            api_key=self.config.api_key.get_secret_value(),
-        )
+class CerebrasInferenceAdapter(OpenAIMixin):
+    config: CerebrasImplConfig
 
     def get_api_key(self) -> str:
         return self.config.api_key.get_secret_value()
 
     def get_base_url(self) -> str:
         return urljoin(self.config.base_url, "v1")
 
-    async def initialize(self) -> None:
-        return
-
-    async def shutdown(self) -> None:
-        pass
-
-    async def _get_params(self, request: ChatCompletionRequest | CompletionRequest) -> dict:
-        if request.sampling_params and isinstance(request.sampling_params.strategy, TopKSamplingStrategy):
-            raise ValueError("`top_k` not supported by Cerebras")
-
-        prompt = ""
-        if isinstance(request, ChatCompletionRequest):
-            prompt = await chat_completion_request_to_prompt(request, self.get_llama_model(request.model))
-        elif isinstance(request, CompletionRequest):
-            prompt = await completion_request_to_prompt(request)
-        else:
-            raise ValueError(f"Unknown request type {type(request)}")
-
-        return {
-            "model": request.model,
-            "prompt": prompt,
-            "stream": request.stream,
-            **get_sampling_options(request.sampling_params),
-        }
-
     async def openai_embeddings(
         self,
         model: str,

llama_stack/providers/remote/inference/cerebras/config.py

@@ -22,7 +22,7 @@ class CerebrasImplConfig(RemoteInferenceProviderConfig):
         description="Base URL for the Cerebras API",
     )
     api_key: SecretStr = Field(
-        default=SecretStr(os.environ.get("CEREBRAS_API_KEY")),
+        default=SecretStr(os.environ.get("CEREBRAS_API_KEY")),  # type: ignore[arg-type]
         description="Cerebras API Key",
     )
 

llama_stack/providers/remote/inference/databricks/__init__.py

@@ -11,6 +11,6 @@ async def get_adapter_impl(config: DatabricksImplConfig, _deps):
     from .databricks import DatabricksInferenceAdapter
 
     assert isinstance(config, DatabricksImplConfig), f"Unexpected config type: {type(config)}"
-    impl = DatabricksInferenceAdapter(config)
+    impl = DatabricksInferenceAdapter(config=config)
     await impl.initialize()
     return impl