Merge branch 'main' into chroma

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

@@ -0,0 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .config import AzureConfig
+
+
+async def get_adapter_impl(config: AzureConfig, _deps):
+    from .azure import AzureInferenceAdapter
+
+    impl = AzureInferenceAdapter(config)
+    await impl.initialize()
+    return impl

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

@@ -0,0 +1,64 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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
+from .models import MODEL_ENTRIES
+
+
+class AzureInferenceAdapter(OpenAIMixin, LiteLLMOpenAIMixin):
+    def __init__(self, config: AzureConfig) -> None:
+        LiteLLMOpenAIMixin.__init__(
+            self,
+            MODEL_ENTRIES,
+            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
+
+    # Delegate the client data handling get_api_key method to LiteLLMOpenAIMixin
+    get_api_key = LiteLLMOpenAIMixin.get_api_key
+
+    def get_base_url(self) -> str:
+        """
+        Get the Azure API base URL.
+
+        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/azure/config.py

@@ -0,0 +1,63 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+import os
+from typing import Any
+
+from pydantic import BaseModel, Field, HttpUrl, SecretStr
+
+from llama_stack.schema_utils import json_schema_type
+
+
+class AzureProviderDataValidator(BaseModel):
+    azure_api_key: SecretStr = Field(
+        description="Azure API key for Azure",
+    )
+    azure_api_base: HttpUrl = Field(
+        description="Azure API base for Azure (e.g., https://your-resource-name.openai.azure.com)",
+    )
+    azure_api_version: str | None = Field(
+        default=None,
+        description="Azure API version for Azure (e.g., 2024-06-01)",
+    )
+    azure_api_type: str | None = Field(
+        default="azure",
+        description="Azure API type for Azure (e.g., azure)",
+    )
+
+
+@json_schema_type
+class AzureConfig(BaseModel):
+    api_key: SecretStr = Field(
+        description="Azure API key for Azure",
+    )
+    api_base: HttpUrl = Field(
+        description="Azure API base for Azure (e.g., https://your-resource-name.openai.azure.com)",
+    )
+    api_version: str | None = Field(
+        default_factory=lambda: os.getenv("AZURE_API_VERSION"),
+        description="Azure API version for Azure (e.g., 2024-12-01-preview)",
+    )
+    api_type: str | None = Field(
+        default_factory=lambda: os.getenv("AZURE_API_TYPE", "azure"),
+        description="Azure API type for Azure (e.g., azure)",
+    )
+
+    @classmethod
+    def sample_run_config(
+        cls,
+        api_key: str = "${env.AZURE_API_KEY:=}",
+        api_base: str = "${env.AZURE_API_BASE:=}",
+        api_version: str = "${env.AZURE_API_VERSION:=}",
+        api_type: str = "${env.AZURE_API_TYPE:=}",
+        **kwargs,
+    ) -> dict[str, Any]:
+        return {
+            "api_key": api_key,
+            "api_base": api_base,
+            "api_version": api_version,
+            "api_type": api_type,
+        }

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

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# 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.model_registry import (
+    ProviderModelEntry,
+)
+
+# https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/models?tabs=global-standard%2Cstandard-chat-completions
+LLM_MODEL_IDS = [
+    "gpt-5",
+    "gpt-5-mini",
+    "gpt-5-nano",
+    "gpt-5-chat",
+    "o1",
+    "o1-mini",
+    "o3-mini",
+    "o4-mini",
+    "gpt-4.1",
+    "gpt-4.1-mini",
+    "gpt-4.1-nano",
+]
+
+SAFETY_MODELS_ENTRIES = list[ProviderModelEntry]()
+
+MODEL_ENTRIES = [ProviderModelEntry(provider_model_id=m) for m in LLM_MODEL_IDS] + SAFETY_MODELS_ENTRIES