chore: make OpenAIMixin maintainable, turn OpenAIMixin into a pydantic.BaseModel

 - 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 w/ new tests
 - update cerebras (remove cerebras-cloud-sdk) and anthropic (custom model listing) inference adapters

Author: mattf

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.",
@@ -179,7 +177,7 @@ def available_providers() -> list[ProviderSpec]:
             api=Api.inference,
             adapter_type="anthropic",
             provider_type="remote::anthropic",
-            pip_packages=["litellm"],
+            pip_packages=["litellm", "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",

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 get_models(self) -> Iterable[str] | None:
+        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:
         """
@@ -38,25 +27,25 @@ def get_base_url(self) -> str:
         """
         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
+    # 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/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

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

@@ -9,11 +9,9 @@
 from databricks.sdk import WorkspaceClient
 
 from llama_stack.apis.inference import (
-    Inference,
     Model,
     OpenAICompletion,
 )
-from llama_stack.apis.models import ModelType
 from llama_stack.log import get_logger
 from llama_stack.providers.utils.inference.openai_mixin import OpenAIMixin
 
@@ -22,31 +20,21 @@
 logger = get_logger(name=__name__, category="inference::databricks")
 
 
-class DatabricksInferenceAdapter(
-    OpenAIMixin,
-    Inference,
-):
+class DatabricksInferenceAdapter(OpenAIMixin):
+    config: DatabricksImplConfig
+
     # source: https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/supported-models
-    embedding_model_metadata = {
+    embedding_model_metadata: dict[str, dict[str, int]] = {
         "databricks-gte-large-en": {"embedding_dimension": 1024, "context_length": 8192},
         "databricks-bge-large-en": {"embedding_dimension": 1024, "context_length": 512},
     }
 
-    def __init__(self, config: DatabricksImplConfig) -> None:
-        self.config = config
-
     def get_api_key(self) -> str:
         return self.config.api_token.get_secret_value()
 
     def get_base_url(self) -> str:
         return f"{self.config.url}/serving-endpoints"
 
-    async def initialize(self) -> None:
-        return
-
-    async def shutdown(self) -> None:
-        pass
-
     async def openai_completion(
         self,
         model: str,
@@ -72,31 +60,13 @@ async def openai_completion(
     ) -> OpenAICompletion:
         raise NotImplementedError()
 
-    async def list_models(self) -> list[Model] | None:
-        self._model_cache = {}  # from OpenAIMixin
-        ws_client = WorkspaceClient(host=self.config.url, token=self.get_api_key())  # TODO: this is not async
-        endpoints = ws_client.serving_endpoints.list()
-        for endpoint in endpoints:
-            model = Model(
-                provider_id=self.__provider_id__,
-                provider_resource_id=endpoint.name,
-                identifier=endpoint.name,
-            )
-            if endpoint.task == "llm/v1/chat":
-                model.model_type = ModelType.llm  # this is redundant, but informative
-            elif endpoint.task == "llm/v1/embeddings":
-                if endpoint.name not in self.embedding_model_metadata:
-                    logger.warning(f"No metadata information available for embedding model {endpoint.name}, skipping.")
-                    continue
-                model.model_type = ModelType.embedding
-                model.metadata = self.embedding_model_metadata[endpoint.name]
-            else:
-                logger.warning(f"Unknown model type, skipping: {endpoint}")
-                continue
-
-            self._model_cache[endpoint.name] = model
-
-        return list(self._model_cache.values())
+    async def get_models(self) -> list[Model] | None:
+        return [
+            endpoint.name
+            for endpoint in WorkspaceClient(
+                host=self.config.url, token=self.get_api_key()
+            ).serving_endpoints.list()  # TODO: this is not async
+        ]
 
     async def should_refresh_models(self) -> bool:
         return False

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

@@ -17,6 +17,6 @@ async def get_adapter_impl(config: FireworksImplConfig, _deps):
     from .fireworks import FireworksInferenceAdapter
 
     assert isinstance(config, FireworksImplConfig), f"Unexpected config type: {type(config)}"
-    impl = FireworksInferenceAdapter(config)
+    impl = FireworksInferenceAdapter(config=config)
     await impl.initialize()
     return impl