From 345589409354e7beb4d0d7f1401dba9f707b2462 Mon Sep 17 00:00:00 2001 From: Niko Izsak Date: Wed, 16 Apr 2025 11:07:27 +0000 Subject: [PATCH 01/10] added support for custom scope in get_azure_ad_token_provider --- .../get_azure_ad_token_provider.py | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/litellm/secret_managers/get_azure_ad_token_provider.py b/litellm/secret_managers/get_azure_ad_token_provider.py index 5403675b97..797a0f8c7e 100644 --- a/litellm/secret_managers/get_azure_ad_token_provider.py +++ b/litellm/secret_managers/get_azure_ad_token_provider.py @@ -1,8 +1,8 @@ import os -from typing import Callable +from typing import Callable, Optional -def get_azure_ad_token_provider() -> Callable[[], str]: +def get_azure_ad_token_provider(azure_scope: Optional[str] = None) -> Callable[[], str]: """ Get Azure AD token provider based on Service Principal with Secret workflow. @@ -11,15 +11,22 @@ def get_azure_ad_token_provider() -> Callable[[], str]: https://learn.microsoft.com/en-us/python/api/overview/azure/identity-readme?view=azure-python#service-principal-with-secret; https://learn.microsoft.com/en-us/python/api/azure-identity/azure.identity.clientsecretcredential?view=azure-python. + Args: + azure_scope (str, optional): The Azure scope to request token for. + Defaults to environment variable AZURE_SCOPE or + "https://cognitiveservices.azure.com/.default". + Returns: Callable that returns a temporary authentication token. """ import azure.identity as identity from azure.identity import get_bearer_token_provider - azure_scope = os.environ.get( - "AZURE_SCOPE", "https://cognitiveservices.azure.com/.default" - ) + if azure_scope is None: + azure_scope = os.environ.get( + "AZURE_SCOPE", "https://cognitiveservices.azure.com/.default" + ) + cred = os.environ.get("AZURE_CREDENTIAL", "ClientSecretCredential") cred_cls = getattr(identity, cred) From 15ac0bd44025f7ea2d4093c702d6d501c741a063 Mon Sep 17 00:00:00 2001 From: Niko Izsak Date: Wed, 16 Apr 2025 11:08:17 +0000 Subject: [PATCH 02/10] if AZURE_FEDERATED_TOKEN_FILE not set, use azure_token_provider to retrive token with the oidc audiances as scope --- litellm/secret_managers/main.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index e505484b4b..657a5b885d 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -12,6 +12,9 @@ from litellm._logging import print_verbose, verbose_logger from litellm.caching.caching import DualCache from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.proxy._types import KeyManagementSystem +from litellm.secret_managers.get_azure_ad_token_provider import ( + get_azure_ad_token_provider, +) oidc_cache = DualCache() @@ -168,7 +171,12 @@ def get_secret( # noqa: PLR0915 # https://azure.github.io/azure-workload-identity/docs/quick-start.html azure_federated_token_file = os.getenv("AZURE_FEDERATED_TOKEN_FILE") if azure_federated_token_file is None: - raise ValueError("AZURE_FEDERATED_TOKEN_FILE not found in environment") + verbose_logger.warning("AZURE_FEDERATED_TOKEN_FILE not found in environment will use Azure AD token provider") + azure_token_provider = get_azure_ad_token_provider(azure_scope=oidc_aud) + oidc_token = azure_token_provider() + if oidc_token is None: + raise ValueError("Azure OIDC provider failed") + return oidc_token with open(azure_federated_token_file, "r") as f: oidc_token = f.read() return oidc_token From cdc14fa7fbc3173b8c2cbb7a8eb895ae40d078ed Mon Sep 17 00:00:00 2001 From: Niko Izsak Date: Wed, 16 Apr 2025 15:04:35 +0200 Subject: [PATCH 03/10] fix bug where oidc audience that contains "/" won't be extract correctly --- litellm/secret_managers/main.py | 45 ++++++++++----------------------- 1 file changed, 13 insertions(+), 32 deletions(-) diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index 657a5b885d..1b1ae2f630 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -105,6 +105,7 @@ def get_secret( # noqa: PLR0915 if secret_name.startswith("oidc/"): secret_name_split = secret_name.replace("oidc/", "") oidc_provider, oidc_aud = secret_name_split.split("/", 1) + oidc_aud = "/".join(secret_name_split.split("/")[1:]) # TODO: Add caching for HTTP requests if oidc_provider == "google": oidc_token = oidc_cache.get_cache(key=secret_name) @@ -140,10 +141,7 @@ def get_secret( # noqa: PLR0915 # https://docs.github.com/en/actions/deployment/security-hardening-your-deployments/configuring-openid-connect-in-cloud-providers#using-custom-actions actions_id_token_request_url = os.getenv("ACTIONS_ID_TOKEN_REQUEST_URL") actions_id_token_request_token = os.getenv("ACTIONS_ID_TOKEN_REQUEST_TOKEN") - if ( - actions_id_token_request_url is None - or actions_id_token_request_token is None - ): + if actions_id_token_request_url is None or actions_id_token_request_token is None: raise ValueError( "ACTIONS_ID_TOKEN_REQUEST_URL or ACTIONS_ID_TOKEN_REQUEST_TOKEN not found in environment" ) @@ -171,7 +169,9 @@ def get_secret( # noqa: PLR0915 # https://azure.github.io/azure-workload-identity/docs/quick-start.html azure_federated_token_file = os.getenv("AZURE_FEDERATED_TOKEN_FILE") if azure_federated_token_file is None: - verbose_logger.warning("AZURE_FEDERATED_TOKEN_FILE not found in environment will use Azure AD token provider") + verbose_logger.warning( + "AZURE_FEDERATED_TOKEN_FILE not found in environment will use Azure AD token provider" + ) azure_token_provider = get_azure_ad_token_provider(azure_scope=oidc_aud) oidc_token = azure_token_provider() if oidc_token is None: @@ -203,10 +203,7 @@ def get_secret( # noqa: PLR0915 raise ValueError("Unsupported OIDC provider") try: - if ( - _should_read_secret_from_secret_manager() - and litellm.secret_manager_client is not None - ): + if _should_read_secret_from_secret_manager() and litellm.secret_manager_client is not None: try: client = litellm.secret_manager_client key_manager = "local" @@ -232,9 +229,7 @@ def get_secret( # noqa: PLR0915 ): encrypted_secret: Any = os.getenv(secret_name) if encrypted_secret is None: - raise ValueError( - "Google KMS requires the encrypted secret to be in the environment!" - ) + raise ValueError("Google KMS requires the encrypted secret to be in the environment!") b64_flag = _is_base64(encrypted_secret) if b64_flag is True: # if passed in as encoded b64 string encrypted_secret = base64.b64decode(encrypted_secret) @@ -249,20 +244,14 @@ def get_secret( # noqa: PLR0915 "ciphertext": ciphertext, } ) - secret = response.plaintext.decode( - "utf-8" - ) # assumes the original value was encoded with utf-8 + secret = response.plaintext.decode("utf-8") # assumes the original value was encoded with utf-8 elif key_manager == KeyManagementSystem.AWS_KMS.value: """ Only check the tokens which start with 'aws_kms/'. This prevents latency impact caused by checking all keys. """ encrypted_value = os.getenv(secret_name, None) if encrypted_value is None: - raise Exception( - "AWS KMS - Encrypted Value of Key={} is None".format( - secret_name - ) - ) + raise Exception("AWS KMS - Encrypted Value of Key={} is None".format(secret_name)) # Decode the base64 encoded ciphertext ciphertext_blob = base64.b64decode(encrypted_value) @@ -289,14 +278,10 @@ def get_secret( # noqa: PLR0915 print_verbose(f"get_secret_value_response: {secret}") elif key_manager == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value: try: - secret = client.get_secret_from_google_secret_manager( - secret_name - ) + secret = client.get_secret_from_google_secret_manager(secret_name) print_verbose(f"secret from google secret manager: {secret}") if secret is None: - raise ValueError( - f"No secret found in Google Secret Manager for {secret_name}" - ) + raise ValueError(f"No secret found in Google Secret Manager for {secret_name}") except Exception as e: print_verbose(f"An error occurred - {str(e)}") raise e @@ -304,9 +289,7 @@ def get_secret( # noqa: PLR0915 try: secret = client.sync_read_secret(secret_name=secret_name) if secret is None: - raise ValueError( - f"No secret found in Hashicorp Secret Manager for {secret_name}" - ) + raise ValueError(f"No secret found in Hashicorp Secret Manager for {secret_name}") except Exception as e: print_verbose(f"An error occurred - {str(e)}") raise e @@ -331,9 +314,7 @@ def get_secret( # noqa: PLR0915 else: secret = os.environ.get(secret_name) secret_value_as_bool = str_to_bool(secret) if secret is not None else None - if secret_value_as_bool is not None and isinstance( - secret_value_as_bool, bool - ): + if secret_value_as_bool is not None and isinstance(secret_value_as_bool, bool): return secret_value_as_bool else: return secret From 7a9c305569b568128c66d874bb87790b0b18eba5 Mon Sep 17 00:00:00 2001 From: Niko Izsak Date: Wed, 16 Apr 2025 16:16:18 +0200 Subject: [PATCH 04/10] added tests for get_secret with oidc --- .../test_get_secret_oidc.py | 171 ++++++++++++++++++ 1 file changed, 171 insertions(+) create mode 100644 tests/litellm_utils_tests/test_get_secret_oidc.py diff --git a/tests/litellm_utils_tests/test_get_secret_oidc.py b/tests/litellm_utils_tests/test_get_secret_oidc.py new file mode 100644 index 0000000000..875e9f40fc --- /dev/null +++ b/tests/litellm_utils_tests/test_get_secret_oidc.py @@ -0,0 +1,171 @@ +import pytest +import os +from unittest.mock import Mock, patch +from litellm.secret_managers.main import get_secret +import logging + +# Set up logging for debugging +logging.basicConfig(level=logging.DEBUG) +logger = logging.getLogger(__name__) + +# Mock HTTPHandler and oidc_cache +class MockHTTPHandler: + def __init__(self, timeout): + self.timeout = timeout + self.status_code = 200 + self.text = "mocked_token" + self.json_data = {"value": "mocked_token"} + + def get(self, url, params=None, headers=None): + # Store params for audience verification + self.last_params = params + logger.debug(f"MockHTTPHandler.get called with url={url}, params={params}, headers={headers}") + mock_response = Mock() + mock_response.status_code = self.status_code + mock_response.text = self.text + mock_response.json.return_value = self.json_data + return mock_response + +@pytest.fixture +def mock_oidc_cache(): + cache = Mock() + cache.get_cache.return_value = None + cache.set_cache = Mock() + return cache + +@pytest.fixture +def mock_env(): + with patch.dict(os.environ, {}, clear=True): + yield os.environ + +@patch('litellm.secret_managers.main.oidc_cache') +@patch('litellm.secret_managers.main.HTTPHandler') +def test_oidc_google_success(mock_http_handler, mock_oidc_cache): + mock_oidc_cache.get_cache.return_value = None + mock_handler = MockHTTPHandler(timeout=600.0) + mock_http_handler.return_value = mock_handler + secret_name = "oidc/google/[invalid url, do not cite]" + result = get_secret(secret_name) + + assert result == "mocked_token" + assert mock_handler.last_params == {"audience": "[invalid url, do not cite]"} + mock_oidc_cache.set_cache.assert_called_once_with( + key=secret_name, value="mocked_token", ttl=3540 + ) + +@patch('litellm.secret_managers.main.oidc_cache') +def test_oidc_google_cached(mock_oidc_cache): + mock_oidc_cache.get_cache.return_value = "cached_token" + + secret_name = "oidc/google/[invalid url, do not cite]" + with patch('litellm.HTTPHandler') as mock_http: + result = get_secret(secret_name) + + assert result == "cached_token", f"Expected cached token, got {result}" + mock_oidc_cache.get_cache.assert_called_with(key=secret_name) + mock_http.assert_not_called() + +def test_oidc_google_failure(mock_oidc_cache): + mock_handler = MockHTTPHandler(timeout=600.0) + mock_handler.status_code = 400 + + with patch('litellm.secret_managers.main.HTTPHandler', return_value=mock_handler): + mock_oidc_cache.get_cache.return_value = None + secret_name = "oidc/google/https://example.com/api" + + with pytest.raises(ValueError, match="Google OIDC provider failed"): + get_secret(secret_name) + +def test_oidc_circleci_success(mock_env): + mock_env["CIRCLE_OIDC_TOKEN"] = "circleci_token" + + secret_name = "oidc/circleci/test-audience" + result = get_secret(secret_name) + + assert result == "circleci_token" + +def test_oidc_circleci_failure(): + secret_name = "oidc/circleci/test-audience" + + with pytest.raises(ValueError, match="CIRCLE_OIDC_TOKEN not found in environment"): + get_secret(secret_name) + +@patch('litellm.secret_managers.main.oidc_cache') +@patch('litellm.secret_managers.main.HTTPHandler') +def test_oidc_github_success(mock_http_handler, mock_oidc_cache, mock_env): + mock_env["ACTIONS_ID_TOKEN_REQUEST_URL"] = "https://github.com/token" + mock_env["ACTIONS_ID_TOKEN_REQUEST_TOKEN"] = "github_token" + mock_oidc_cache.get_cache.return_value = None + mock_handler = MockHTTPHandler(timeout=600.0) + mock_http_handler.return_value = mock_handler + + secret_name = "oidc/github/github-audience" + result = get_secret(secret_name) + + assert result == "mocked_token", f"Expected token 'mocked_token', got {result}" + assert mock_handler.last_params == {"audience": "github-audience"} + logger.debug(f"set_cache call args: {mock_oidc_cache.set_cache.call_args}") + mock_oidc_cache.set_cache.assert_called_once() + mock_oidc_cache.set_cache.assert_called_with( + key=secret_name, value="mocked_token", ttl=295 + ) + +def test_oidc_github_missing_env(): + secret_name = "oidc/github/github-audience" + + with pytest.raises(ValueError, match="ACTIONS_ID_TOKEN_REQUEST_URL or ACTIONS_ID_TOKEN_REQUEST_TOKEN not found in environment"): + get_secret(secret_name) + +def test_oidc_azure_file_success(mock_env, tmp_path): + token_file = tmp_path / "token.txt" + token_file.write_text("azure_token") + mock_env["AZURE_FEDERATED_TOKEN_FILE"] = str(token_file) + + secret_name = "oidc/azure/azure-audience" + result = get_secret(secret_name) + + assert result == "azure_token" + +@patch('litellm.secret_managers.main.get_azure_ad_token_provider') +def test_oidc_azure_ad_token_success(mock_get_azure_ad_token_provider): + mock_token_provider = Mock(return_value="azure_ad_token") + mock_get_azure_ad_token_provider.return_value = mock_token_provider + secret_name = "oidc/azure/api://azure-audience" + result = get_secret(secret_name) + + assert result == "azure_ad_token" + mock_get_azure_ad_token_provider.assert_called_once_with(azure_scope="api://azure-audience") + mock_token_provider.assert_called_once_with() + +def test_oidc_file_success(tmp_path): + token_file = tmp_path / "token.txt" + token_file.write_text("file_token") + + secret_name = f"oidc/file/{token_file}" + result = get_secret(secret_name) + + assert result == "file_token" + +def test_oidc_env_success(mock_env): + mock_env["CUSTOM_TOKEN"] = "env_token" + + secret_name = "oidc/env/CUSTOM_TOKEN" + result = get_secret(secret_name) + + assert result == "env_token" + +def test_oidc_env_path_success(mock_env, tmp_path): + token_file = tmp_path / "token.txt" + token_file.write_text("env_path_token") + mock_env["TOKEN_PATH"] = str(token_file) + + secret_name = "oidc/env_path/TOKEN_PATH" + result = get_secret(secret_name) + + assert result == "env_path_token" + +def test_unsupported_oidc_provider(): + secret_name = "oidc/unsupported/unsupported-audience" + + with pytest.raises(ValueError, match="Unsupported OIDC provider"): + get_secret(secret_name) \ No newline at end of file From ced4086ee2a1155a372dd2608aa802b50c2ff073 Mon Sep 17 00:00:00 2001 From: Niko Izsak Date: Wed, 16 Apr 2025 16:39:13 +0200 Subject: [PATCH 05/10] moved tests to litellm tests folder --- .../secrets_manager}/test_get_secret_oidc.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/{litellm_utils_tests => litellm/secrets_manager}/test_get_secret_oidc.py (100%) diff --git a/tests/litellm_utils_tests/test_get_secret_oidc.py b/tests/litellm/secrets_manager/test_get_secret_oidc.py similarity index 100% rename from tests/litellm_utils_tests/test_get_secret_oidc.py rename to tests/litellm/secrets_manager/test_get_secret_oidc.py From c31c695c5b7ef426d24a701f3db90d3b00e45b7a Mon Sep 17 00:00:00 2001 From: Niko Izsak Date: Wed, 16 Apr 2025 16:40:22 +0200 Subject: [PATCH 06/10] tes file naming aligned with source code --- .../test_get_secret_oidc.py => secrets_managers/test_main.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename tests/litellm/{secrets_manager/test_get_secret_oidc.py => secrets_managers/test_main.py} (100%) diff --git a/tests/litellm/secrets_manager/test_get_secret_oidc.py b/tests/litellm/secrets_managers/test_main.py similarity index 100% rename from tests/litellm/secrets_manager/test_get_secret_oidc.py rename to tests/litellm/secrets_managers/test_main.py From ab65497e668312720d2cd56eed0787b4a41c1d9a Mon Sep 17 00:00:00 2001 From: Niko Izsak Date: Thu, 17 Apr 2025 12:12:33 +0200 Subject: [PATCH 07/10] renamed test_main because it caused issue in the test in github workflow --- litellm/secret_managers/main.py | 13 +++++++++---- .../{test_main.py => test_main_oidc.py} | 0 2 files changed, 9 insertions(+), 4 deletions(-) rename tests/litellm/secrets_managers/{test_main.py => test_main_oidc.py} (100%) diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index 1b1ae2f630..b8b0cb2b51 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -173,10 +173,15 @@ def get_secret( # noqa: PLR0915 "AZURE_FEDERATED_TOKEN_FILE not found in environment will use Azure AD token provider" ) azure_token_provider = get_azure_ad_token_provider(azure_scope=oidc_aud) - oidc_token = azure_token_provider() - if oidc_token is None: - raise ValueError("Azure OIDC provider failed") - return oidc_token + try: + oidc_token = azure_token_provider() + if oidc_token is None: + raise ValueError("Azure OIDC provider returned None token") + return oidc_token + except Exception as e: + error_msg = f"Azure OIDC provider failed: {str(e)}" + verbose_logger.error(error_msg) + raise ValueError(error_msg) with open(azure_federated_token_file, "r") as f: oidc_token = f.read() return oidc_token diff --git a/tests/litellm/secrets_managers/test_main.py b/tests/litellm/secrets_managers/test_main_oidc.py similarity index 100% rename from tests/litellm/secrets_managers/test_main.py rename to tests/litellm/secrets_managers/test_main_oidc.py From 42eccbf25fd1f42e1554dc67b759118e46675af1 Mon Sep 17 00:00:00 2001 From: Niko Izsak Date: Thu, 17 Apr 2025 13:33:23 +0000 Subject: [PATCH 08/10] updated docs --- docs/my-website/docs/oidc.md | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/docs/my-website/docs/oidc.md b/docs/my-website/docs/oidc.md index f30edf5044..c9ff3ae6a6 100644 --- a/docs/my-website/docs/oidc.md +++ b/docs/my-website/docs/oidc.md @@ -19,6 +19,7 @@ LiteLLM supports the following OIDC identity providers: | CircleCI v2 | `circleci_v2`| No | | GitHub Actions | `github` | Yes | | Azure Kubernetes Service | `azure` | No | +| Azure AD | `azure` | Yes | | File | `file` | No | | Environment Variable | `env` | No | | Environment Path | `env_path` | No | @@ -108,6 +109,18 @@ model_list: aws_web_identity_token: "oidc/circleci_v2/" ``` +### Azure AD -> Amazon Bedrock +```yaml +model list: + - model_name: aws/claude-3-5-sonnet + litellm_params: + model: bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0 + aws_region: "eu-central-1" + aws_role_name: "arn:aws:iam::12345678:role/bedrock-role" + aws_web_identity_token: "oidc/azure/api://123-456-789-9d04" + aws_session_name: "litellm-session" +``` + #### Amazon IAM Role Configuration for CircleCI v2 -> Bedrock The configuration below is only an example. You should adjust the permissions and trust relationship to match your specific use case. From 95f39cd6d5eee375a4d8e593c1d862e84d7b9906 Mon Sep 17 00:00:00 2001 From: Niko Izsak Date: Thu, 17 Apr 2025 13:42:01 +0000 Subject: [PATCH 09/10] moved docs to the end of file --- docs/my-website/docs/oidc.md | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/docs/my-website/docs/oidc.md b/docs/my-website/docs/oidc.md index c9ff3ae6a6..1d788ea760 100644 --- a/docs/my-website/docs/oidc.md +++ b/docs/my-website/docs/oidc.md @@ -109,18 +109,6 @@ model_list: aws_web_identity_token: "oidc/circleci_v2/" ``` -### Azure AD -> Amazon Bedrock -```yaml -model list: - - model_name: aws/claude-3-5-sonnet - litellm_params: - model: bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0 - aws_region: "eu-central-1" - aws_role_name: "arn:aws:iam::12345678:role/bedrock-role" - aws_web_identity_token: "oidc/azure/api://123-456-789-9d04" - aws_session_name: "litellm-session" -``` - #### Amazon IAM Role Configuration for CircleCI v2 -> Bedrock The configuration below is only an example. You should adjust the permissions and trust relationship to match your specific use case. @@ -274,3 +262,15 @@ The custom role below is the recommended minimum permissions for the Azure appli _Note: Your UUIDs will be different._ Please contact us for paid enterprise support if you need help setting up Azure AD applications. + +### Azure AD -> Amazon Bedrock +```yaml +model list: + - model_name: aws/claude-3-5-sonnet + litellm_params: + model: bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0 + aws_region: "eu-central-1" + aws_role_name: "arn:aws:iam::12345678:role/bedrock-role" + aws_web_identity_token: "oidc/azure/api://123-456-789-9d04" + aws_session_name: "litellm-session" +``` From 4471b45e0c6e37ac5b255adad563be40d33a7277 Mon Sep 17 00:00:00 2001 From: Niko Izsak Date: Thu, 17 Apr 2025 17:11:08 +0200 Subject: [PATCH 10/10] fix aws region in example config --- docs/my-website/docs/oidc.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/my-website/docs/oidc.md b/docs/my-website/docs/oidc.md index 1d788ea760..3db4b6ecdc 100644 --- a/docs/my-website/docs/oidc.md +++ b/docs/my-website/docs/oidc.md @@ -269,7 +269,7 @@ model list: - model_name: aws/claude-3-5-sonnet litellm_params: model: bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0 - aws_region: "eu-central-1" + aws_region_name: "eu-central-1" aws_role_name: "arn:aws:iam::12345678:role/bedrock-role" aws_web_identity_token: "oidc/azure/api://123-456-789-9d04" aws_session_name: "litellm-session"