diff --git a/litellm/integrations/azure_cost_management/azure_cost_management_client.py b/litellm/integrations/azure_cost_management/azure_cost_management_client.py index f92e6d6de27..c098fc92807 100644 --- a/litellm/integrations/azure_cost_management/azure_cost_management_client.py +++ b/litellm/integrations/azure_cost_management/azure_cost_management_client.py @@ -30,6 +30,7 @@ class AzureCostManagementConfig: client_id: str client_secret: str = field(repr=False) api_version: str = "2023-11-01" + management_base_url: str = "https://management.azure.com" @classmethod def from_env(cls, subscription_id: str) -> AzureCostManagementConfig: @@ -52,6 +53,7 @@ class AzureCostManagementConfig: tenant_id=tenant_id or "", client_id=client_id or "", client_secret=client_secret or "", + management_base_url=os.getenv("AZURE_MANAGEMENT_BASE_URL", "https://management.azure.com"), ) @@ -105,7 +107,7 @@ class AzureCostManagementClient: valid response, typically due to reporting lag or zero utilization). """ url = ( - "https://management.azure.com/subscriptions/" + f"{self._config.management_base_url}/subscriptions/" f"{self._config.subscription_id}" "/providers/Microsoft.CostManagement/query" ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index bb6c4747d02..8c268da497a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -7464,6 +7464,24 @@ def giveup(e): return result +class _DemoAzureCostFetcher: + """Env-gated stub used only when AZURE_PTU_DEMO_USD is set. + + Returns the same USD amount for every reservation/day so the Usage page can show + a plausible azure_billing figure during customer demos without provisioning a real + PTU or waiting for Azure's 24-72h billing lag. Presence of the env var is + intentionally the only opt-in; it prints a loud warning so nobody ships to prod + with it accidentally set. + """ + + def __init__(self, amount_usd: float) -> None: + self._amount = amount_usd + self.last_currency = "USD" + + async def get_daily_cost(self, resource_id: str, day: Any) -> float: # noqa: ARG002 # demo signature must match AzureCostFetcher Protocol + return self._amount + + def _build_azure_cost_fetcher_if_enabled() -> Optional[Any]: """Return an AzureCostManagementClient when config + creds are present, else None. @@ -7472,9 +7490,31 @@ def _build_azure_cost_fetcher_if_enabled() -> Optional[Any]: piece is missing, azure_billing reservations no-op with a warning inside the rollup and manual reservations continue to accrue via the formula. + When ``AZURE_PTU_DEMO_USD`` is set, a stub fetcher returning that amount is + used instead of the real client. Demo-only; a WARNING is logged every call so + the mode is impossible to miss in production logs. + Evaluated at rollup call time so runtime config changes take effect without a proxy restart. """ + demo_env = os.getenv("AZURE_PTU_DEMO_USD") + if demo_env: + try: + amount = float(demo_env) + except ValueError: + verbose_proxy_logger.warning( + "AZURE_PTU_DEMO_USD=%r is not a valid float; skipping demo fetcher.", + demo_env, + ) + else: + verbose_proxy_logger.warning( + "AZURE_PTU_DEMO_USD=%.2f is set; azure_billing reservations will return " + "this stub USD amount instead of hitting Azure Cost Management. " + "Unset this env var before shipping to production.", + amount, + ) + return _DemoAzureCostFetcher(amount) + azure_ptu_billing = general_settings.get("azure_ptu_billing") or {} subscription_id = azure_ptu_billing.get("subscription_id") if not subscription_id: diff --git a/tests/test_litellm/integrations/azure_cost_management/test_azure_cost_management_client.py b/tests/test_litellm/integrations/azure_cost_management/test_azure_cost_management_client.py index 385ae15f3cc..6636fd9ab90 100644 --- a/tests/test_litellm/integrations/azure_cost_management/test_azure_cost_management_client.py +++ b/tests/test_litellm/integrations/azure_cost_management/test_azure_cost_management_client.py @@ -155,6 +155,43 @@ def test_client_secret_not_in_repr(): assert "client_secret" not in repr(cfg) +def test_config_management_base_url_defaults_to_public_cloud(): + cfg = AzureCostManagementConfig(subscription_id="s", tenant_id="t", client_id="c", client_secret="x") + assert cfg.management_base_url == "https://management.azure.com" + + +def test_config_management_base_url_reads_env_override(monkeypatch): + """Sovereign clouds + testability: AZURE_MANAGEMENT_BASE_URL can point at a different host.""" + monkeypatch.setenv("AZURE_TENANT_ID", "t") + monkeypatch.setenv("AZURE_CLIENT_ID", "c") + monkeypatch.setenv("AZURE_CLIENT_SECRET", "x") + monkeypatch.setenv("AZURE_MANAGEMENT_BASE_URL", "https://management.usgovcloudapi.net") + + cfg = AzureCostManagementConfig.from_env(subscription_id="s") + + assert cfg.management_base_url == "https://management.usgovcloudapi.net" + + +@pytest.mark.asyncio +async def test_get_daily_cost_hits_management_base_url_from_config(): + """Client must use the base URL from config (sovereign cloud support).""" + payload = {"properties": {"columns": [{"name": "Cost"}], "rows": []}} + http = _http_ok(payload) + cfg = AzureCostManagementConfig( + subscription_id="sub-x", + tenant_id="t", + client_id="c", + client_secret="x", + management_base_url="http://localhost:18443", + ) + client = AzureCostManagementClient(config=cfg, http_handler=http, token_provider=lambda: "tkn") + + await client.get_daily_cost("/subs/x/deploy/y", date(2026, 7, 15)) + + kwargs = http.post.await_args.kwargs + assert kwargs["url"].startswith("http://localhost:18443/subscriptions/sub-x/") + + @pytest.mark.asyncio async def test_get_daily_cost_wraps_network_errors(): """httpx.RequestError family (ConnectError, ReadTimeout, RemoteProtocolError) must surface as AzureCostManagementError."""