"""Redis-powered Session backend. Usage:: from agents.extensions.memory import RedisSession # Create from Redis URL session = RedisSession.from_url( session_id="user-123", url="redis://localhost:6379/0", ) # Or pass an existing Redis client that your application already manages session = RedisSession( session_id="user-123", redis_client=my_redis_client, ) await Runner.run(agent, "Hello", session=session) """ from __future__ import annotations import asyncio import json import time from typing import Any from urllib.parse import urlparse try: import redis.asyncio as redis from redis.asyncio import Redis except ImportError as e: raise ImportError( "RedisSession requires the 'redis' package. Install it with: pip install redis" ) from e from ...items import TResponseInputItem from ...memory.session import SessionABC class RedisSession(SessionABC): """Redis implementation of :pyclass:`agents.memory.session.Session`.""" def __init__( self, session_id: str, *, redis_client: Redis, key_prefix: str = "agents:session", ttl: int | None = None, ): """Initializes a new RedisSession. Args: session_id (str): Unique identifier for the conversation. redis_client (Redis[bytes]): A pre-configured Redis async client. key_prefix (str, optional): Prefix for Redis keys to avoid collisions. Defaults to "agents:session". ttl (int | None, optional): Time-to-live in seconds for session data. If None, data persists indefinitely. Defaults to None. """ self.session_id = session_id self._redis = redis_client self._key_prefix = key_prefix self._ttl = ttl self._lock = asyncio.Lock() self._owns_client = False # Track if we own the Redis client # Redis key patterns self._session_key = f"{self._key_prefix}:{self.session_id}" self._messages_key = f"{self._session_key}:messages" self._counter_key = f"{self._session_key}:counter" @classmethod def from_url( cls, session_id: str, *, url: str, redis_kwargs: dict[str, Any] | None = None, **kwargs: Any, ) -> RedisSession: """Create a session from a Redis URL string. Args: session_id (str): Conversation ID. url (str): Redis URL, e.g. "redis://localhost:6379/0" or "rediss://host:6380". redis_kwargs (dict[str, Any] | None): Additional keyword arguments forwarded to redis.asyncio.from_url. **kwargs: Additional keyword arguments forwarded to the main constructor (e.g., key_prefix, ttl, etc.). Returns: RedisSession: An instance of RedisSession connected to the specified Redis server. """ redis_kwargs = redis_kwargs or {} # Parse URL to determine if we need SSL parsed = urlparse(url) if parsed.scheme == "rediss": redis_kwargs.setdefault("ssl", True) redis_client = redis.from_url(url, **redis_kwargs) session = cls(session_id, redis_client=redis_client, **kwargs) session._owns_client = True # We created the client, so we own it return session async def _serialize_item(self, item: TResponseInputItem) -> str: """Serialize an item to JSON string. Can be overridden by subclasses.""" return json.dumps(item, separators=(",", ":")) async def _deserialize_item(self, item: str) -> TResponseInputItem: """Deserialize a JSON string to an item. Can be overridden by subclasses.""" return json.loads(item) # type: ignore[no-any-return] # json.loads returns Any but we know the structure async def _get_next_id(self) -> int: """Get the next message ID using Redis INCR for atomic increment.""" result = await self._redis.incr(self._counter_key) return int(result) async def _set_ttl_if_configured(self, *keys: str) -> None: """Set TTL on keys if configured.""" if self._ttl is not None: pipe = self._redis.pipeline() for key in keys: pipe.expire(key, self._ttl) await pipe.execute() # ------------------------------------------------------------------ # Session protocol implementation # ------------------------------------------------------------------ async def get_items(self, limit: int | None = None) -> list[TResponseInputItem]: """Retrieve the conversation history for this session. Args: limit: Maximum number of items to retrieve. If None, retrieves all items. When specified, returns the latest N items in chronological order. Returns: List of input items representing the conversation history """ async with self._lock: if limit is None: # Get all messages in chronological order raw_messages = await self._redis.lrange(self._messages_key, 0, -1) # type: ignore[misc] # Redis library returns Union[Awaitable[T], T] in async context else: if limit <= 0: return [] # Get the latest N messages (Redis list is ordered chronologically) # Use negative indices to get from the end - Redis uses -N to -1 for last N items raw_messages = await self._redis.lrange(self._messages_key, -limit, -1) # type: ignore[misc] # Redis library returns Union[Awaitable[T], T] in async context items: list[TResponseInputItem] = [] for raw_msg in raw_messages: try: # Handle both bytes (default) and str (decode_responses=True) Redis clients if isinstance(raw_msg, bytes): msg_str = raw_msg.decode("utf-8") else: msg_str = raw_msg # Already a string item = await self._deserialize_item(msg_str) items.append(item) except (json.JSONDecodeError, UnicodeDecodeError): # Skip corrupted messages continue return items async def add_items(self, items: list[TResponseInputItem]) -> None: """Add new items to the conversation history. Args: items: List of input items to add to the history """ if not items: return async with self._lock: pipe = self._redis.pipeline() # Set session metadata with current timestamp pipe.hset( self._session_key, mapping={ "session_id": self.session_id, "created_at": str(int(time.time())), "updated_at": str(int(time.time())), }, ) # Add all items to the messages list serialized_items = [] for item in items: serialized = await self._serialize_item(item) serialized_items.append(serialized) if serialized_items: pipe.rpush(self._messages_key, *serialized_items) # Update the session timestamp pipe.hset(self._session_key, "updated_at", str(int(time.time()))) # Execute all commands await pipe.execute() # Set TTL if configured await self._set_ttl_if_configured( self._session_key, self._messages_key, self._counter_key ) async def pop_item(self) -> TResponseInputItem | None: """Remove and return the most recent item from the session. Returns: The most recent item if it exists, None if the session is empty """ async with self._lock: # Use RPOP to atomically remove and return the rightmost (most recent) item raw_msg = await self._redis.rpop(self._messages_key) # type: ignore[misc] # Redis library returns Union[Awaitable[T], T] in async context if raw_msg is None: return None try: # Handle both bytes (default) and str (decode_responses=True) Redis clients if isinstance(raw_msg, bytes): msg_str = raw_msg.decode("utf-8") else: msg_str = raw_msg # Already a string return await self._deserialize_item(msg_str) except (json.JSONDecodeError, UnicodeDecodeError): # Return None for corrupted messages (already removed) return None async def clear_session(self) -> None: """Clear all items for this session.""" async with self._lock: # Delete all keys associated with this session await self._redis.delete( self._session_key, self._messages_key, self._counter_key, ) async def close(self) -> None: """Close the Redis connection. Only closes the connection if this session owns the Redis client (i.e., created via from_url). If the client was injected externally, the caller is responsible for managing its lifecycle. """ if self._owns_client: await self._redis.aclose() async def ping(self) -> bool: """Test Redis connectivity. Returns: True if Redis is reachable, False otherwise. """ try: await self._redis.ping() return True except Exception: return False