From 438774db3c2fd829de63c6b85ed174dc6fe0e380 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Mon, 23 Mar 2026 09:31:09 -0700 Subject: [PATCH] refactor(redis): move GCPIAMCredentialProvider to its own file Extract GCPIAMCredentialProvider and _generate_gcp_iam_access_token into litellm/_redis_credential_provider.py. _redis.py imports them from there, keeping the public API unchanged. --- litellm/_redis.py | 54 +++------------------------ litellm/_redis_credential_provider.py | 51 +++++++++++++++++++++++++ tests/test_litellm/test_redis.py | 8 ++-- 3 files changed, 61 insertions(+), 52 deletions(-) create mode 100644 litellm/_redis_credential_provider.py diff --git a/litellm/_redis.py b/litellm/_redis.py index 472b3c2c80d..b42bad1a651 100644 --- a/litellm/_redis.py +++ b/litellm/_redis.py @@ -7,13 +7,12 @@ # # Thank you users! We ❤️ you! - Krrish & Ishaan -import asyncio import inspect import json # s/o [@Frank Colson](https://www.linkedin.com/in/frank-colson-422b9b183/) for this redis implementation import os -from typing import Callable, List, Optional, Tuple, Union +from typing import Callable, List, Optional, Union import redis # type: ignore import redis.asyncio as async_redis # type: ignore @@ -108,31 +107,10 @@ def _redis_kwargs_from_environment(): return return_dict -def _generate_gcp_iam_access_token(service_account: str) -> str: - """ - Generate GCP IAM access token for Redis authentication. - - Args: - service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com' - - Returns: - Access token string for GCP IAM authentication - """ - try: - from google.cloud import iam_credentials_v1 - except ImportError: - raise ImportError( - "google-cloud-iam is required for GCP IAM Redis authentication. " - "Install it with: pip install google-cloud-iam" - ) - - client = iam_credentials_v1.IAMCredentialsClient() - request = iam_credentials_v1.GenerateAccessTokenRequest( - name=service_account, - scope=["https://www.googleapis.com/auth/cloud-platform"], - ) - response = client.generate_access_token(request=request) - return str(response.access_token) +from litellm._redis_credential_provider import ( + GCPIAMCredentialProvider, + _generate_gcp_iam_access_token, +) def create_gcp_iam_redis_connect_func( @@ -179,28 +157,6 @@ def create_gcp_iam_redis_connect_func( return iam_connect -class GCPIAMCredentialProvider: - """ - redis.credentials.CredentialProvider implementation that generates a fresh GCP IAM - token on every new connection. This fixes the 1-hour token expiry issue for async - Redis cluster clients, which previously generated the token once at startup and - cached it as a static password. - """ - - def __init__(self, gcp_service_account: str) -> None: - self._gcp_service_account = gcp_service_account - - def get_credentials(self) -> Tuple[str]: - token = _generate_gcp_iam_access_token(self._gcp_service_account) - return (token,) - - async def get_credentials_async(self) -> Tuple[str]: - token = await asyncio.to_thread( - _generate_gcp_iam_access_token, self._gcp_service_account - ) - return (token,) - - def get_redis_url_from_environment(): if "REDIS_URL" in os.environ: return os.environ["REDIS_URL"] diff --git a/litellm/_redis_credential_provider.py b/litellm/_redis_credential_provider.py new file mode 100644 index 00000000000..e2906087adf --- /dev/null +++ b/litellm/_redis_credential_provider.py @@ -0,0 +1,51 @@ +import asyncio +from typing import Tuple + + +def _generate_gcp_iam_access_token(service_account: str) -> str: + """ + Generate GCP IAM access token for Redis authentication. + + Args: + service_account: GCP service account in format 'projects/-/serviceAccounts/name@project.iam.gserviceaccount.com' + + Returns: + Access token string for GCP IAM authentication + """ + try: + from google.cloud import iam_credentials_v1 + except ImportError: + raise ImportError( + "google-cloud-iam is required for GCP IAM Redis authentication. " + "Install it with: pip install google-cloud-iam" + ) + + client = iam_credentials_v1.IAMCredentialsClient() + request = iam_credentials_v1.GenerateAccessTokenRequest( + name=service_account, + scope=["https://www.googleapis.com/auth/cloud-platform"], + ) + response = client.generate_access_token(request=request) + return str(response.access_token) + + +class GCPIAMCredentialProvider: + """ + redis.credentials.CredentialProvider implementation that generates a fresh GCP IAM + token on every new connection. This fixes the 1-hour token expiry issue for async + Redis cluster clients, which previously generated the token once at startup and + cached it as a static password. + """ + + def __init__(self, gcp_service_account: str) -> None: + self._gcp_service_account = gcp_service_account + + def get_credentials(self) -> Tuple[str]: + token = _generate_gcp_iam_access_token(self._gcp_service_account) + return (token,) + + async def get_credentials_async(self) -> Tuple[str]: + token = await asyncio.to_thread( + _generate_gcp_iam_access_token, self._gcp_service_account + ) + return (token,) diff --git a/tests/test_litellm/test_redis.py b/tests/test_litellm/test_redis.py index da0d350da2f..a683572bde2 100644 --- a/tests/test_litellm/test_redis.py +++ b/tests/test_litellm/test_redis.py @@ -5,11 +5,11 @@ import pytest import redis.asyncio as async_redis from litellm._redis import ( - GCPIAMCredentialProvider, _get_redis_cluster_kwargs, get_redis_async_client, get_redis_url_from_environment, ) +from litellm._redis_credential_provider import GCPIAMCredentialProvider def test_get_redis_url_from_environment_single_url(monkeypatch): @@ -200,7 +200,8 @@ def test_gcp_iam_credential_provider_get_credentials(): service_account = "projects/-/serviceAccounts/test@project.iam.gserviceaccount.com" with patch( - "litellm._redis._generate_gcp_iam_access_token", return_value="tok-1" + "litellm._redis_credential_provider._generate_gcp_iam_access_token", + return_value="tok-1", ) as mock_gen: provider = GCPIAMCredentialProvider(service_account) creds = provider.get_credentials() @@ -215,7 +216,8 @@ def test_gcp_iam_credential_provider_regenerates_token_on_each_call(): tokens = ["tok-1", "tok-2", "tok-3"] with patch( - "litellm._redis._generate_gcp_iam_access_token", side_effect=tokens + "litellm._redis_credential_provider._generate_gcp_iam_access_token", + side_effect=tokens, ) as mock_gen: provider = GCPIAMCredentialProvider(service_account) results = [provider.get_credentials() for _ in range(3)]