agent/agent-framework
2
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: add RedisContextProvider (#716)

Browse Source 
* Setting up

* Readme

* Add redis tests path to all-tests

* First pass integration

* Keep provider convention

* First pass integration

* add redis integration tests

* update README.md

* Add basic sample for redis integration

* Add partitioning, add partition-aware tests, improve sample script

* Fix code quality check

* Try to resolve pytest check

* Try to identify if pytest is the cause of failed checks

* Re-enable tests

* Rename redis test file

* Removing some tests to narrow down issue

* Revert, no difference

* Delete temp files

* Starting refactor of RedisProvider

* Build dynamic schema builder, still need to do dynamic embedding model config

* Add scope control

* Complete first pass functionality with OpenAI + HF vectors -> Tests, Samples, Demo to follow

* Fix code quality

* attempt to identify rootcause of failed test

* attempt to identify rootcause of failed test

* Attempt to resolve code quality fail

* Update pyproject.toml for foundry to pin     azure-ai-projects == 1.1.0b3,azure-ai-agents == 1.2.0b3

* Add tests for redisprovider

* Remove invalid tests

* Add API key handling for openai vectorizer

* Update uv.locl

* Use master uv.lock

* Begin sample file, add lazy index creation, fix faulty override

* Index drop and reinit depends on drop_redis_index not overwrite

* Add samples, threading included, escaped queries, verify threading works, sample README.md

* Refactor filters

* Opinionated vars

* Allow filter expression combination

* Try inline stubs for mypy

* Address mypy errors

* Better docstrings, tweaks for feedback

* Tweak example 3 in redis_threads.py sample

* adjust confusing name

* Enrich docstrings

* Add descriptions and comments to samples, externalize vectorizer choice, remove nltk and sentencetransformers dependnecy

* Add descriptions and comments to samples, externalize vectorizer choice, remove nltk and sentencetransformers dependnecy

* Incorporate initial feedback from dmytrostruk

* Fix uv.lock

* Attempt to resolve conflict

* Use remote .tomls

* Sanity check

* fix tests

* Remove hardcoded API key from samples

* Fix incorrect env var

* Make add and redis_search private

* Fix tests relying on private funcs

* Expand tests

* Explainer comments to each test

* Add a 'get_conversation_history' function to RedisProvider - This just returns messages in sequential order. Added 'created_at_*' timestamps to facilitate sequential recovery. function has to be manually invoked by user

* Add agent-framework-redis to  python/pyproject.toml

* Remove get_conversation_history

* improve redis context provider with pydantic techniques and safe index handling patterns

* add RedisChatMessageStore

* remove integration test :(

* fix mypy error

* Remove unused params

* Redo schema validation to be order-invariant, handle attrs (previously throwing errors due to strict ==)

* Expand explanation

* Add ChatMessageStore example

* Fix comments in redis_conversation.py

* Resolving uv.lock conflict, update to match main

* Fix test in redis provider

* Apply suggestion from @ekzhu

* Update python/packages/main/pyproject.toml

---------

Co-authored-by: Tyler Hutcherson <tyler.hutcherson@redis.com>
Co-authored-by: Eric Zhu <ekzhu@users.noreply.github.com>
This commit is contained in:
Han
2025-09-23 08:36:27 +08:00
committed by GitHub
Unverified
parent 1ef24d3e91
commit ca810076e8
23 changed files with 3406 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/redis/agent_framework_redis/__init__.py
+16
View File
@@ -0,0 +1,16 @@
# Copyright (c) Microsoft. All rights reserved.
import importlib.metadata
from ._chat_message_store import RedisChatMessageStore
from ._provider import RedisProvider
try:
__version__ = importlib.metadata.version(__name__)
except importlib.metadata.PackageNotFoundError:
__version__ = "0.0.0" # Fallback for development mode
__all__ = [
"RedisChatMessageStore",
"RedisProvider",
"__version__",
]
python/packages/redis/agent_framework_redis/_chat_message_store.py
+507
View File
@@ -0,0 +1,507 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import json
from collections.abc import Sequence
from typing import Any
from uuid import uuid4
import redis.asyncio as redis
from agent_framework import ChatMessage
from agent_framework._pydantic import AFBaseModel
class RedisStoreState(AFBaseModel):
"""State model for serializing and deserializing Redis chat message store data."""
thread_id: str
redis_url: str | None = None
key_prefix: str = "chat_messages"
max_messages: int | None = None
class RedisChatMessageStore:
"""Redis-backed implementation of ChatMessageStore 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,
thread_id: str | None = None,
key_prefix: str = "chat_messages",
max_messages: int | None = None,
messages: Sequence[ChatMessage] | None = None,
) -> None:
"""Initialize the Redis chat message store.
Creates a Redis-backed chat message store for a specific conversation thread.
The store will automatically create a Redis connection and manage message
persistence using Redis List operations.
Args:
redis_url: Redis connection URL (e.g., "redis://localhost:6379").
Required for establishing Redis connection.
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 redis_url is None (Redis connection is required).
redis.ConnectionError: If unable to connect to Redis server.
"""
# Validate required parameters
if redis_url is None:
raise ValueError("redis_url is required for Redis connection")
# Store configuration
self.redis_url = redis_url
self.thread_id = thread_id or f"thread_{uuid4()}"
self.key_prefix = key_prefix
self.max_messages = max_messages
# Initialize Redis client with connection pooling and async support
self._redis_client = redis.from_url(redis_url, decode_responses=True) # type: ignore[no-untyped-call]
# 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[ChatMessage]) -> 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 ChatMessage 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[ChatMessage]) -> None:
"""Add messages to the Redis store (ChatMessageStore protocol method).
This method implements the required ChatMessageStore 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 ChatMessage 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:
```python
messages = [ChatMessage(role=Role.USER, text="Hello"), ChatMessage(role=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[ChatMessage]:
"""Get all messages from the store in chronological order (ChatMessageStore protocol method).
This method implements the required ChatMessageStore protocol for retrieving messages.
Returns all messages stored in Redis, ordered from oldest (index 0) to newest (index -1).
Returns:
List of ChatMessage objects in chronological order (oldest first).
Returns empty list if no messages exist or if Redis connection fails.
Example:
```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 ChatMessage
message = self._deserialize_message(serialized_message)
messages.append(message)
return messages
async def serialize_state(self, **kwargs: Any) -> Any:
"""Serialize the current store state for persistence (ChatMessageStore protocol method).
This method implements the required ChatMessageStore protocol for state serialization.
Captures the Redis connection configuration and thread information needed to
reconstruct the store and reconnect to the same conversation data.
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.model_dump(**kwargs)
async def deserialize_state(self, serialized_store_state: Any, **kwargs: Any) -> None:
"""Deserialize state data into this store instance (ChatMessageStore protocol method).
This method implements the required ChatMessageStore 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.
**kwargs: Additional arguments passed to Pydantic model validation.
"""
if not serialized_store_state:
return
# Validate and parse the serialized state using Pydantic
state = RedisStoreState.model_validate(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:
```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: ChatMessage) -> str:
"""Serialize a ChatMessage to JSON string.
Args:
message: ChatMessage to serialize.
Returns:
JSON string representation of the message.
"""
# Convert ChatMessage to dictionary using Pydantic serialization
message_dict = message.model_dump()
# Serialize to compact JSON (no extra whitespace for Redis efficiency)
return json.dumps(message_dict, separators=(",", ":"))
def _deserialize_message(self, serialized_message: str) -> ChatMessage:
"""Deserialize a JSON string to ChatMessage.
Args:
serialized_message: JSON string representation of a message.
Returns:
ChatMessage object.
"""
# Parse JSON string back to dictionary
message_dict = json.loads(serialized_message)
# Reconstruct ChatMessage using Pydantic validation
return ChatMessage.model_validate(message_dict)
# ============================================================================
# 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) -> ChatMessage:
"""Get a message by index using Redis LINDEX.
Args:
index: The index of the message to retrieve.
Returns:
The ChatMessage 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: ChatMessage) -> None:
"""Set a message at the specified index using Redis LSET.
Args:
index: The index at which to set the message.
item: The ChatMessage 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: ChatMessage) -> None:
"""Append a message to the end of the store.
Args:
item: The ChatMessage 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: ChatMessage) -> 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 ChatMessage 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("ChatMessage not found in store")
async def remove(self, item: ChatMessage) -> 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 ChatMessage 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("ChatMessage not found in store")
async def extend(self, items: Sequence[ChatMessage]) -> None:
"""Extend the store by appending all messages from the iterable.
Args:
items: Sequence of ChatMessage 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})"
)
python/packages/redis/agent_framework_redis/_provider.py
+547
View File
@@ -0,0 +1,547 @@
# Copyright (c) Microsoft. All rights reserved.
from __future__ import annotations
import sys
from collections.abc import MutableSequence, Sequence
from functools import reduce
from operator import and_
from typing import Any, Literal, cast
from agent_framework import ChatMessage, Context, ContextProvider, Role, TextContent
from agent_framework.exceptions import (
ServiceInitializationError,
ServiceInvalidRequestError,
)
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
else:
from typing_extensions import Self # pragma: no cover
import json
import numpy as np
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
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.
"""
# Connection and indexing
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
_per_operation_thread_id: str | None = None
_token_escaper: TokenEscaper = TokenEscaper()
_conversation_id: str | None = None
_index_initialized: bool = False
_schema_dict: dict[str, Any] | None = None
def model_post_init(self, __context: Any) -> None:
"""Post-initialization hook to set up computed fields after Pydantic initialization.
This is called automatically by Pydantic after the model is initialized.
"""
# Create Redis index using the cached schema_dict property
self.redis_index = 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.
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.
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.
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
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
async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
"""Called when a new message is being added to the thread.
Validates scope, normalizes allowed roles, and persists messages to Redis via add().
Args:
thread_id: The ID of the thread or None.
new_messages: New messages to add.
"""
self._validate_filters()
self._validate_per_operation_thread_id(thread_id)
self._per_operation_thread_id = self._per_operation_thread_id or thread_id
messages_list = [new_messages] if isinstance(new_messages, ChatMessage) else list(new_messages)
messages: list[dict[str, Any]] = []
for message in messages_list:
if (
message.role.value in {Role.USER.value, Role.ASSISTANT.value, Role.SYSTEM.value}
and message.text
and message.text.strip()
):
shaped: dict[str, Any] = {
"role": message.role.value,
"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)
async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> 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.
Returns:
Context: Context object containing instructions with memories.
"""
self._validate_filters()
messages_list = [messages] if isinstance(messages, ChatMessage) 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")
)
content = TextContent(f"{self.context_prompt}\n{line_separated_memories}") if line_separated_memories else None
return Context(contents=[content] if content 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."
)