Add base_url to AzureOpenAI (#17996)

run-llama · Mar 3, 2025 · cb16cf6 · cb16cf6
1 parent e8f73a4
commit cb16cf6
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 2 deletions.
diff --git a/...dex-integrations/llms/llama-index-llms-azure-openai/llama_index/llms/azure_openai/base.py b/...dex-integrations/llms/llama-index-llms-azure-openai/llama_index/llms/azure_openai/base.py
@@ -11,6 +11,7 @@
     resolve_from_aliases,
 )
 from llama_index.llms.openai import OpenAI
+from llama_index.llms.openai.utils import DEFAULT_OPENAI_API_BASE
 from openai import AsyncAzureOpenAI
 from openai import AzureOpenAI as SyncAzureOpenAI
 from openai.lib.azure import AzureADTokenProvider
@@ -92,10 +93,13 @@ class AzureOpenAI(OpenAI):
     use_azure_ad: bool = Field(
         description="Indicates if Microsoft Entra ID (former Azure AD) is used for token authentication"
     )
-
     azure_ad_token_provider: Optional[AzureADTokenProvider] = Field(
         default=None, description="Callback function to provide Azure Entra ID token."
     )
+    api_base: Optional[str] = Field(
+        default=None,
+        description="The Azure Base URL to use. Useful for proxies on top of Azure OpenAI.",
+    )
 
     _azure_ad_token: Any = PrivateAttr(default=None)
     _client: SyncAzureOpenAI = PrivateAttr()
@@ -113,6 +117,7 @@ def __init__(
         reuse_client: bool = True,
         api_key: Optional[str] = None,
         api_version: Optional[str] = None,
+        api_base: Optional[str] = None,
         # azure specific
         azure_endpoint: Optional[str] = None,
         azure_deployment: Optional[str] = None,
@@ -157,6 +162,7 @@ def __init__(
             api_key=api_key,
             azure_endpoint=azure_endpoint,
             azure_deployment=azure_deployment,
+            api_base=api_base,
             azure_ad_token_provider=azure_ad_token_provider,
             use_azure_ad=use_azure_ad,
             api_version=api_version,
@@ -171,6 +177,10 @@ def __init__(
             **kwargs,
         )
 
+        # reset api_base to None if it is the default
+        if self.api_base == DEFAULT_OPENAI_API_BASE:
+            self.api_base = None
+
     @model_validator(mode="before")
     def validate_env(cls, values: Dict[str, Any]) -> Dict[str, Any]:
         """Validate necessary credentials are set."""
@@ -235,6 +245,7 @@ def _get_credential_kwargs(
             "timeout": self.timeout,
             "azure_endpoint": self.azure_endpoint,
             "azure_deployment": self.azure_deployment,
+            "base_url": self.api_base,
             "azure_ad_token_provider": self.azure_ad_token_provider,
             "api_version": self.api_version,
             "default_headers": self.default_headers,

diff --git a/llama-index-integrations/llms/llama-index-llms-azure-openai/pyproject.toml b/llama-index-integrations/llms/llama-index-llms-azure-openai/pyproject.toml
@@ -29,7 +29,7 @@ exclude = ["**/BUILD"]
 license = "MIT"
 name = "llama-index-llms-azure-openai"
 readme = "README.md"
-version = "0.3.0"
+version = "0.3.1"
 
 [tool.poetry.dependencies]
 python = ">=3.9,<4.0"