Python: [BREAKING] PR2 — Wire context provider pipeline, remove old types, update all consumers (#3850)

* PR2: Wire context provider pipeline and update all internal consumers

- Replace AgentThread with AgentSession across all packages
- Replace ContextProvider with BaseContextProvider across all packages
- Replace context_provider param with context_providers (Sequence)
- Replace thread= with session= in run() signatures
- Replace get_new_thread() with create_session()
- Add get_session(service_session_id) to agent interface
- DurableAgentThread -> DurableAgentSession
- Remove _notify_thread_of_new_messages from WorkflowAgent
- Wire before_run/after_run context provider pipeline in RawAgent
- Auto-inject InMemoryHistoryProvider when no providers configured

* fix: update all tests for context provider pipeline, fix lazy-loaders, remove old test files

* refactor: update all sample files for context provider pipeline (AgentThread→AgentSession, ContextProvider→BaseContextProvider)

* fix: update remaining ag-ui references (client docstring, getting_started sample)

* fix: make get_session service_session_id keyword-only to avoid confusion with session_id

* refactor: rename _RunContext.thread_messages to session_messages

* refactor: remove _threads.py, _memory.py, and old provider files; migrate devui to use plain message lists

* rename: remove _new_ prefix from test files

* refactor: rewrite SlidingWindowChatMessageStore as SlidingWindowHistoryProvider(InMemoryHistoryProvider)

* fix: read full history from session state directly instead of reaching into provider internals

* fix: update stale .pyi stubs, sample imports, and README references for new provider types

* fix: remove stale message_store, _notify_thread_of_new_messages, and session_id.key references in samples

* refactor: merge context_providers and sessions sample folders into sessions, remove aggregate_context_provider

* refactor: UserInfoMemory stores state in session.state instead of instance attributes

* feat: add Pydantic BaseModel support to session state serialization

Pydantic models stored in session.state are now automatically serialized
via model_dump() and restored via model_validate() during to_dict()/from_dict()
round-trips. Models are auto-registered on first serialization; use
register_state_type() for cold-start deserialization.

Also export register_state_type as a public API.

* fix mem0

* Update sample README links and descriptions for session terminology

- Replace 'thread' with 'session' in sample descriptions across all READMEs
- Update file links for renamed samples (mem0_sessions, redis_sessions, etc.)
- Fix Threads section → Sessions section in main samples/README.md
- Update tools, middleware, workflows, durabletask, azure_functions READMEs
- Update architecture diagrams in concepts/tools/README.md
- Update migration guides (autogen, semantic-kernel)

* Fix broken Redis README link to renamed sample

* Fix Mem0 OSS client search: pass scoping params as direct kwargs

AsyncMemory (OSS) expects user_id/agent_id/run_id as direct kwargs,
while AsyncMemoryClient (Platform) expects them in a filters dict.
Adds tests for both client types.

Port of fix from #3844 to new Mem0ContextProvider.

* Fix rebase issues: restore missing _conversation_state.py and checkpoint decode logic

- Add back _conversation_state.py (encode/decode_chat_messages) lost in rebase
- Fix on_checkpoint_restore to decode cache/conversation with decode_chat_messages
- Fix on_checkpoint_restore to use decode_checkpoint_value for pending requests
- Add tests/workflow/__init__.py for relative import support
- Fix test_agent_executor checkpoint selection (checkpoints[1] not superstep)

* Add STORES_BY_DEFAULT ClassVar to skip redundant InMemoryHistoryProvider injection

Chat clients that store history server-side by default (OpenAI Responses API,
Azure AI Agent) now declare STORES_BY_DEFAULT = True. The agent checks this
during auto-injection and skips InMemoryHistoryProvider unless the user
explicitly sets store=False.

* Fix broken markdown links in azure_ai and redis READMEs

* Fix getting-started samples to use session API instead of removed thread/ContextProvider API

* updates to workflow as agent

* fix group chat import

* Rename Thread→Session throughout, fix service_session_id propagation, remove stale AGUIThread

- Fix: Propagate conversation_id from ChatResponse back to session.service_session_id
  in both streaming and non-streaming paths in _agents.py
- Rename AgentThreadException → AgentSessionException
- Remove stale AGUIThread from ag_ui lazy-loader
- Rename use_service_thread → use_service_session in ag-ui package
- Rename test functions from *_thread_* to *_session_*
- Rename sample files from *_thread* to *_session*
- Update docstrings and comments: thread → session
- Update _mcp.py kwargs filter: add 'session' alongside 'thread'
- Fix ContinuationToken docstring example: thread=thread → session=session
- Fix _clients.py docstring: 'Agent threads' → 'Agent sessions'

* Fix broken markdown links after thread→session file renames

* fix azure ai test
This commit is contained in:
Eduard van Valkenburg
2026-02-12 22:00:32 +01:00
committed by GitHub
Unverified
parent 0c67dbbce5
commit 1e350ea22f
312 changed files with 6669 additions and 11423 deletions
+2 -2
View File
@@ -30,10 +30,10 @@ The `RedisChatMessageStore` provides persistent conversation storage using Redis
#### Basic Usage Examples
See the complete [Redis chat message store examples](../../samples/02-agents/conversations/redis_chat_message_store_thread.py) including:
See the complete [Redis history provider examples](../../samples/02-agents/conversations/redis_chat_message_store_session.py) including:
- User session management
- Conversation persistence across restarts
- Thread serialization and deserialization
- Session serialization and deserialization
- Automatic message trimming
- Error handling patterns
@@ -1,10 +1,8 @@
# Copyright (c) Microsoft. All rights reserved.
import importlib.metadata
from ._chat_message_store import RedisChatMessageStore
from ._context_provider import _RedisContextProvider
from ._history_provider import _RedisHistoryProvider
from ._provider import RedisProvider
from ._context_provider import RedisContextProvider
from ._history_provider import RedisHistoryProvider
try:
__version__ = importlib.metadata.version(__name__)
@@ -12,9 +10,7 @@ except importlib.metadata.PackageNotFoundError:
__version__ = "0.0.0" # Fallback for development mode
__all__ = [
"RedisChatMessageStore",
"RedisProvider",
"_RedisContextProvider",
"_RedisHistoryProvider",
"RedisContextProvider",
"RedisHistoryProvider",
"__version__",
]
@@ -1,595 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
from collections.abc import Sequence
from typing import Any
from uuid import uuid4
import redis.asyncio as redis
from agent_framework import Message
from agent_framework._serialization import SerializationMixin
from redis.credentials import CredentialProvider
class RedisStoreState(SerializationMixin):
"""State model for serializing and deserializing Redis chat message store data."""
def __init__(
self,
thread_id: str,
redis_url: str | None = None,
key_prefix: str = "chat_messages",
max_messages: int | None = None,
) -> None:
"""State model for serializing and deserializing Redis chat message store data."""
self.thread_id = thread_id
self.redis_url = redis_url
self.key_prefix = key_prefix
self.max_messages = max_messages
class RedisChatMessageStore:
"""Redis-backed implementation of ChatMessageStoreProtocol using Redis Lists.
This implementation provides persistent, thread-safe chat message storage using Redis Lists.
Messages are stored as JSON-serialized strings in chronological order, with each conversation
thread isolated by a unique Redis key.
Key Features:
============
- **Persistent Storage**: Messages survive application restarts and crashes
- **Thread Isolation**: Each conversation thread has its own Redis key namespace
- **Auto Message Limits**: Configurable automatic trimming of old messages using LTRIM
- **Performance Optimized**: Uses native Redis operations for efficiency
- **State Serialization**: Full compatibility with Agent Framework thread serialization
- **Initial Message Support**: Pre-load conversations with existing message history
- **Production Ready**: Atomic operations, error handling, connection pooling
Redis Operations:
- RPUSH: Add messages to the end of the list (chronological order)
- LRANGE: Retrieve messages in chronological order
- LTRIM: Maintain message limits by trimming old messages
- DELETE: Clear all messages for a thread
"""
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,
messages: Sequence[Message] | None = None,
) -> None:
"""Initialize the Redis chat message store.
Creates a Redis-backed chat message store for a specific conversation thread.
Supports both traditional URL-based authentication and Azure Managed Redis
with credential provider.
Args:
redis_url: Redis connection URL (e.g., "redis://localhost:6379").
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}
key_prefix: Prefix for Redis keys to namespace different applications.
Defaults to 'chat_messages'. Useful for multi-tenant scenarios.
max_messages: Maximum number of messages to retain in Redis.
When exceeded, oldest messages are automatically trimmed using LTRIM.
None means unlimited storage.
messages: Initial messages to pre-populate the conversation.
These are added to Redis on first access if the Redis key is empty.
Useful for resuming conversations or seeding with context.
Raises:
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")
if credential_provider is not None and host is None:
raise ValueError("host is required when using credential_provider")
# Store configuration
self.thread_id = thread_id or f"thread_{uuid4()}"
self.key_prefix = key_prefix
self.max_messages = max_messages
# 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 []
self._initial_messages_added = False
@property
def redis_key(self) -> str:
"""Get the Redis key for this thread's messages.
The key format is: {key_prefix}:{thread_id}
Returns:
Redis key string used for storing this thread's messages.
Example:
For key_prefix="chat_messages" and thread_id="user_123_session_456":
Returns "chat_messages:user_123_session_456"
"""
return f"{self.key_prefix}:{self.thread_id}"
async def _ensure_initial_messages_added(self) -> None:
"""Ensure initial messages are added to Redis if not already present.
This method is called before any Redis operations to guarantee that
initial messages provided during construction are persisted to Redis.
"""
if not self._initial_messages or self._initial_messages_added:
return
# Check if Redis key already has messages (prevents duplicate additions)
existing_count = await self._redis_client.llen(self.redis_key) # type: ignore[misc] # type: ignore[misc]
if existing_count == 0:
# Add initial messages using atomic pipeline operation
await self._add_redis_messages(self._initial_messages)
# Mark as completed and free memory
self._initial_messages_added = True
self._initial_messages.clear()
async def _add_redis_messages(self, messages: Sequence[Message]) -> None:
"""Add multiple messages to Redis using atomic pipeline operation.
This internal method efficiently adds multiple messages to the Redis list
using a single atomic transaction to ensure consistency.
Args:
messages: Sequence of Message objects to add to Redis.
"""
if not messages:
return
# Pre-serialize all messages for efficient pipeline operation
serialized_messages = [self._serialize_message(message) for message in messages]
# Use Redis pipeline for atomic batch operation
async with self._redis_client.pipeline(transaction=True) as pipe:
for serialized_message in serialized_messages:
await pipe.rpush(self.redis_key, serialized_message) # type: ignore[misc]
await pipe.execute()
async def add_messages(self, messages: Sequence[Message]) -> None:
"""Add messages to the Redis store (ChatMessageStoreProtocol protocol method).
This method implements the required ChatMessageStoreProtocol protocol for adding messages.
Messages are appended to the Redis list in chronological order, with automatic
trimming if message limits are configured.
Args:
messages: Sequence of Message objects to add to the store.
Can be empty (no-op) or contain multiple messages.
Thread Safety:
- Atomic pipeline ensures all messages are added together
- LTRIM operation is atomic for consistent message limits
Example:
.. code-block:: python
messages = [Message(role="user", text="Hello"), Message(role="assistant", text="Hi there!")]
await store.add_messages(messages)
"""
if not messages:
return
# Ensure any initial messages are persisted first
await self._ensure_initial_messages_added()
# Add new messages using atomic pipeline operation
await self._add_redis_messages(messages)
# Apply message limit if configured (automatic cleanup)
if self.max_messages is not None:
current_count = await self._redis_client.llen(self.redis_key) # type: ignore[misc]
if current_count > self.max_messages:
# Keep only the most recent max_messages using LTRIM
await self._redis_client.ltrim(self.redis_key, -self.max_messages, -1) # type: ignore[misc]
async def list_messages(self) -> list[Message]:
"""Get all messages from the store in chronological order (ChatMessageStoreProtocol protocol method).
This method implements the required ChatMessageStoreProtocol protocol for retrieving messages.
Returns all messages stored in Redis, ordered from oldest (index 0) to newest (index -1).
Returns:
List of Message objects in chronological order (oldest first).
Returns empty list if no messages exist or if Redis connection fails.
Example:
.. code-block:: python
# Get all conversation history
messages = await store.list_messages()
"""
# Ensure any initial messages are persisted to Redis first
await self._ensure_initial_messages_added()
messages = []
# Retrieve all messages from Redis list (oldest to newest)
redis_messages = await self._redis_client.lrange(self.redis_key, 0, -1) # type: ignore[misc]
if redis_messages:
for serialized_message in redis_messages:
# Deserialize each JSON message back to Message
message = self._deserialize_message(serialized_message)
messages.append(message)
return messages
async def serialize(self, **kwargs: Any) -> Any:
"""Serialize the current store state for persistence (ChatMessageStoreProtocol protocol method).
This method implements the required ChatMessageStoreProtocol protocol for state serialization.
Captures the Redis connection configuration and thread information needed to
reconstruct the store and reconnect to the same conversation data.
Keyword Args:
**kwargs: Additional arguments passed to Pydantic model_dump() for serialization.
Common options: exclude_none=True, by_alias=True
Returns:
Dictionary containing serialized store configuration that can be persisted
to databases, files, or other storage mechanisms.
"""
state = RedisStoreState(
thread_id=self.thread_id,
redis_url=self.redis_url,
key_prefix=self.key_prefix,
max_messages=self.max_messages,
)
return state.to_dict(exclude_none=False, **kwargs)
@classmethod
async def deserialize(cls, serialized_store_state: Any, **kwargs: Any) -> RedisChatMessageStore:
"""Deserialize state data into a new store instance (ChatMessageStoreProtocol protocol method).
This method implements the required ChatMessageStoreProtocol protocol for state deserialization.
Creates a new RedisChatMessageStore instance from previously serialized data,
allowing the store to reconnect to the same conversation data in Redis.
Args:
serialized_store_state: Previously serialized state data from serialize_state().
Should be a dictionary with thread_id, redis_url, etc.
Keyword Args:
**kwargs: Additional arguments passed to Pydantic model validation.
Returns:
A new RedisChatMessageStore instance configured from the serialized state.
Raises:
ValueError: If required fields are missing or invalid in the serialized state.
"""
if not serialized_store_state:
raise ValueError("serialized_store_state is required for deserialization")
# Validate and parse the serialized state using Pydantic
state = RedisStoreState.from_dict(serialized_store_state, **kwargs)
# Create and return a new store instance with the deserialized configuration
return cls(
redis_url=state.redis_url,
thread_id=state.thread_id,
key_prefix=state.key_prefix,
max_messages=state.max_messages,
)
async def update_from_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
"""Deserialize state data into this store instance (ChatMessageStoreProtocol protocol method).
This method implements the required ChatMessageStoreProtocol protocol for state deserialization.
Restores the store configuration from previously serialized data, allowing the store
to reconnect to the same conversation data in Redis.
Args:
serialized_store_state: Previously serialized state data from serialize_state().
Should be a dictionary with thread_id, redis_url, etc.
Keyword Args:
**kwargs: Additional arguments passed to Pydantic model validation.
"""
if not serialized_store_state:
return
# Validate and parse the serialized state using Pydantic
state = RedisStoreState.from_dict(serialized_store_state, **kwargs)
# Update store configuration from deserialized state
self.thread_id = state.thread_id
if state.redis_url is not None:
self.redis_url = state.redis_url
self.key_prefix = state.key_prefix
self.max_messages = state.max_messages
# Recreate Redis client if the URL changed
if state.redis_url and state.redis_url != getattr(self, "_last_redis_url", None):
self._redis_client = redis.from_url(state.redis_url, decode_responses=True) # type: ignore[no-untyped-call]
self._last_redis_url = state.redis_url
# Reset initial message state since we're connecting to existing data
self._initial_messages_added = False
async def clear(self) -> None:
"""Remove all messages from the store.
Permanently deletes all messages for this conversation thread by removing
the Redis key. This operation cannot be undone.
Warning:
- This permanently deletes all conversation history
- Consider exporting messages before clearing if backup is needed
Example:
.. code-block:: python
# Clear conversation history
await store.clear()
# Verify messages are gone
messages = await store.list_messages()
assert len(messages) == 0
"""
await self._redis_client.delete(self.redis_key)
def _serialize_message(self, message: Message) -> str:
"""Serialize a Message to JSON string.
Args:
message: Message to serialize.
Returns:
JSON string representation of the message.
"""
# Serialize to compact JSON (no extra whitespace for Redis efficiency)
return message.to_json(separators=(",", ":"))
def _deserialize_message(self, serialized_message: str) -> Message:
"""Deserialize a JSON string to Message.
Args:
serialized_message: JSON string representation of a message.
Returns:
Message object.
"""
# Reconstruct Message using custom deserialization
return Message.from_json(serialized_message)
# ============================================================================
# List-like Convenience Methods (Redis-optimized async versions)
# ============================================================================
def __bool__(self) -> bool:
"""Return True since the store always exists once created.
This method is called by Python's truthiness checks (if store:).
Since a RedisChatMessageStore instance always represents a valid store,
this always returns True.
Returns:
Always True - the store exists and is ready for operations.
Note:
This is used by the Agent Framework to check if a message store
is configured: `if thread.message_store:`
"""
return True
async def __len__(self) -> int:
"""Return the number of messages in the Redis store.
Provides efficient message counting using Redis LLEN command.
This is the async equivalent of Python's built-in len() function.
Returns:
The count of messages currently stored in Redis.
"""
await self._ensure_initial_messages_added()
return await self._redis_client.llen(self.redis_key) # type: ignore[misc,no-any-return]
async def getitem(self, index: int) -> Message:
"""Get a message by index using Redis LINDEX.
Args:
index: The index of the message to retrieve.
Returns:
The Message at the specified index.
Raises:
IndexError: If the index is out of range.
"""
await self._ensure_initial_messages_added()
# Use Redis LINDEX for efficient single-item access
serialized_message = await self._redis_client.lindex(self.redis_key, index) # type: ignore[misc]
if serialized_message is None:
raise IndexError("list index out of range")
return self._deserialize_message(serialized_message)
async def setitem(self, index: int, item: Message) -> None:
"""Set a message at the specified index using Redis LSET.
Args:
index: The index at which to set the message.
item: The Message to set at the specified index.
Raises:
IndexError: If the index is out of range.
"""
await self._ensure_initial_messages_added()
# Validate index exists using LLEN
current_count = await self._redis_client.llen(self.redis_key) # type: ignore[misc]
if index < 0:
index = current_count + index
if index < 0 or index >= current_count:
raise IndexError("list index out of range")
# Use Redis LSET for efficient single-item update
serialized_message = self._serialize_message(item)
await self._redis_client.lset(self.redis_key, index, serialized_message) # type: ignore[misc]
async def append(self, item: Message) -> None:
"""Append a message to the end of the store.
Args:
item: The Message to append.
"""
await self.add_messages([item])
async def count(self) -> int:
"""Return the number of messages in the Redis store.
Returns:
The count of messages currently stored in Redis.
"""
await self._ensure_initial_messages_added()
return await self._redis_client.llen(self.redis_key) # type: ignore[misc,no-any-return]
async def index(self, item: Message) -> int:
"""Return the index of the first occurrence of the specified message.
Uses Redis LINDEX to iterate through the list without loading all messages.
Still O(N) but more memory efficient for large lists.
Args:
item: The Message to find.
Returns:
The index of the first occurrence of the message.
Raises:
ValueError: If the message is not found in the store.
"""
await self._ensure_initial_messages_added()
target_serialized = self._serialize_message(item)
list_length = await self._redis_client.llen(self.redis_key) # type: ignore[misc]
# Iterate through Redis list using LINDEX
for i in range(list_length):
redis_message = await self._redis_client.lindex(self.redis_key, i) # type: ignore[misc]
if redis_message == target_serialized:
return i
raise ValueError("Message not found in store")
async def remove(self, item: Message) -> None:
"""Remove the first occurrence of the specified message from the store.
Uses Redis LREM command for efficient removal by value.
O(N) but performed natively in Redis without data transfer.
Args:
item: The Message to remove.
Raises:
ValueError: If the message is not found in the store.
"""
await self._ensure_initial_messages_added()
# Serialize the message to match Redis storage format
target_serialized = self._serialize_message(item)
# Use LREM to remove first occurrence (count=1)
removed_count = await self._redis_client.lrem(self.redis_key, 1, target_serialized) # type: ignore[misc]
if removed_count == 0:
raise ValueError("Message not found in store")
async def extend(self, items: Sequence[Message]) -> None:
"""Extend the store by appending all messages from the iterable.
Args:
items: Sequence of Message objects to append.
"""
await self.add_messages(items)
async def ping(self) -> bool:
"""Test the Redis connection.
Returns:
True if the connection is successful, False otherwise.
"""
try:
await self._redis_client.ping() # type: ignore[misc]
return True
except Exception:
return False
async def aclose(self) -> None:
"""Close the Redis connection.
This method provides a clean way to close the underlying Redis connection
when the store is no longer needed. This is particularly useful in samples
and applications where explicit resource cleanup is desired.
"""
await self._redis_client.aclose() # type: ignore[misc]
def __repr__(self) -> str:
"""String representation of the store."""
return (
f"RedisChatMessageStore(thread_id='{self.thread_id}', "
f"redis_key='{self.redis_key}', max_messages={self.max_messages})"
)
@@ -2,9 +2,8 @@
"""New-pattern Redis context provider using BaseContextProvider.
This module provides ``_RedisContextProvider``, a side-by-side implementation of
:class:`RedisProvider` built on the new :class:`BaseContextProvider` hooks pattern.
It will be renamed to ``RedisContextProvider`` in PR2 when the old class is removed.
This module provides ``RedisContextProvider``, built on the new
:class:`BaseContextProvider` hooks pattern.
"""
from __future__ import annotations
@@ -43,17 +42,11 @@ if TYPE_CHECKING:
from agent_framework._agents import SupportsAgentRun
class _RedisContextProvider(BaseContextProvider):
class RedisContextProvider(BaseContextProvider):
"""Redis context provider using the new BaseContextProvider hooks pattern.
Stores context in Redis and retrieves scoped context via full-text or
optional hybrid vector search. This is the new-pattern equivalent of
:class:`RedisProvider`.
Note:
This class uses a temporary ``_`` prefix to coexist with the existing
:class:`RedisProvider`. It will be renamed to ``RedisContextProvider``
in PR2.
optional hybrid vector search.
"""
DEFAULT_CONTEXT_PROMPT = "## Memories\nConsider the following memories when answering user questions:"
@@ -429,4 +422,4 @@ class _RedisContextProvider(BaseContextProvider):
"""Async context manager exit."""
__all__ = ["_RedisContextProvider"]
__all__ = ["RedisContextProvider"]
@@ -2,9 +2,8 @@
"""New-pattern Redis history provider using BaseHistoryProvider.
This module provides ``_RedisHistoryProvider``, a side-by-side implementation of
:class:`RedisMessageStore` built on the new :class:`BaseHistoryProvider` hooks pattern.
It will be renamed to ``RedisHistoryProvider`` in PR2 when the old class is removed.
This module provides ``RedisHistoryProvider``, built on the new
:class:`BaseHistoryProvider` hooks pattern.
"""
from __future__ import annotations
@@ -18,17 +17,11 @@ from agent_framework._sessions import BaseHistoryProvider
from redis.credentials import CredentialProvider
class _RedisHistoryProvider(BaseHistoryProvider):
class RedisHistoryProvider(BaseHistoryProvider):
"""Redis-backed history provider using the new BaseHistoryProvider hooks pattern.
Stores conversation history in Redis Lists, with each session isolated by a
unique Redis key. This is the new-pattern equivalent of
:class:`RedisMessageStore`.
Note:
This class uses a temporary ``_`` prefix to coexist with the existing
:class:`RedisMessageStore`. It will be renamed to ``RedisHistoryProvider``
in PR2.
unique Redis key.
"""
def __init__(
@@ -181,4 +174,4 @@ class _RedisHistoryProvider(BaseHistoryProvider):
await self._redis_client.aclose() # type: ignore[misc]
__all__ = ["_RedisHistoryProvider"]
__all__ = ["RedisHistoryProvider"]
@@ -1,595 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import json
import sys
from collections.abc import MutableSequence, Sequence
from functools import reduce
from operator import and_
from typing import Any, Literal, cast
import numpy as np
from agent_framework import Context, ContextProvider, Message
from agent_framework.exceptions import (
AgentException,
ServiceInitializationError,
ServiceInvalidRequestError,
)
from redisvl.index import AsyncSearchIndex
from redisvl.query import FilterQuery, HybridQuery, TextQuery
from redisvl.query.filter import FilterExpression, Tag
from redisvl.utils.token_escaper import TokenEscaper
from redisvl.utils.vectorize import BaseVectorizer
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
else:
from typing_extensions import Self # pragma: no cover
if sys.version_info >= (3, 12):
from typing import override # type: ignore # pragma: no cover
else:
from typing_extensions import override # type: ignore[import] # pragma: no cover
class RedisProvider(ContextProvider):
"""Redis context provider with dynamic, filterable schema.
Stores context in Redis and retrieves scoped context.
Uses full-text or optional hybrid vector search to ground model responses.
"""
def __init__(
self,
redis_url: str = "redis://localhost:6379",
index_name: str = "context",
prefix: str = "context",
# Redis vectorizer configuration (optional, injected by client)
redis_vectorizer: BaseVectorizer | None = None,
vector_field_name: str | None = None,
vector_algorithm: Literal["flat", "hnsw"] | None = None,
vector_distance_metric: Literal["cosine", "ip", "l2"] | None = None,
# Partition fields (indexed for filtering)
application_id: str | None = None,
agent_id: str | None = None,
user_id: str | None = None,
thread_id: str | None = None,
scope_to_per_operation_thread_id: bool = False,
# Prompt and runtime
context_prompt: str = ContextProvider.DEFAULT_CONTEXT_PROMPT,
redis_index: Any = None,
overwrite_index: bool = False,
):
"""Create a Redis Context Provider.
Args:
redis_url: The Redis server URL.
index_name: The name of the Redis index.
prefix: The prefix for all keys in the Redis database.
redis_vectorizer: The vectorizer to use for Redis.
vector_field_name: The name of the vector field in Redis.
vector_algorithm: The algorithm to use for vector search.
vector_distance_metric: The distance metric to use for vector search.
application_id: The application ID to scope the context.
agent_id: The agent ID to scope the context.
user_id: The user ID to scope the context.
thread_id: The thread ID to scope the context.
scope_to_per_operation_thread_id: Whether to scope to the per-operation thread ID.
context_prompt: The context prompt to use for the provider.
redis_index: The Redis index to use for the provider.
overwrite_index: Whether to overwrite the existing Redis index.
"""
self.redis_url = redis_url
self.index_name = index_name
self.prefix = prefix
if redis_vectorizer is not None and not isinstance(redis_vectorizer, BaseVectorizer):
raise AgentException(
f"The redis vectorizer is not a valid type, got: {type(redis_vectorizer)}, expected: BaseVectorizer."
)
self.redis_vectorizer = redis_vectorizer
self.vector_field_name = vector_field_name
self.vector_algorithm: Literal["flat", "hnsw"] | None = vector_algorithm
self.vector_distance_metric: Literal["cosine", "ip", "l2"] | None = vector_distance_metric
self.application_id = application_id
self.agent_id = agent_id
self.user_id = user_id
self.thread_id = thread_id
self.scope_to_per_operation_thread_id = scope_to_per_operation_thread_id
self.context_prompt = context_prompt
self.overwrite_index = overwrite_index
self._per_operation_thread_id: str | None = None
self._token_escaper: TokenEscaper = TokenEscaper()
self._conversation_id: str | None = None
self._index_initialized: bool = False
self._schema_dict: dict[str, Any] | None = None
self.redis_index = redis_index or AsyncSearchIndex.from_dict(
self.schema_dict, redis_url=self.redis_url, validate_on_load=True
)
@property
def schema_dict(self) -> dict[str, Any]:
"""Get the Redis schema dictionary, computing and caching it on first access."""
if self._schema_dict is None:
# Get vector configuration from vectorizer if available
vector_dims = self.redis_vectorizer.dims if self.redis_vectorizer is not None else None
vector_datatype = self.redis_vectorizer.dtype if self.redis_vectorizer is not None else None
self._schema_dict = self._build_schema_dict(
index_name=self.index_name,
prefix=self.prefix,
vector_field_name=self.vector_field_name,
vector_dims=vector_dims,
vector_datatype=vector_datatype,
vector_algorithm=self.vector_algorithm,
vector_distance_metric=self.vector_distance_metric,
)
return self._schema_dict
def _build_filter_from_dict(self, filters: dict[str, str | None]) -> Any | None:
"""Builds a combined filter expression from simple equality tags.
This ANDs non-empty tag filters and is used to scope all operations to app/agent/user/thread partitions.
Args:
filters: Mapping of field name to value; falsy values are ignored.
Returns:
A combined filter expression or None if no filters are provided.
"""
parts = [Tag(k) == v for k, v in filters.items() if v]
return reduce(and_, parts) if parts else None
def _build_schema_dict(
self,
*,
index_name: str,
prefix: str,
vector_field_name: str | None,
vector_dims: int | None,
vector_datatype: str | None,
vector_algorithm: Literal["flat", "hnsw"] | None,
vector_distance_metric: Literal["cosine", "ip", "l2"] | None,
) -> dict[str, Any]:
"""Builds the RediSearch schema configuration dictionary.
Defines text and tag fields for messages plus an optional vector field enabling KNN/hybrid search.
Keyword Args:
index_name: Index name.
prefix: Key prefix.
vector_field_name: Vector field name or None.
vector_dims: Vector dimensionality or None.
vector_datatype: Vector datatype or None.
vector_algorithm: Vector index algorithm or None.
vector_distance_metric: Vector distance metric or None.
Returns:
Dict representing the index and fields configuration.
"""
fields: list[dict[str, Any]] = [
{"name": "role", "type": "tag"},
{"name": "mime_type", "type": "tag"},
{"name": "content", "type": "text"},
# Conversation tracking
{"name": "conversation_id", "type": "tag"},
{"name": "message_id", "type": "tag"},
{"name": "author_name", "type": "tag"},
# Partition fields (TAG for fast filtering)
{"name": "application_id", "type": "tag"},
{"name": "agent_id", "type": "tag"},
{"name": "user_id", "type": "tag"},
{"name": "thread_id", "type": "tag"},
]
# Add vector field only if configured (keeps provider runnable with no params)
if vector_field_name is not None and vector_dims is not None:
fields.append({
"name": vector_field_name,
"type": "vector",
"attrs": {
"algorithm": (vector_algorithm or "hnsw"),
"dims": int(vector_dims),
"distance_metric": (vector_distance_metric or "cosine"),
"datatype": (vector_datatype or "float32"),
},
})
return {
"index": {
"name": index_name,
"prefix": prefix,
"key_separator": ":",
"storage_type": "hash",
},
"fields": fields,
}
async def _ensure_index(self) -> None:
"""Initialize the search index.
- Connect to existing index if it exists and schema matches
- Create new index if it doesn't exist
- Overwrite if requested via overwrite_index=True
- Validate schema compatibility to prevent accidental data loss
"""
if self._index_initialized:
return
# Check if index already exists
index_exists = await self.redis_index.exists()
if not self.overwrite_index and index_exists:
# Validate schema compatibility before connecting
await self._validate_schema_compatibility()
# Create the index (will connect to existing or create new)
await self.redis_index.create(overwrite=self.overwrite_index, drop=False)
self._index_initialized = True
async def _validate_schema_compatibility(self) -> None:
"""Validate that existing index schema matches current configuration.
Raises ServiceInitializationError if schemas don't match, with helpful guidance.
self._build_schema_dict returns a minimal schema while Redis returns an expanded
schema with all defaults filled in. To compare for incompatibilities, compare
significant parts of the schema by creating signatures with normalized default values.
"""
# Defaults for attr normalization
TAG_DEFAULTS = {"separator": ",", "case_sensitive": False, "withsuffixtrie": False}
TEXT_DEFAULTS = {"weight": 1.0, "no_stem": False}
def _significant_index(i: dict[str, Any]) -> dict[str, Any]:
return {k: i.get(k) for k in ("name", "prefix", "key_separator", "storage_type")}
def _sig_tag(attrs: dict[str, Any] | None) -> dict[str, Any]:
a = {**TAG_DEFAULTS, **(attrs or {})}
return {k: a[k] for k in ("separator", "case_sensitive", "withsuffixtrie")}
def _sig_text(attrs: dict[str, Any] | None) -> dict[str, Any]:
a = {**TEXT_DEFAULTS, **(attrs or {})}
return {k: a[k] for k in ("weight", "no_stem")}
def _sig_vector(attrs: dict[str, Any] | None) -> dict[str, Any]:
a = {**(attrs or {})}
# Require these to exist if vector field is present
return {k: a.get(k) for k in ("algorithm", "dims", "distance_metric", "datatype")}
def _schema_signature(schema: dict[str, Any]) -> dict[str, Any]:
# Order-independent, minimal signature
sig: dict[str, Any] = {"index": _significant_index(schema.get("index", {})), "fields": {}}
for f in schema.get("fields", []):
name, ftype = f.get("name"), f.get("type")
if not name:
continue
if ftype == "tag":
sig["fields"][name] = {"type": "tag", "attrs": _sig_tag(f.get("attrs"))}
elif ftype == "text":
sig["fields"][name] = {"type": "text", "attrs": _sig_text(f.get("attrs"))}
elif ftype == "vector":
sig["fields"][name] = {"type": "vector", "attrs": _sig_vector(f.get("attrs"))}
else:
# Unknown field types: compare by type only
sig["fields"][name] = {"type": ftype}
return sig
existing_index = await AsyncSearchIndex.from_existing(self.index_name, redis_url=self.redis_url)
existing_schema = existing_index.schema.to_dict()
current_schema = self.schema_dict
existing_sig = _schema_signature(existing_schema)
current_sig = _schema_signature(current_schema)
if existing_sig != current_sig:
# Add sigs to error message
raise ServiceInitializationError(
"Existing Redis index schema is incompatible with the current configuration.\n"
f"Existing (significant): {json.dumps(existing_sig, indent=2, sort_keys=True)}\n"
f"Current (significant): {json.dumps(current_sig, indent=2, sort_keys=True)}\n"
"Set overwrite_index=True to rebuild if this change is intentional."
)
async def _add(
self,
*,
data: dict[str, Any] | list[dict[str, Any]],
metadata: dict[str, Any] | None = None,
) -> None:
"""Inserts one or many documents with partition fields populated.
Fills default partition fields, optionally embeds content when configured, and loads documents in a batch.
Keyword Args:
data: Single document or list of documents to insert.
metadata: Optional metadata dictionary (unused placeholder).
Raises:
ServiceInvalidRequestError: If required fields are missing or invalid.
"""
# Ensure provider has at least one scope set (symmetry with Mem0Provider)
self._validate_filters()
await self._ensure_index()
docs = data if isinstance(data, list) else [data]
prepared: list[dict[str, Any]] = []
for doc in docs:
d = dict(doc) # shallow copy
# Partition defaults
d.setdefault("application_id", self.application_id)
d.setdefault("agent_id", self.agent_id)
d.setdefault("user_id", self.user_id)
d.setdefault("thread_id", self._effective_thread_id)
# Conversation defaults
d.setdefault("conversation_id", self._conversation_id)
# Logical requirement
if "content" not in d:
raise ServiceInvalidRequestError("add() requires a 'content' field in data")
# Vector field requirement (only if schema has one)
if self.vector_field_name:
d.setdefault(self.vector_field_name, None)
prepared.append(d)
# Batch embed contents for every message
if self.redis_vectorizer and self.vector_field_name:
text_list = [d["content"] for d in prepared]
embeddings = await self.redis_vectorizer.aembed_many(text_list, batch_size=len(text_list))
for i, d in enumerate(prepared):
vec = np.asarray(embeddings[i], dtype=np.float32).tobytes()
field_name: str = self.vector_field_name
d[field_name] = vec
# Load all at once if supported
await self.redis_index.load(prepared)
return
async def _redis_search(
self,
text: str,
*,
text_scorer: str = "BM25STD",
filter_expression: Any | None = None,
return_fields: list[str] | None = None,
num_results: int = 10,
alpha: float = 0.7,
) -> list[dict[str, Any]]:
"""Runs a text or hybrid vector-text search with optional filters.
Builds a TextQuery or HybridQuery and automatically ANDs partition filters to keep results scoped and safe.
Args:
text: Query text.
Keyword Args:
text_scorer: Scorer to use for text ranking.
filter_expression: Additional filter expression to AND with partition filters.
return_fields: Fields to return in results.
num_results: Maximum number of results.
alpha: Hybrid balancing parameter when vectors are enabled.
Returns:
List of result dictionaries.
Raises:
ServiceInvalidRequestError: If input is invalid or the query fails.
"""
# Enforce presence of at least one provider-level filter (symmetry with Mem0Provider)
await self._ensure_index()
self._validate_filters()
q = (text or "").strip()
if not q:
raise ServiceInvalidRequestError("text_search() requires non-empty text")
num_results = max(int(num_results or 10), 1)
combined_filter = self._build_filter_from_dict({
"application_id": self.application_id,
"agent_id": self.agent_id,
"user_id": self.user_id,
"thread_id": self._effective_thread_id,
"conversation_id": self._conversation_id,
})
if filter_expression is not None:
combined_filter = (combined_filter & filter_expression) if combined_filter else filter_expression
# Choose return fields
return_fields = (
return_fields
if return_fields is not None
else ["content", "role", "application_id", "agent_id", "user_id", "thread_id"]
)
try:
if self.redis_vectorizer and self.vector_field_name:
# Build hybrid query: combine full-text and vector similarity
vector = await self.redis_vectorizer.aembed(q)
query = HybridQuery(
text=q,
text_field_name="content",
vector=vector,
vector_field_name=self.vector_field_name,
text_scorer=text_scorer,
filter_expression=combined_filter,
alpha=alpha,
dtype=self.redis_vectorizer.dtype,
num_results=num_results,
return_fields=return_fields,
stopwords=None,
)
hybrid_results = await self.redis_index.query(query)
return cast(list[dict[str, Any]], hybrid_results)
# Text-only search
query = TextQuery(
text=q,
text_field_name="content",
text_scorer=text_scorer,
filter_expression=combined_filter,
num_results=num_results,
return_fields=return_fields,
stopwords=None,
)
text_results = await self.redis_index.query(query)
return cast(list[dict[str, Any]], text_results)
except Exception as exc: # pragma: no cover - surface as framework error
raise ServiceInvalidRequestError(f"Redis text search failed: {exc}") from exc
async def search_all(self, page_size: int = 200) -> list[dict[str, Any]]:
"""Returns all documents in the index.
Streams results via pagination to avoid excessive memory and response sizes.
Args:
page_size: Page size used for pagination under the hood.
Returns:
List of all documents.
"""
out: list[dict[str, Any]] = []
async for batch in self.redis_index.paginate(
FilterQuery(FilterExpression("*"), return_fields=[], num_results=page_size),
page_size=page_size,
):
out.extend(batch)
return out
@property
def _effective_thread_id(self) -> str | None:
"""Resolves the active thread id.
Returns per-operation thread id when scoping is enabled; otherwise the provider's thread id.
"""
return self._per_operation_thread_id if self.scope_to_per_operation_thread_id else self.thread_id
@override
async def thread_created(self, thread_id: str | None) -> None:
"""Called when a new thread is created.
Captures the per-operation thread id when scoping is enabled to enforce single-thread usage.
Args:
thread_id: The ID of the thread or None.
"""
self._validate_per_operation_thread_id(thread_id)
self._per_operation_thread_id = self._per_operation_thread_id or thread_id
# Track current conversation id (Agent passes conversation_id here)
self._conversation_id = thread_id or self._conversation_id
@override
async def invoked(
self,
request_messages: Message | Sequence[Message],
response_messages: Message | Sequence[Message] | None = None,
invoke_exception: Exception | None = None,
**kwargs: Any,
) -> None:
self._validate_filters()
request_messages_list = [request_messages] if isinstance(request_messages, Message) else list(request_messages)
response_messages_list = (
[response_messages]
if isinstance(response_messages, Message)
else list(response_messages)
if response_messages
else []
)
messages_list = [*request_messages_list, *response_messages_list]
messages: list[dict[str, Any]] = []
for message in messages_list:
if message.role in {"user", "assistant", "system"} and message.text and message.text.strip():
shaped: dict[str, Any] = {
"role": message.role,
"content": message.text,
"conversation_id": self._conversation_id,
"message_id": message.message_id,
"author_name": message.author_name,
}
messages.append(shaped)
if messages:
await self._add(data=messages)
@override
async def invoking(self, messages: Message | MutableSequence[Message], **kwargs: Any) -> Context:
"""Called before invoking the model to provide scoped context.
Concatenates recent messages into a query, fetches matching memories from Redis.
Prepends them as instructions.
Args:
messages: List of new messages in the thread.
Keyword Args:
**kwargs: not used at present at present.
Returns:
Context: Context object containing instructions with memories.
"""
self._validate_filters()
messages_list = [messages] if isinstance(messages, Message) else list(messages)
input_text = "\n".join(msg.text for msg in messages_list if msg and msg.text and msg.text.strip())
memories = await self._redis_search(text=input_text)
line_separated_memories = "\n".join(
str(memory.get("content", "")) for memory in memories if memory.get("content")
)
return Context(
messages=[Message(role="user", text=f"{self.context_prompt}\n{line_separated_memories}")]
if line_separated_memories
else None
)
async def __aenter__(self) -> Self:
"""Async context manager entry.
No special setup is required; provided for symmetry with the Mem0 provider.
"""
return self
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
"""Async context manager exit.
No cleanup is required; indexes/keys remain unless explicitly cleared.
"""
return
def _validate_filters(self) -> None:
"""Validates that at least one filter is provided.
Prevents unbounded operations by requiring a partition filter before reads or writes.
Raises:
ServiceInitializationError: If no filters are provided.
"""
if not self.agent_id and not self.user_id and not self.application_id and not self.thread_id:
raise ServiceInitializationError(
"At least one of the filters: agent_id, user_id, application_id, or thread_id is required."
)
def _validate_per_operation_thread_id(self, thread_id: str | None) -> None:
"""Validates that a new thread ID doesn't conflict when scoped.
Prevents cross-thread data leakage by enforcing single-thread usage when per-operation scoping is enabled.
Args:
thread_id: The new thread ID or None.
Raises:
ValueError: If a new thread ID conflicts with the existing one.
"""
if (
self.scope_to_per_operation_thread_id
and thread_id
and self._per_operation_thread_id
and thread_id != self._per_operation_thread_id
):
raise ValueError(
"RedisProvider can only be used with one thread, when scope_to_per_operation_thread_id is True."
)
@@ -1,6 +1,6 @@
# Copyright (c) Microsoft. All rights reserved.
"""Tests for _RedisContextProvider and _RedisHistoryProvider."""
"""Tests for RedisContextProvider and RedisHistoryProvider."""
from __future__ import annotations
@@ -12,8 +12,8 @@ from agent_framework import AgentResponse, Message
from agent_framework._sessions import AgentSession, SessionContext
from agent_framework.exceptions import ServiceInitializationError
from agent_framework_redis._context_provider import _RedisContextProvider
from agent_framework_redis._history_provider import _RedisHistoryProvider
from agent_framework_redis._context_provider import RedisContextProvider
from agent_framework_redis._history_provider import RedisHistoryProvider
# ---------------------------------------------------------------------------
# Shared fixtures
@@ -63,13 +63,13 @@ def mock_redis_client():
# ===========================================================================
# _RedisContextProvider tests
# RedisContextProvider tests
# ===========================================================================
class TestRedisContextProviderInit:
def test_basic_construction(self, patch_index_from_dict: MagicMock): # noqa: ARG002
provider = _RedisContextProvider(source_id="ctx", user_id="u1")
provider = RedisContextProvider(source_id="ctx", user_id="u1")
assert provider.source_id == "ctx"
assert provider.user_id == "u1"
assert provider.redis_url == "redis://localhost:6379"
@@ -77,7 +77,7 @@ class TestRedisContextProviderInit:
assert provider.prefix == "context"
def test_custom_params(self, patch_index_from_dict: MagicMock): # noqa: ARG002
provider = _RedisContextProvider(
provider = RedisContextProvider(
source_id="ctx",
redis_url="redis://custom:6380",
index_name="my_idx",
@@ -95,31 +95,31 @@ class TestRedisContextProviderInit:
assert provider.context_prompt == "Custom prompt"
def test_default_context_prompt(self, patch_index_from_dict: MagicMock): # noqa: ARG002
provider = _RedisContextProvider(source_id="ctx", user_id="u1")
provider = RedisContextProvider(source_id="ctx", user_id="u1")
assert "Memories" in provider.context_prompt
def test_invalid_vectorizer_raises(self, patch_index_from_dict: MagicMock): # noqa: ARG002
from agent_framework.exceptions import AgentException
with pytest.raises(AgentException, match="not a valid type"):
_RedisContextProvider(source_id="ctx", user_id="u1", redis_vectorizer="bad") # type: ignore[arg-type]
RedisContextProvider(source_id="ctx", user_id="u1", redis_vectorizer="bad") # type: ignore[arg-type]
class TestRedisContextProviderValidateFilters:
def test_no_filters_raises(self, patch_index_from_dict: MagicMock): # noqa: ARG002
provider = _RedisContextProvider(source_id="ctx")
provider = RedisContextProvider(source_id="ctx")
with pytest.raises(ServiceInitializationError, match="(?i)at least one"):
provider._validate_filters()
def test_any_single_filter_ok(self, patch_index_from_dict: MagicMock): # noqa: ARG002
for kwargs in [{"user_id": "u"}, {"agent_id": "a"}, {"application_id": "app"}]:
provider = _RedisContextProvider(source_id="ctx", **kwargs)
provider = RedisContextProvider(source_id="ctx", **kwargs)
provider._validate_filters() # should not raise
class TestRedisContextProviderSchema:
def test_schema_has_expected_fields(self, patch_index_from_dict: MagicMock): # noqa: ARG002
provider = _RedisContextProvider(source_id="ctx", user_id="u1")
provider = RedisContextProvider(source_id="ctx", user_id="u1")
schema = provider.schema_dict
field_names = [f["name"] for f in schema["fields"]]
for expected in ("role", "content", "conversation_id", "message_id", "application_id", "agent_id", "user_id"):
@@ -128,7 +128,7 @@ class TestRedisContextProviderSchema:
assert schema["index"]["prefix"] == "context"
def test_schema_no_vector_without_vectorizer(self, patch_index_from_dict: MagicMock): # noqa: ARG002
provider = _RedisContextProvider(source_id="ctx", user_id="u1")
provider = RedisContextProvider(source_id="ctx", user_id="u1")
field_types = [f["type"] for f in provider.schema_dict["fields"]]
assert "vector" not in field_types
@@ -140,7 +140,7 @@ class TestRedisContextProviderBeforeRun:
patch_index_from_dict: MagicMock, # noqa: ARG002
):
mock_index.query = AsyncMock(return_value=[{"content": "Memory A"}, {"content": "Memory B"}])
provider = _RedisContextProvider(source_id="ctx", user_id="u1")
provider = RedisContextProvider(source_id="ctx", user_id="u1")
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", contents=["test query"])], session_id="s1")
@@ -157,7 +157,7 @@ class TestRedisContextProviderBeforeRun:
mock_index: AsyncMock,
patch_index_from_dict: MagicMock, # noqa: ARG002
):
provider = _RedisContextProvider(source_id="ctx", user_id="u1")
provider = RedisContextProvider(source_id="ctx", user_id="u1")
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", contents=[" "])], session_id="s1")
@@ -172,7 +172,7 @@ class TestRedisContextProviderBeforeRun:
patch_index_from_dict: MagicMock, # noqa: ARG002
):
mock_index.query = AsyncMock(return_value=[])
provider = _RedisContextProvider(source_id="ctx", user_id="u1")
provider = RedisContextProvider(source_id="ctx", user_id="u1")
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", contents=["hello"])], session_id="s1")
@@ -187,7 +187,7 @@ class TestRedisContextProviderAfterRun:
mock_index: AsyncMock,
patch_index_from_dict: MagicMock, # noqa: ARG002
):
provider = _RedisContextProvider(source_id="ctx", user_id="u1")
provider = RedisContextProvider(source_id="ctx", user_id="u1")
session = AgentSession(session_id="test-session")
response = AgentResponse(messages=[Message(role="assistant", contents=["response text"])])
ctx = SessionContext(input_messages=[Message(role="user", contents=["user input"])], session_id="s1")
@@ -206,7 +206,7 @@ class TestRedisContextProviderAfterRun:
mock_index: AsyncMock,
patch_index_from_dict: MagicMock, # noqa: ARG002
):
provider = _RedisContextProvider(source_id="ctx", user_id="u1")
provider = RedisContextProvider(source_id="ctx", user_id="u1")
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", contents=[" "])], session_id="s1")
@@ -219,7 +219,7 @@ class TestRedisContextProviderAfterRun:
mock_index: AsyncMock,
patch_index_from_dict: MagicMock, # noqa: ARG002
):
provider = _RedisContextProvider(source_id="ctx", application_id="app", agent_id="ag", user_id="u1")
provider = RedisContextProvider(source_id="ctx", application_id="app", agent_id="ag", user_id="u1")
session = AgentSession(session_id="test-session")
ctx = SessionContext(input_messages=[Message(role="user", contents=["hello"])], session_id="s1")
@@ -235,13 +235,13 @@ class TestRedisContextProviderAfterRun:
class TestRedisContextProviderContextManager:
async def test_aenter_returns_self(self, patch_index_from_dict: MagicMock): # noqa: ARG002
provider = _RedisContextProvider(source_id="ctx", user_id="u1")
provider = RedisContextProvider(source_id="ctx", user_id="u1")
async with provider as p:
assert p is provider
# ===========================================================================
# _RedisHistoryProvider tests
# RedisHistoryProvider tests
# ===========================================================================
@@ -249,7 +249,7 @@ class TestRedisHistoryProviderInit:
def test_basic_construction(self, mock_redis_client: MagicMock):
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider("memory", redis_url="redis://localhost:6379")
provider = RedisHistoryProvider("memory", redis_url="redis://localhost:6379")
assert provider.source_id == "memory"
assert provider.key_prefix == "chat_messages"
@@ -261,7 +261,7 @@ class TestRedisHistoryProviderInit:
def test_custom_params(self, mock_redis_client: MagicMock):
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider(
provider = RedisHistoryProvider(
"mem",
redis_url="redis://localhost:6379",
key_prefix="custom",
@@ -279,12 +279,12 @@ class TestRedisHistoryProviderInit:
def test_no_redis_url_or_credential_raises(self):
with pytest.raises(ValueError, match="Either redis_url or credential_provider must be provided"):
_RedisHistoryProvider("mem")
RedisHistoryProvider("mem")
def test_both_url_and_credential_raises(self):
mock_cred = MagicMock()
with pytest.raises(ValueError, match="mutually exclusive"):
_RedisHistoryProvider(
RedisHistoryProvider(
"mem",
redis_url="redis://localhost:6379",
credential_provider=mock_cred,
@@ -294,13 +294,13 @@ class TestRedisHistoryProviderInit:
def test_credential_provider_without_host_raises(self):
mock_cred = MagicMock()
with pytest.raises(ValueError, match="host is required"):
_RedisHistoryProvider("mem", credential_provider=mock_cred)
RedisHistoryProvider("mem", credential_provider=mock_cred)
def test_credential_provider_with_host(self):
mock_cred = MagicMock()
with patch("agent_framework_redis._history_provider.redis.Redis") as mock_redis_cls:
mock_redis_cls.return_value = MagicMock()
provider = _RedisHistoryProvider("mem", credential_provider=mock_cred, host="myhost")
provider = RedisHistoryProvider("mem", credential_provider=mock_cred, host="myhost")
mock_redis_cls.assert_called_once_with(
host="myhost",
@@ -317,7 +317,7 @@ class TestRedisHistoryProviderRedisKey:
def test_key_format(self, mock_redis_client: MagicMock):
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider("mem", redis_url="redis://localhost:6379", key_prefix="msgs")
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379", key_prefix="msgs")
assert provider._redis_key("session-123") == "msgs:session-123"
assert provider._redis_key(None) == "msgs:default"
@@ -331,7 +331,7 @@ class TestRedisHistoryProviderGetMessages:
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
messages = await provider.get_messages("s1")
assert len(messages) == 2
@@ -345,7 +345,7 @@ class TestRedisHistoryProviderGetMessages:
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
messages = await provider.get_messages("s1")
assert messages == []
@@ -355,7 +355,7 @@ class TestRedisHistoryProviderSaveMessages:
async def test_saves_serialized_messages(self, mock_redis_client: MagicMock):
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
msgs = [Message(role="user", contents=["Hello"]), Message(role="assistant", contents=["Hi"])]
await provider.save_messages("s1", msgs)
@@ -367,7 +367,7 @@ class TestRedisHistoryProviderSaveMessages:
async def test_empty_messages_noop(self, mock_redis_client: MagicMock):
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
await provider.save_messages("s1", [])
mock_redis_client.pipeline.assert_not_called()
@@ -377,7 +377,7 @@ class TestRedisHistoryProviderSaveMessages:
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider("mem", redis_url="redis://localhost:6379", max_messages=10)
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379", max_messages=10)
await provider.save_messages("s1", [Message(role="user", contents=["msg"])])
@@ -388,7 +388,7 @@ class TestRedisHistoryProviderSaveMessages:
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider("mem", redis_url="redis://localhost:6379", max_messages=10)
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379", max_messages=10)
await provider.save_messages("s1", [Message(role="user", contents=["msg"])])
@@ -399,7 +399,7 @@ class TestRedisHistoryProviderClear:
async def test_clear_calls_delete(self, mock_redis_client: MagicMock):
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
await provider.clear("session-1")
mock_redis_client.delete.assert_called_once_with("chat_messages:session-1")
@@ -414,7 +414,7 @@ class TestRedisHistoryProviderBeforeAfterRun:
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
session = AgentSession(session_id="test")
ctx = SessionContext(input_messages=[Message(role="user", contents=["new msg"])], session_id="s1")
@@ -428,7 +428,7 @@ class TestRedisHistoryProviderBeforeAfterRun:
async def test_after_run_stores_input_and_response(self, mock_redis_client: MagicMock):
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
provider = RedisHistoryProvider("mem", redis_url="redis://localhost:6379")
session = AgentSession(session_id="test")
ctx = SessionContext(input_messages=[Message(role="user", contents=["hi"])], session_id="s1")
@@ -443,7 +443,7 @@ class TestRedisHistoryProviderBeforeAfterRun:
async def test_after_run_skips_when_no_messages(self, mock_redis_client: MagicMock):
with patch("agent_framework_redis._history_provider.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
provider = _RedisHistoryProvider(
provider = RedisHistoryProvider(
"mem", redis_url="redis://localhost:6379", store_inputs=False, store_outputs=False
)
@@ -1,621 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from agent_framework import Content, Message
from agent_framework_redis import RedisChatMessageStore
class TestRedisChatMessageStore:
"""Unit tests for RedisChatMessageStore using mocked Redis client.
These tests use mocked Redis operations to verify the logic and behavior
of the RedisChatMessageStore without requiring a real Redis server.
"""
@pytest.fixture
def sample_messages(self):
"""Sample chat messages for testing."""
return [
Message(role="user", text="Hello", message_id="msg1"),
Message(role="assistant", text="Hi there!", message_id="msg2"),
Message(role="user", text="How are you?", message_id="msg3"),
]
@pytest.fixture
def mock_redis_client(self):
"""Mock Redis client with all required methods."""
client = MagicMock()
# Core list operations
client.lrange = AsyncMock(return_value=[])
client.llen = AsyncMock(return_value=0)
client.lindex = AsyncMock(return_value=None)
client.lset = AsyncMock(return_value=True)
client.lrem = AsyncMock(return_value=0)
client.lpop = AsyncMock(return_value=None)
client.rpop = AsyncMock(return_value=None)
client.ltrim = AsyncMock(return_value=True)
client.delete = AsyncMock(return_value=1)
# Pipeline operations
mock_pipeline = AsyncMock()
mock_pipeline.rpush = AsyncMock()
mock_pipeline.execute = AsyncMock()
client.pipeline.return_value.__aenter__.return_value = mock_pipeline
return client
@pytest.fixture
def redis_store(self, mock_redis_client):
"""Redis chat message store with mocked client."""
with patch("agent_framework_redis._chat_message_store.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
store = RedisChatMessageStore(redis_url="redis://localhost:6379", thread_id="test_thread_123")
store._redis_client = mock_redis_client
return store
def test_init_with_thread_id(self):
"""Test initialization with explicit thread ID."""
thread_id = "user123_session456"
with patch("agent_framework_redis._chat_message_store.redis.from_url"):
store = RedisChatMessageStore(redis_url="redis://localhost:6379", thread_id=thread_id)
assert store.thread_id == thread_id
assert store.redis_url == "redis://localhost:6379"
assert store.key_prefix == "chat_messages"
assert store.redis_key == f"chat_messages:{thread_id}"
def test_init_auto_generate_thread_id(self):
"""Test initialization with auto-generated thread ID."""
with patch("agent_framework_redis._chat_message_store.redis.from_url"):
store = RedisChatMessageStore(redis_url="redis://localhost:6379")
assert store.thread_id is not None
assert store.thread_id.startswith("thread_")
assert len(store.thread_id) > 10 # Should be a UUID
def test_init_with_custom_prefix(self):
"""Test initialization with custom key prefix."""
with patch("agent_framework_redis._chat_message_store.redis.from_url"):
store = RedisChatMessageStore(
redis_url="redis://localhost:6379", thread_id="test123", key_prefix="custom_messages"
)
assert store.redis_key == "custom_messages:test123"
def test_init_with_max_messages(self):
"""Test initialization with message limit."""
with patch("agent_framework_redis._chat_message_store.redis.from_url"):
store = RedisChatMessageStore(redis_url="redis://localhost:6379", thread_id="test123", max_messages=100)
assert store.max_messages == 100
def test_init_with_redis_url_required(self):
"""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"):
store = RedisChatMessageStore(
redis_url="redis://localhost:6379", thread_id="test123", messages=sample_messages
)
assert store._initial_messages == sample_messages
async def test_add_messages_single(self, redis_store, mock_redis_client, sample_messages):
"""Test adding a single message using pipeline operations."""
message = sample_messages[0]
await redis_store.add_messages([message])
# Verify pipeline operations were called
mock_redis_client.pipeline.assert_called_with(transaction=True)
# Get the pipeline mock and verify it was used correctly
pipeline_mock = mock_redis_client.pipeline.return_value.__aenter__.return_value
pipeline_mock.rpush.assert_called()
pipeline_mock.execute.assert_called()
async def test_add_messages_multiple(self, redis_store, mock_redis_client, sample_messages):
"""Test adding multiple messages using pipeline operations."""
await redis_store.add_messages(sample_messages)
# Verify pipeline operations
mock_redis_client.pipeline.assert_called_with(transaction=True)
# Verify rpush was called for each message
pipeline_mock = mock_redis_client.pipeline.return_value.__aenter__.return_value
assert pipeline_mock.rpush.call_count == len(sample_messages)
async def test_add_messages_with_max_limit(self, mock_redis_client):
"""Test adding messages with max limit triggers trimming."""
with patch("agent_framework_redis._chat_message_store.redis.from_url") as mock_from_url:
mock_from_url.return_value = mock_redis_client
# Mock llen to return count that exceeds limit after adding
mock_redis_client.llen.return_value = 5
store = RedisChatMessageStore(redis_url="redis://localhost:6379", thread_id="test123", max_messages=3)
store._redis_client = mock_redis_client
message = Message(role="user", text="Test")
await store.add_messages([message])
# Should trim after adding to keep only last 3 messages
mock_redis_client.ltrim.assert_called_once_with("chat_messages:test123", -3, -1)
async def test_list_messages_empty(self, redis_store, mock_redis_client):
"""Test listing messages when store is empty."""
mock_redis_client.lrange.return_value = []
messages = await redis_store.list_messages()
assert messages == []
mock_redis_client.lrange.assert_called_once_with("chat_messages:test_thread_123", 0, -1)
async def test_list_messages_with_data(self, redis_store, mock_redis_client, sample_messages):
"""Test listing messages with data in Redis."""
# Create proper serialized messages using the actual serialization method
test_messages = [
Message(role="user", text="Hello", message_id="msg1"),
Message(role="assistant", text="Hi there!", message_id="msg2"),
]
serialized_messages = [redis_store._serialize_message(msg) for msg in test_messages]
mock_redis_client.lrange.return_value = serialized_messages
messages = await redis_store.list_messages()
assert len(messages) == 2
assert messages[0].role == "user"
assert messages[0].text == "Hello"
assert messages[1].role == "assistant"
assert messages[1].text == "Hi there!"
async def test_list_messages_with_initial_messages(self, sample_messages):
"""Test that initial messages are added to Redis and retrieved correctly."""
with patch("agent_framework_redis._chat_message_store.redis.from_url") as mock_from_url:
mock_redis_client = MagicMock()
mock_redis_client.llen = AsyncMock(return_value=0) # Redis key is empty
mock_redis_client.lrange = AsyncMock(return_value=[])
# Mock pipeline for adding initial messages
mock_pipeline = AsyncMock()
mock_pipeline.rpush = AsyncMock()
mock_pipeline.execute = AsyncMock()
mock_redis_client.pipeline.return_value.__aenter__.return_value = mock_pipeline
mock_from_url.return_value = mock_redis_client
store = RedisChatMessageStore(
redis_url="redis://localhost:6379",
thread_id="test123",
messages=sample_messages[:1], # One initial message
)
store._redis_client = mock_redis_client
# Mock Redis to return the initial message after it's added
initial_message_json = store._serialize_message(sample_messages[0])
mock_redis_client.lrange.return_value = [initial_message_json]
messages = await store.list_messages()
assert len(messages) == 1
assert messages[0].text == "Hello"
# Verify initial message was added to Redis via pipeline
mock_pipeline.rpush.assert_called()
async def test_initial_messages_not_added_if_key_exists(self, sample_messages):
"""Test that initial messages are not added if Redis key already has data."""
with patch("agent_framework_redis._chat_message_store.redis.from_url") as mock_from_url:
mock_redis_client = MagicMock()
mock_redis_client.llen = AsyncMock(return_value=5) # Key already has messages
mock_redis_client.lrange = AsyncMock(return_value=[])
# Pipeline should not be called since key already exists
mock_pipeline = AsyncMock()
mock_pipeline.rpush = AsyncMock()
mock_pipeline.execute = AsyncMock()
mock_redis_client.pipeline.return_value.__aenter__.return_value = mock_pipeline
mock_from_url.return_value = mock_redis_client
store = RedisChatMessageStore(
redis_url="redis://localhost:6379",
thread_id="test123",
messages=sample_messages[:1], # One initial message
)
store._redis_client = mock_redis_client
await store.list_messages()
# Should check length but not add messages since key exists
mock_redis_client.llen.assert_called()
mock_pipeline.rpush.assert_not_called()
async def test_serialize_state(self, redis_store):
"""Test state serialization."""
state = await redis_store.serialize()
expected_state = {
"type": "redis_store_state",
"thread_id": "test_thread_123",
"redis_url": "redis://localhost:6379",
"key_prefix": "chat_messages",
"max_messages": None,
}
assert state == expected_state
async def test_deserialize_state(self, redis_store):
"""Test state deserialization."""
serialized_state = {
"thread_id": "restored_thread_456",
"redis_url": "redis://localhost:6380",
"key_prefix": "restored_messages",
"max_messages": 50,
}
await redis_store.update_from_state(serialized_state)
assert redis_store.thread_id == "restored_thread_456"
assert redis_store.redis_url == "redis://localhost:6380"
assert redis_store.key_prefix == "restored_messages"
assert redis_store.max_messages == 50
async def test_deserialize_state_empty(self, redis_store):
"""Test deserializing empty state doesn't change anything."""
original_thread_id = redis_store.thread_id
await redis_store.update_from_state(None)
assert redis_store.thread_id == original_thread_id
async def test_clear_messages(self, redis_store, mock_redis_client):
"""Test clearing all messages."""
await redis_store.clear()
mock_redis_client.delete.assert_called_once_with("chat_messages:test_thread_123")
async def test_message_serialization_roundtrip(self, sample_messages):
"""Test message serialization and deserialization roundtrip."""
with patch("agent_framework_redis._chat_message_store.redis.from_url"):
store = RedisChatMessageStore(redis_url="redis://localhost:6379", thread_id="test123")
message = sample_messages[0]
# Test serialization
serialized = store._serialize_message(message)
assert isinstance(serialized, str)
# Test deserialization
deserialized = store._deserialize_message(serialized)
assert deserialized.role == message.role
assert deserialized.text == message.text
assert deserialized.message_id == message.message_id
async def test_message_serialization_with_complex_content(self):
"""Test serialization of messages with complex content."""
with patch("agent_framework_redis._chat_message_store.redis.from_url"):
store = RedisChatMessageStore(redis_url="redis://localhost:6379", thread_id="test123")
# Message with multiple content types
message = Message(
role="assistant",
contents=[Content.from_text(text="Hello"), Content.from_text(text="World")],
author_name="TestBot",
message_id="complex_msg",
additional_properties={"metadata": "test"},
)
serialized = store._serialize_message(message)
deserialized = store._deserialize_message(serialized)
assert deserialized.role == "assistant"
assert deserialized.text == "Hello World"
assert deserialized.author_name == "TestBot"
assert deserialized.message_id == "complex_msg"
assert deserialized.additional_properties == {"metadata": "test"}
async def test_redis_connection_error_handling(self):
"""Test handling Redis connection errors in add_messages."""
with patch("agent_framework_redis._chat_message_store.redis.from_url") as mock_from_url:
mock_client = MagicMock()
# Mock pipeline to raise exception during execution
mock_pipeline = AsyncMock()
mock_pipeline.rpush = AsyncMock()
mock_pipeline.execute = AsyncMock(side_effect=Exception("Connection failed"))
mock_client.pipeline.return_value.__aenter__.return_value = mock_pipeline
mock_from_url.return_value = mock_client
store = RedisChatMessageStore(redis_url="redis://localhost:6379", thread_id="test123")
store._redis_client = mock_client
message = Message(role="user", text="Test")
# Should propagate Redis connection errors
with pytest.raises(Exception, match="Connection failed"):
await store.add_messages([message])
async def test_getitem(self, redis_store, mock_redis_client, sample_messages):
"""Test getitem method using Redis LINDEX."""
# Mock LINDEX to return specific messages
serialized_msg0 = redis_store._serialize_message(sample_messages[0])
serialized_msg1 = redis_store._serialize_message(sample_messages[1])
def mock_lindex(key, index):
if index == 0:
return serialized_msg0
if index == -1 or index == 1:
return serialized_msg1
return None
mock_redis_client.lindex = AsyncMock(side_effect=mock_lindex)
# Test positive index
message = await redis_store.getitem(0)
assert message.text == "Hello"
# Test negative index
message = await redis_store.getitem(-1)
assert message.text == "Hi there!"
async def test_getitem_index_error(self, redis_store, mock_redis_client):
"""Test getitem raises IndexError for invalid index."""
mock_redis_client.lindex = AsyncMock(return_value=None)
with pytest.raises(IndexError):
await redis_store.getitem(0)
async def test_setitem(self, redis_store, mock_redis_client, sample_messages):
"""Test setitem method using Redis LSET."""
mock_redis_client.llen.return_value = 2
mock_redis_client.lset = AsyncMock()
new_message = Message(role="user", text="Updated message")
await redis_store.setitem(0, new_message)
mock_redis_client.lset.assert_called_once()
call_args = mock_redis_client.lset.call_args
assert call_args[0][0] == "chat_messages:test_thread_123"
assert call_args[0][1] == 0
async def test_setitem_index_error(self, redis_store, mock_redis_client):
"""Test setitem raises IndexError for invalid index."""
mock_redis_client.llen.return_value = 0
new_message = Message(role="user", text="Test")
with pytest.raises(IndexError):
await redis_store.setitem(0, new_message)
async def test_append(self, redis_store, mock_redis_client):
"""Test append method delegates to add_messages."""
message = Message(role="user", text="Appended message")
await redis_store.append(message)
# Should call pipeline operations via add_messages
mock_redis_client.pipeline.assert_called_with(transaction=True)
# Verify the message was added via pipeline
pipeline_mock = mock_redis_client.pipeline.return_value.__aenter__.return_value
pipeline_mock.rpush.assert_called()
pipeline_mock.execute.assert_called()
async def test_count(self, redis_store, mock_redis_client):
"""Test count method."""
mock_redis_client.llen.return_value = 5
count = await redis_store.count()
assert count == 5
mock_redis_client.llen.assert_called_with("chat_messages:test_thread_123")
async def test_len_method(self, redis_store, mock_redis_client):
"""Test async __len__ method."""
mock_redis_client.llen.return_value = 3
length = await redis_store.__len__()
assert length == 3
mock_redis_client.llen.assert_called_with("chat_messages:test_thread_123")
def test_bool_method(self, redis_store):
"""Test __bool__ method always returns True."""
# Store should always be truthy
assert bool(redis_store) is True
assert redis_store.__bool__() is True
# Should work in if statements (this is what Agent Framework uses)
if redis_store:
assert True # Should reach this
else:
raise AssertionError("Store should be truthy")
async def test_index_found(self, redis_store, mock_redis_client, sample_messages):
"""Test index method when message is found using Redis LINDEX."""
mock_redis_client.llen.return_value = 2
# Mock LINDEX to return messages at each position
serialized_msg0 = redis_store._serialize_message(sample_messages[0])
serialized_msg1 = redis_store._serialize_message(sample_messages[1])
def mock_lindex(key, index):
if index == 0:
return serialized_msg0
if index == 1:
return serialized_msg1
return None
mock_redis_client.lindex = AsyncMock(side_effect=mock_lindex)
index = await redis_store.index(sample_messages[1])
assert index == 1
# Should have called lindex twice (index 0, then index 1)
assert mock_redis_client.lindex.call_count == 2
async def test_index_not_found(self, redis_store, mock_redis_client, sample_messages):
"""Test index method when message is not found."""
mock_redis_client.llen.return_value = 1
mock_redis_client.lindex = AsyncMock(return_value="different_message")
with pytest.raises(ValueError, match="Message not found in store"):
await redis_store.index(sample_messages[0])
async def test_remove(self, redis_store, mock_redis_client, sample_messages):
"""Test remove method using Redis LREM."""
mock_redis_client.lrem = AsyncMock(return_value=1) # 1 element removed
await redis_store.remove(sample_messages[0])
# Should use LREM to remove the message
expected_serialized = redis_store._serialize_message(sample_messages[0])
mock_redis_client.lrem.assert_called_once_with("chat_messages:test_thread_123", 1, expected_serialized)
async def test_remove_not_found(self, redis_store, mock_redis_client, sample_messages):
"""Test remove method when message is not found."""
mock_redis_client.lrem = AsyncMock(return_value=0) # 0 elements removed
with pytest.raises(ValueError, match="Message not found in store"):
await redis_store.remove(sample_messages[0])
async def test_extend(self, redis_store, mock_redis_client, sample_messages):
"""Test extend method delegates to add_messages."""
await redis_store.extend(sample_messages[:2])
# Should call pipeline operations via add_messages
mock_redis_client.pipeline.assert_called_with(transaction=True)
# Verify rpush was called for each message
pipeline_mock = mock_redis_client.pipeline.return_value.__aenter__.return_value
assert pipeline_mock.rpush.call_count >= 2
async def test_serialize_with_agent_thread(self, redis_store, sample_messages):
"""Test that RedisChatMessageStore can be serialized within an AgentThread.
This test verifies the fix for issue #1991 where calling thread.serialize()
with a RedisChatMessageStore would fail with "Messages should be a list" error.
"""
from agent_framework import AgentThread
thread = AgentThread(message_store=redis_store)
await thread.on_new_messages(sample_messages)
serialized = await thread.serialize()
assert serialized is not None
assert "chat_message_store_state" in serialized
assert serialized["chat_message_store_state"] is not None
@@ -1,425 +0,0 @@
# Copyright (c) Microsoft. All rights reserved.
from typing import Any
from unittest.mock import AsyncMock, MagicMock, patch
import numpy as np
import pytest
from agent_framework import Message
from agent_framework.exceptions import AgentException, ServiceInitializationError
from redisvl.utils.vectorize import CustomTextVectorizer
from agent_framework_redis import RedisProvider
CUSTOM_VECTORIZER = CustomTextVectorizer(embed=lambda x: [1.0, 2.0, 3.0], dtype="float32")
@pytest.fixture
def mock_index() -> AsyncMock:
idx = AsyncMock()
idx.create = AsyncMock()
idx.load = AsyncMock()
idx.query = AsyncMock()
idx.exists = AsyncMock(return_value=False)
async def _paginate_generator(*_args: Any, **_kwargs: Any):
# Default empty generator; override per-test as needed
if False: # pragma: no cover
yield []
return
idx.paginate = _paginate_generator
return idx
@pytest.fixture
def patch_index_from_dict(mock_index: AsyncMock):
with patch("agent_framework_redis._provider.AsyncSearchIndex") as mock_cls:
mock_cls.from_dict = MagicMock(return_value=mock_index)
# Mock from_existing to return a mock with matching schema by default
# This prevents schema validation errors in tests that don't specifically test schema validation
async def mock_from_existing(index_name, redis_url):
mock_existing = AsyncMock()
# Return a schema that will match whatever the provider generates
# This is a bit of a hack, but allows existing tests to continue working
mock_existing.schema.to_dict = MagicMock(
side_effect=lambda: mock_cls.from_dict.call_args[0][0] if mock_cls.from_dict.call_args else {}
)
return mock_existing
mock_cls.from_existing = AsyncMock(side_effect=mock_from_existing)
yield mock_cls
@pytest.fixture
def patch_queries():
calls: dict[str, Any] = {"TextQuery": [], "HybridQuery": [], "FilterExpression": []}
def _mk_query(kind: str):
class _Q: # simple marker object with captured kwargs
def __init__(self, **kwargs):
self.kind = kind
self.kwargs = kwargs
return _Q
with (
patch(
"agent_framework_redis._provider.TextQuery",
side_effect=lambda **k: calls["TextQuery"].append(k) or _mk_query("text")(**k),
) as text_q,
patch(
"agent_framework_redis._provider.HybridQuery",
side_effect=lambda **k: calls["HybridQuery"].append(k) or _mk_query("hybrid")(**k),
) as hybrid_q,
patch(
"agent_framework_redis._provider.FilterExpression",
side_effect=lambda s: calls["FilterExpression"].append(s) or ("FE", s),
) as filt,
):
yield {"calls": calls, "TextQuery": text_q, "HybridQuery": hybrid_q, "FilterExpression": filt}
class TestRedisProviderInitialization:
# Verifies the provider can be imported from the package
def test_import(self):
from agent_framework_redis._provider import RedisProvider
assert RedisProvider is not None
# Constructing without filters should not raise; filters are enforced at call-time
def test_init_without_filters_ok(self, patch_index_from_dict): # noqa: ARG002
provider = RedisProvider()
assert provider.user_id is None
assert provider.agent_id is None
assert provider.application_id is None
assert provider.thread_id is None
# Schema should omit vector field when no vector configuration is provided
def test_schema_without_vector_field(self, patch_index_from_dict):
RedisProvider(user_id="u1")
# Inspect schema passed to from_dict
args, kwargs = patch_index_from_dict.from_dict.call_args
schema = args[0]
assert isinstance(schema, dict)
names = [f["name"] for f in schema["fields"]]
types = [f["type"] for f in schema["fields"]]
assert "content" in names
assert "text" in types
assert "vector" not in types
class TestRedisProviderMessages:
@pytest.fixture
def sample_messages(self) -> list[Message]:
return [
Message(role="user", text="Hello, how are you?"),
Message(role="assistant", text="I'm doing well, thank you!"),
Message(role="system", text="You are a helpful assistant"),
]
# Writes require at least one scoping filter to avoid unbounded operations
async def test_messages_adding_requires_filters(self, patch_index_from_dict): # noqa: ARG002
provider = RedisProvider()
with pytest.raises(ServiceInitializationError):
await provider.invoked("thread123", Message(role="user", text="Hello"))
# Captures the per-operation thread id when provided
async def test_thread_created_sets_per_operation_id(self, patch_index_from_dict): # noqa: ARG002
provider = RedisProvider(user_id="u1")
await provider.thread_created("t1")
assert provider._per_operation_thread_id == "t1"
# Enforces single-thread usage when scope_to_per_operation_thread_id is True
async def test_thread_created_conflict_when_scoped(self, patch_index_from_dict): # noqa: ARG002
provider = RedisProvider(user_id="u1", scope_to_per_operation_thread_id=True)
provider._per_operation_thread_id = "t1"
with pytest.raises(ValueError) as exc:
await provider.thread_created("t2")
assert "only be used with one thread" in str(exc.value)
# Aggregates all results from the async paginator into a flat list
async def test_search_all_paginates(self, mock_index: AsyncMock, patch_index_from_dict): # noqa: ARG002
async def gen(_q, page_size: int = 200): # noqa: ARG001, ANN001
yield [{"id": 1}]
yield [{"id": 2}, {"id": 3}]
mock_index.paginate = gen
provider = RedisProvider(user_id="u1")
res = await provider.search_all(page_size=2)
assert res == [{"id": 1}, {"id": 2}, {"id": 3}]
class TestRedisProviderModelInvoking:
# Reads require at least one scoping filter to avoid unbounded operations
async def test_model_invoking_requires_filters(self, patch_index_from_dict): # noqa: ARG002
provider = RedisProvider()
with pytest.raises(ServiceInitializationError):
await provider.invoking(Message(role="user", text="Hi"))
# Ensures text-only search path is used and context is composed from hits
async def test_textquery_path_and_context_contents(
self, mock_index: AsyncMock, patch_index_from_dict, patch_queries
): # noqa: ARG002
# Arrange: text-only search
mock_index.query = AsyncMock(return_value=[{"content": "A"}, {"content": "B"}])
provider = RedisProvider(user_id="u1")
# Act
ctx = await provider.invoking([Message(role="user", text="q1")])
# Assert: TextQuery used (not HybridQuery), filter_expression included
assert patch_queries["TextQuery"].call_count == 1
assert patch_queries["HybridQuery"].call_count == 0
kwargs = patch_queries["calls"]["TextQuery"][0]
assert kwargs["text"] == "q1"
assert kwargs["text_field_name"] == "content"
assert kwargs["num_results"] == 10
assert "filter_expression" in kwargs
# Context contains memories joined after the default prompt
assert ctx.messages is not None and len(ctx.messages) == 1
text = ctx.messages[0].text
assert text.endswith("A\nB")
# When no results are returned, Context should have no contents
async def test_model_invoking_empty_results_returns_empty_context(
self, mock_index: AsyncMock, patch_index_from_dict, patch_queries
): # noqa: ARG002
mock_index.query = AsyncMock(return_value=[])
provider = RedisProvider(user_id="u1")
ctx = await provider.invoking([Message(role="user", text="any")])
assert ctx.messages == []
# Ensures hybrid vector-text search is used when a vectorizer and vector field are configured
async def test_hybridquery_path_with_vectorizer(self, mock_index: AsyncMock, patch_index_from_dict, patch_queries): # noqa: ARG002
mock_index.query = AsyncMock(return_value=[{"content": "Hit"}])
provider = RedisProvider(user_id="u1", redis_vectorizer=CUSTOM_VECTORIZER, vector_field_name="vec")
ctx = await provider.invoking([Message(role="user", text="hello")])
# Assert: HybridQuery used with vector and vector field
assert patch_queries["HybridQuery"].call_count == 1
k = patch_queries["calls"]["HybridQuery"][0]
assert k["text"] == "hello"
assert k["vector_field_name"] == "vec"
assert k["vector"] == [1.0, 2.0, 3.0]
assert k["dtype"] == "float32"
assert k["num_results"] == 10
assert "filter_expression" in k
# Context assembled from returned memories
assert ctx.messages and "Hit" in ctx.messages[0].text
class TestRedisProviderContextManager:
# Verifies async context manager returns self for chaining
async def test_async_context_manager_returns_self(self, patch_index_from_dict): # noqa: ARG002
provider = RedisProvider(user_id="u1")
async with provider as ctx:
assert ctx is provider
# Exit should be a no-op and not raise
async def test_aexit_noop(self, patch_index_from_dict): # noqa: ARG002
provider = RedisProvider(user_id="u1")
assert await provider.__aexit__(None, None, None) is None
class TestMessagesAddingBehavior:
# Adds messages while injecting partition defaults and preserving allowed roles
async def test_messages_adding_adds_partition_defaults_and_roles(
self, mock_index: AsyncMock, patch_index_from_dict
): # noqa: ARG002
provider = RedisProvider(
application_id="app",
agent_id="agent",
user_id="u1",
scope_to_per_operation_thread_id=True,
)
msgs = [
Message(role="user", text="u"),
Message(role="assistant", text="a"),
Message(role="system", text="s"),
]
await provider.invoked(msgs)
# Ensure load invoked with shaped docs containing defaults
assert mock_index.load.await_count == 1
(loaded_args, _kwargs) = mock_index.load.call_args
docs = loaded_args[0]
assert isinstance(docs, list) and len(docs) == 3
for d in docs:
assert d["role"] in {"user", "assistant", "system"}
assert d["content"] in {"u", "a", "s"}
assert d["application_id"] == "app"
assert d["agent_id"] == "agent"
assert d["user_id"] == "u1"
# Skips blank text and disallowed roles (e.g., TOOL) when adding messages
async def test_messages_adding_ignores_blank_and_disallowed_roles(
self, mock_index: AsyncMock, patch_index_from_dict
): # noqa: ARG002
provider = RedisProvider(user_id="u1", scope_to_per_operation_thread_id=True)
msgs = [
Message(role="user", text=" "),
Message(role="tool", text="tool output"),
]
await provider.invoked(msgs)
# No valid messages -> no load
assert mock_index.load.await_count == 0
class TestIndexCreationPublicCalls:
# Ensures index is created only once when drop=True on first public write call
async def test_messages_adding_triggers_index_create_once_when_drop_true(
self, mock_index: AsyncMock, patch_index_from_dict
): # noqa: ARG002
provider = RedisProvider(user_id="u1")
await provider.invoked(Message(role="user", text="m1"))
await provider.invoked(Message(role="user", text="m2"))
# create only on first call
assert mock_index.create.await_count == 1
# Ensures index is created when drop=False and the index does not exist on first read
async def test_model_invoking_triggers_create_when_drop_false_and_not_exists(
self, mock_index: AsyncMock, patch_index_from_dict
): # noqa: ARG002
mock_index.exists = AsyncMock(return_value=False)
provider = RedisProvider(user_id="u1")
mock_index.query = AsyncMock(return_value=[{"content": "C"}])
await provider.invoking([Message(role="user", text="q")])
assert mock_index.create.await_count == 1
class TestThreadCreatedAdditional:
# Allows None or same thread id repeatedly; different id raises when scoped
async def test_thread_created_allows_none_and_same_id(self, patch_index_from_dict): # noqa: ARG002
provider = RedisProvider(user_id="u1", scope_to_per_operation_thread_id=True)
# None is allowed
await provider.thread_created(None)
# Same id is allowed repeatedly
await provider.thread_created("t1")
await provider.thread_created("t1")
# Different id should raise
with pytest.raises(ValueError):
await provider.thread_created("t2")
class TestVectorPopulation:
# When vectorizer configured, invoked should embed content and populate the vector field
async def test_messages_adding_populates_vector_field_when_vectorizer_present(
self, mock_index: AsyncMock, patch_index_from_dict
): # noqa: ARG002
provider = RedisProvider(
user_id="u1",
scope_to_per_operation_thread_id=True,
redis_vectorizer=CUSTOM_VECTORIZER,
vector_field_name="vec",
)
await provider.invoked(Message(role="user", text="hello"))
assert mock_index.load.await_count == 1
(loaded_args, _kwargs) = mock_index.load.call_args
docs = loaded_args[0]
assert isinstance(docs, list) and len(docs) == 1
vec = docs[0].get("vec")
assert isinstance(vec, (bytes, bytearray))
assert len(vec) == 3 * np.dtype(np.float32).itemsize
class TestRedisProviderSchemaVectors:
# Adds a vector field when vectorizer supplies dims implicitly
def test_schema_with_vector_field_and_dims_inferred(self, patch_index_from_dict): # noqa: ARG002
RedisProvider(user_id="u1", redis_vectorizer=CUSTOM_VECTORIZER, vector_field_name="vec")
args, _ = patch_index_from_dict.from_dict.call_args
schema = args[0]
names = [f["name"] for f in schema["fields"]]
types = {f["name"]: f["type"] for f in schema["fields"]}
assert "vec" in names
assert types["vec"] == "vector"
# Raises when redis_vectorizer is not the correct type
def test_init_invalid_vectorizer(self, patch_index_from_dict): # noqa: ARG002
class DummyVectorizer:
pass
with pytest.raises(AgentException):
RedisProvider(user_id="u1", redis_vectorizer=DummyVectorizer(), vector_field_name="vec")
class TestEnsureIndex:
# Creates index once and marks _index_initialized to prevent duplicate calls
async def test_ensure_index_creates_once(self, mock_index: AsyncMock, patch_index_from_dict): # noqa: ARG002
# Mock index doesn't exist, so it will be created
mock_index.exists = AsyncMock(return_value=False)
provider = RedisProvider(user_id="u1", overwrite_index=False)
assert provider._index_initialized is False
await provider._ensure_index()
assert mock_index.create.await_count == 1
assert provider._index_initialized is True
# Second call should not create again due to _index_initialized flag
await provider._ensure_index()
assert mock_index.create.await_count == 1
# Creates index with overwrite=True when overwrite_index=True
async def test_ensure_index_with_overwrite_true(self, mock_index: AsyncMock, patch_index_from_dict): # noqa: ARG002
mock_index.exists = AsyncMock(return_value=True)
provider = RedisProvider(user_id="u1", overwrite_index=True)
await provider._ensure_index()
# Should call create with overwrite=True, drop=False
mock_index.create.assert_called_once_with(overwrite=True, drop=False)
# Creates index with overwrite=False when index doesn't exist
async def test_ensure_index_create_if_missing(self, mock_index: AsyncMock, patch_index_from_dict): # noqa: ARG002
mock_index.exists = AsyncMock(return_value=False)
provider = RedisProvider(user_id="u1", overwrite_index=False)
await provider._ensure_index()
# Should call create with overwrite=False, drop=False
mock_index.create.assert_called_once_with(overwrite=False, drop=False)
# Validates schema compatibility when index exists and overwrite=False
async def test_ensure_index_schema_validation_success(self, mock_index: AsyncMock, patch_index_from_dict): # noqa: ARG002
mock_index.exists = AsyncMock(return_value=True)
provider = RedisProvider(user_id="u1", overwrite_index=False)
# Mock existing index with matching schema
expected_schema = provider.schema_dict
patch_index_from_dict.from_existing.return_value.schema.to_dict.return_value = expected_schema
await provider._ensure_index()
# Should validate schema and proceed to create
patch_index_from_dict.from_existing.assert_called_once_with("context", redis_url="redis://localhost:6379")
mock_index.create.assert_called_once_with(overwrite=False, drop=False)
# Raises ServiceInitializationError when schemas don't match
async def test_ensure_index_schema_validation_failure(self, mock_index: AsyncMock, patch_index_from_dict): # noqa: ARG002
mock_index.exists = AsyncMock(return_value=True)
provider = RedisProvider(user_id="u1", overwrite_index=False)
# Override the mock to return a different schema after provider is created
async def mock_from_existing_different(index_name, redis_url):
mock_existing = AsyncMock()
mock_existing.schema.to_dict = MagicMock(return_value={"different": "schema"})
return mock_existing
patch_index_from_dict.from_existing = AsyncMock(side_effect=mock_from_existing_different)
with pytest.raises(ServiceInitializationError) as exc:
await provider._ensure_index()
assert "incompatible with the current configuration" in str(exc.value)
assert "overwrite_index=True" in str(exc.value)
# Should not call create when schema validation fails
mock_index.create.assert_not_called()