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

* azure redis support

* small fixes

* azure managed redis sample

* fixes
This commit is contained in:
Giles Odigwe
2025-12-18 09:10:55 -08:00
committed by GitHub
Unverified
parent ff9343d7cc
commit 37b4cfd024
5 changed files with 297 additions and 16 deletions
@@ -9,6 +9,7 @@ from uuid import uuid4
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 @@ class RedisChatMessageStore:
"""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 @@ class RedisChatMessageStore:
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 required parameters
if redis_url is None:
raise ValueError("redis_url is required for Redis connection")
# 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")
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 []
@@ -93,11 +93,118 @@ class TestRedisChatMessageStore:
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"):
@@ -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 fulltext 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) peroperation thread scope where `scope_to_per_operation_thread_id=True` binds memory to a single thread for the providers 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) peroperation 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
@@ -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())
@@ -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)