Python: Add Azure Managed Redis Support with Credential Provider (#2887)

giles17 · web-flow · commit 37b4cfd024a3 · 2025-12-18T17:10:55.000Z
* azure redis support

* small fixes

* azure managed redis sample

* fixes
diff --git a/python/packages/redis/agent_framework_redis/_chat_message_store.py b/python/packages/redis/agent_framework_redis/_chat_message_store.py
@@ -9,6 +9,7 @@
 import redis.asyncio as redis
 from agent_framework import ChatMessage
 from agent_framework._serialization import SerializationMixin
+from redis.credentials import CredentialProvider
 
 
 class RedisStoreState(SerializationMixin):
@@ -55,6 +56,11 @@ class RedisChatMessageStore:
     def __init__(
         self,
         redis_url: str | None = None,
+        credential_provider: CredentialProvider | None = None,
+        host: str | None = None,
+        port: int = 6380,
+        ssl: bool = True,
+        username: str | None = None,
         thread_id: str | None = None,
         key_prefix: str = "chat_messages",
         max_messages: int | None = None,
@@ -63,12 +69,19 @@ def __init__(
         """Initialize the Redis chat message store.
 
         Creates a Redis-backed chat message store for a specific conversation thread.
-        The store will automatically create a Redis connection and manage message
-        persistence using Redis List operations.
+        Supports both traditional URL-based authentication and Azure Managed Redis
+        with credential provider.
 
         Args:
             redis_url: Redis connection URL (e.g., "redis://localhost:6379").
-                      Required for establishing Redis connection.
+                      Used for traditional authentication. Mutually exclusive with credential_provider.
+            credential_provider: Redis credential provider (redis.credentials.CredentialProvider) for
+                                Azure AD authentication. Requires host parameter. Mutually exclusive with redis_url.
+            host: Redis host name (e.g., "myredis.redis.cache.windows.net").
+                 Required when using credential_provider.
+            port: Redis port number. Defaults to 6380 (Azure Redis SSL port).
+            ssl: Enable SSL/TLS connection. Defaults to True.
+            username: Redis username. Defaults to None.
             thread_id: Unique identifier for this conversation thread.
                       If not provided, a UUID will be auto-generated.
                       This becomes part of the Redis key: {key_prefix}:{thread_id}
@@ -82,23 +95,58 @@ def __init__(
                      Useful for resuming conversations or seeding with context.
 
         Raises:
-            ValueError: If redis_url is None (Redis connection is required).
-            redis.ConnectionError: If unable to connect to Redis server.
+            ValueError: If neither redis_url nor credential_provider is provided.
+            ValueError: If both redis_url and credential_provider are provided.
+            ValueError: If credential_provider is used without host parameter.
+
+        Examples:
+            Traditional connection:
+                store = RedisChatMessageStore(
+                    redis_url="redis://localhost:6379",
+                    thread_id="conversation_123"
+                )
+
+            Azure Managed Redis with credential provider:
+                from redis.credentials import CredentialProvider
+                from azure.identity.aio import DefaultAzureCredential
+
+                store = RedisChatMessageStore(
+                    credential_provider=CredentialProvider(DefaultAzureCredential()),
+                    host="myredis.redis.cache.windows.net",
+                    thread_id="conversation_123"
+                )
+        """
+        # Validate connection parameters
+        if redis_url is None and credential_provider is None:
+            raise ValueError("Either redis_url or credential_provider must be provided")
 
+        if redis_url is not None and credential_provider is not None:
+            raise ValueError("redis_url and credential_provider are mutually exclusive")
 
-        """
-        # Validate required parameters
-        if redis_url is None:
-            raise ValueError("redis_url is required for Redis connection")
+        if credential_provider is not None and host is None:
+            raise ValueError("host is required when using credential_provider")
 
         # Store configuration
-        self.redis_url = redis_url
         self.thread_id = thread_id or f"thread_{uuid4()}"
         self.key_prefix = key_prefix
         self.max_messages = max_messages
 
-        # Initialize Redis client with connection pooling and async support
-        self._redis_client = redis.from_url(redis_url, decode_responses=True)  # type: ignore[no-untyped-call]
+        # Initialize Redis client based on authentication method
+        if credential_provider is not None and host is not None:
+            # Azure AD authentication with credential provider
+            self.redis_url = None  # Not using URL-based auth
+            self._redis_client = redis.Redis(
+                host=host,
+                port=port,
+                ssl=ssl,
+                username=username,
+                credential_provider=credential_provider,
+                decode_responses=True,
+            )
+        else:
+            # Traditional URL-based authentication
+            self.redis_url = redis_url
+            self._redis_client = redis.from_url(redis_url, decode_responses=True)  # type: ignore[no-untyped-call]
 
         # Handle initial messages (will be moved to Redis on first access)
         self._initial_messages = list(messages) if messages else []
diff --git a/python/packages/redis/tests/test_redis_chat_message_store.py b/python/packages/redis/tests/test_redis_chat_message_store.py
@@ -93,11 +93,118 @@ def test_init_with_max_messages(self):
         assert store.max_messages == 100
 
     def test_init_with_redis_url_required(self):
-        """Test that redis_url is required for initialization."""
-        with pytest.raises(ValueError, match="redis_url is required for Redis connection"):
-            # Should raise an exception since redis_url is required
+        """Test that either redis_url or credential_provider is required."""
+        with pytest.raises(ValueError, match="Either redis_url or credential_provider must be provided"):
             RedisChatMessageStore(thread_id="test123")
 
+    def test_init_with_credential_provider(self):
+        """Test initialization with credential_provider."""
+        mock_credential_provider = MagicMock()
+
+        with patch("agent_framework_redis._chat_message_store.redis.Redis") as mock_redis_class:
+            mock_redis_instance = MagicMock()
+            mock_redis_class.return_value = mock_redis_instance
+
+            store = RedisChatMessageStore(
+                credential_provider=mock_credential_provider,
+                host="myredis.redis.cache.windows.net",
+                thread_id="test123",
+            )
+
+            # Verify Redis.Redis was called with correct parameters
+            mock_redis_class.assert_called_once_with(
+                host="myredis.redis.cache.windows.net",
+                port=6380,
+                ssl=True,
+                username=None,
+                credential_provider=mock_credential_provider,
+                decode_responses=True,
+            )
+            # Verify store instance is properly initialized
+            assert store.thread_id == "test123"
+            assert store.redis_url is None  # Should be None for credential provider auth
+            assert store.key_prefix == "chat_messages"
+            assert store.max_messages is None
+
+    def test_init_with_credential_provider_custom_port(self):
+        """Test initialization with credential_provider and custom port."""
+        mock_credential_provider = MagicMock()
+
+        with patch("agent_framework_redis._chat_message_store.redis.Redis") as mock_redis_class:
+            mock_redis_instance = MagicMock()
+            mock_redis_class.return_value = mock_redis_instance
+
+            store = RedisChatMessageStore(
+                credential_provider=mock_credential_provider,
+                host="myredis.redis.cache.windows.net",
+                port=6379,
+                ssl=False,
+                username="admin",
+                thread_id="test123",
+            )
+
+            # Verify custom parameters were passed
+            mock_redis_class.assert_called_once_with(
+                host="myredis.redis.cache.windows.net",
+                port=6379,
+                ssl=False,
+                username="admin",
+                credential_provider=mock_credential_provider,
+                decode_responses=True,
+            )
+            # Verify store instance is properly initialized
+            assert store.thread_id == "test123"
+            assert store.redis_url is None  # Should be None for credential provider auth
+            assert store.key_prefix == "chat_messages"
+
+    def test_init_credential_provider_requires_host(self):
+        """Test that credential_provider requires host parameter."""
+        mock_credential_provider = MagicMock()
+
+        with pytest.raises(ValueError, match="host is required when using credential_provider"):
+            RedisChatMessageStore(
+                credential_provider=mock_credential_provider,
+                thread_id="test123",
+            )
+
+    def test_init_mutually_exclusive_params(self):
+        """Test that redis_url and credential_provider are mutually exclusive."""
+        mock_credential_provider = MagicMock()
+
+        with pytest.raises(ValueError, match="redis_url and credential_provider are mutually exclusive"):
+            RedisChatMessageStore(
+                redis_url="redis://localhost:6379",
+                credential_provider=mock_credential_provider,
+                host="myredis.redis.cache.windows.net",
+                thread_id="test123",
+            )
+
+    async def test_serialize_with_credential_provider(self):
+        """Test that serialization works correctly with credential provider authentication."""
+        mock_credential_provider = MagicMock()
+
+        with patch("agent_framework_redis._chat_message_store.redis.Redis") as mock_redis_class:
+            mock_redis_instance = MagicMock()
+            mock_redis_class.return_value = mock_redis_instance
+
+            store = RedisChatMessageStore(
+                credential_provider=mock_credential_provider,
+                host="myredis.redis.cache.windows.net",
+                thread_id="test123",
+                key_prefix="custom_prefix",
+                max_messages=100,
+            )
+
+            # Serialize the store state
+            state = await store.serialize()
+
+            # Verify serialization includes correct values
+            assert state["thread_id"] == "test123"
+            assert state["redis_url"] is None  # Should be None for credential provider auth
+            assert state["key_prefix"] == "custom_prefix"
+            assert state["max_messages"] == 100
+            assert state["type"] == "redis_store_state"
+
     def test_init_with_initial_messages(self, sample_messages):
         """Test initialization with initial messages."""
         with patch("agent_framework_redis._chat_message_store.redis.from_url"):
diff --git a/python/samples/getting_started/context_providers/redis/README.md b/python/samples/getting_started/context_providers/redis/README.md
@@ -8,8 +8,11 @@ This folder contains an example demonstrating how to use the Redis context provi
 
 | File | Description |
 |------|-------------|
+| [`azure_redis_conversation.py`](azure_redis_conversation.py) | Demonstrates conversation persistence with RedisChatMessageStore and Azure Redis with Azure AD (Entra ID) authentication using credential provider. |
 | [`redis_basics.py`](redis_basics.py) | Shows standalone provider usage and agent integration. Demonstrates writing messages to Redis, retrieving context via full‑text or hybrid vector search, and persisting preferences across threads. Also includes a simple tool example whose outputs are remembered. |
-| [`redis_threads.py`](redis_threads.py) | Demonstrates thread scoping. Includes: (1) global thread scope with a fixed `thread_id` shared across operations; (2) per‑operation thread scope where `scope_to_per_operation_thread_id=True` binds memory to a single thread for the provider’s lifetime; and (3) multiple agents with isolated memory via different `agent_id` values. |
+| [`redis_conversation.py`](redis_conversation.py) | Simple example showing conversation persistence with RedisChatMessageStore using traditional connection string authentication. |
+| [`redis_threads.py`](redis_threads.py) | Demonstrates thread scoping. Includes: (1) global thread scope with a fixed `thread_id` shared across operations; (2) per‑operation thread scope where `scope_to_per_operation_thread_id=True` binds memory to a single thread for the provider's lifetime; and (3) multiple agents with isolated memory via different `agent_id` values. |
+
 
 ## Prerequisites
 
diff --git a/python/samples/getting_started/context_providers/redis/azure_redis_conversation.py b/python/samples/getting_started/context_providers/redis/azure_redis_conversation.py
@@ -0,0 +1,123 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""Azure Managed Redis Chat Message Store with Azure AD Authentication
+
+This example demonstrates how to use Azure Managed Redis with Azure AD authentication
+to persist conversational details using RedisChatMessageStore.
+
+Requirements:
+  - Azure Managed Redis instance with Azure AD authentication enabled
+  - Azure credentials configured (az login or managed identity)
+  - agent-framework-redis: pip install agent-framework-redis
+  - azure-identity: pip install azure-identity
+
+Environment Variables:
+  - AZURE_REDIS_HOST: Your Azure Managed Redis host (e.g., myredis.redis.cache.windows.net)
+  - OPENAI_API_KEY: Your OpenAI API key
+  - OPENAI_CHAT_MODEL_ID: OpenAI model (e.g., gpt-4o-mini)
+  - AZURE_USER_OBJECT_ID: Your Azure AD User Object ID for authentication
+"""
+
+import asyncio
+import os
+
+from agent_framework.openai import OpenAIChatClient
+from agent_framework.redis import RedisChatMessageStore
+from azure.identity.aio import AzureCliCredential
+from redis.credentials import CredentialProvider
+
+
+class AzureCredentialProvider(CredentialProvider):
+    """Credential provider for Azure AD authentication with Redis Enterprise."""
+
+    def __init__(self, azure_credential: AzureCliCredential, user_object_id: str):
+        self.azure_credential = azure_credential
+        self.user_object_id = user_object_id
+
+    async def get_credentials_async(self) -> tuple[str] | tuple[str, str]:
+        """Get Azure AD token for Redis authentication.
+
+        Returns (username, token) where username is the Azure user's Object ID.
+        """
+        token = await self.azure_credential.get_token("https://redis.azure.com/.default")
+        return (self.user_object_id, token.token)
+
+
+async def main() -> None:
+    redis_host = os.environ.get("AZURE_REDIS_HOST")
+    if not redis_host:
+        print("ERROR: Set AZURE_REDIS_HOST environment variable")
+        return
+
+    # For Azure Redis with Entra ID, username must be your Object ID
+    user_object_id = os.environ.get("AZURE_USER_OBJECT_ID")
+    if not user_object_id:
+        print("ERROR: Set AZURE_USER_OBJECT_ID environment variable")
+        print("Get your Object ID from the Azure Portal")
+        return
+
+    # Create Azure CLI credential provider (uses 'az login' credentials)
+    azure_credential = AzureCliCredential()
+    credential_provider = AzureCredentialProvider(azure_credential, user_object_id)
+
+    thread_id = "azure_test_thread"
+
+    # Factory for creating Azure Redis chat message store
+    chat_message_store_factory = lambda: RedisChatMessageStore(
+        credential_provider=credential_provider,
+        host=redis_host,
+        port=10000,
+        ssl=True,
+        thread_id=thread_id,
+        key_prefix="chat_messages",
+        max_messages=100,
+    )
+
+    # Create chat client
+    client = OpenAIChatClient()
+
+    # Create agent with Azure Redis store
+    agent = client.create_agent(
+        name="AzureRedisAssistant",
+        instructions="You are a helpful assistant.",
+        chat_message_store_factory=chat_message_store_factory,
+    )
+
+    # Conversation
+    query = "Remember that I enjoy gumbo"
+    result = await agent.run(query)
+    print("User: ", query)
+    print("Agent: ", result)
+
+    # Ask the agent to recall the stored preference; it should retrieve from memory
+    query = "What do I enjoy?"
+    result = await agent.run(query)
+    print("User: ", query)
+    print("Agent: ", result)
+
+    query = "What did I say to you just now?"
+    result = await agent.run(query)
+    print("User: ", query)
+    print("Agent: ", result)
+
+    query = "Remember that I have a meeting at 3pm tomorrow"
+    result = await agent.run(query)
+    print("User: ", query)
+    print("Agent: ", result)
+
+    query = "Tulips are red"
+    result = await agent.run(query)
+    print("User: ", query)
+    print("Agent: ", result)
+
+    query = "What was the first thing I said to you this conversation?"
+    result = await agent.run(query)
+    print("User: ", query)
+    print("Agent: ", result)
+
+    # Cleanup
+    await azure_credential.close()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())
diff --git a/python/samples/getting_started/context_providers/redis/redis_conversation.py b/python/samples/getting_started/context_providers/redis/redis_conversation.py
@@ -91,7 +91,7 @@ async def main() -> None:
     print("User: ", query)
     print("Agent: ", result)
 
-    query = "Remember that anyone who does not clean shrimp will be eaten by a shark"
+    query = "Remember that I have a meeting at 3pm tomorro"
     result = await agent.run(query)
     print("User: ", query)
     print("Agent: ", result)
