mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
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:
committed by
GitHub
Unverified
parent
ff9343d7cc
commit
37b4cfd024
@@ -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 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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user