attune-ai 2.1.4__py3-none-any.whl → 2.2.0__py3-none-any.whl

attune/memory/short_term/base.py

@@ -0,0 +1,467 @@
+"""Core CRUD operations and connection management for short-term memory.
+
+This module provides the foundational Redis operations:
+- Connection management with retry logic
+- Basic get/set/delete/keys operations
+- Health check (ping) and statistics
+- Cleanup and lifecycle management
+
+The BaseOperations class is designed to be composed into the main
+RedisShortTermMemory facade, providing backward compatibility while
+enabling modular testing and maintenance.
+
+Target Methods (extracted from original RedisShortTermMemory):
+    - __init__ (initialization logic)
+    - client property
+    - _create_client_with_retry
+    - _execute_with_retry
+    - _get
+    - _set
+    - _delete
+    - _keys
+    - ping
+    - get_stats
+    - close
+
+Dependencies:
+    - RedisConfig for configuration
+    - RedisMetrics for operation tracking
+    - structlog for logging
+
+Copyright 2025 Smart-AI-Memory
+Licensed under Fair Source License 0.9
+"""
+
+from __future__ import annotations
+
+import os
+import time
+from collections.abc import Callable
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+import structlog
+
+from attune.memory.types import RedisConfig, RedisMetrics
+
+if TYPE_CHECKING:
+    pass
+
+logger = structlog.get_logger(__name__)
+
+# Redis availability check
+try:
+    import redis
+    from redis.exceptions import ConnectionError as RedisConnectionError
+    from redis.exceptions import TimeoutError as RedisTimeoutError
+
+    REDIS_AVAILABLE = True
+except ImportError:
+    REDIS_AVAILABLE = False
+    redis = None  # type: ignore
+    RedisConnectionError = Exception  # type: ignore
+    RedisTimeoutError = Exception  # type: ignore
+
+
+class BaseOperations:
+    """Core CRUD operations and connection management.
+
+    This class provides the foundational Redis operations that other
+    modules build upon. It handles:
+
+    - Connection creation with exponential backoff retry
+    - Basic get/set/delete/keys operations
+    - Health checks and statistics
+    - Resource cleanup
+
+    Example:
+        >>> from attune.memory.short_term.base import BaseOperations
+        >>> from attune.memory.types import RedisConfig
+        >>> config = RedisConfig(use_mock=True)
+        >>> base = BaseOperations(config=config)
+        >>> base._set("key", "value")
+        True
+        >>> base._get("key")
+        'value'
+
+    Attributes:
+        use_mock: Whether using mock storage instead of Redis
+        _config: Redis configuration
+        _metrics: Operation metrics tracker
+        _client: Redis client instance (None if mock)
+        _mock_storage: In-memory storage for mock mode
+    """
+
+    # Key prefixes for namespacing (shared across all operations)
+    PREFIX_WORKING = "empathy:working:"
+    PREFIX_STAGED = "empathy:staged:"
+    PREFIX_CONFLICT = "empathy:conflict:"
+    PREFIX_SESSION = "empathy:session:"
+    PREFIX_PUBSUB = "empathy:pubsub:"
+    PREFIX_STREAM = "empathy:stream:"
+    PREFIX_TIMELINE = "empathy:timeline:"
+    PREFIX_QUEUE = "empathy:queue:"
+
+    def __init__(
+        self,
+        host: str = "localhost",
+        port: int = 6379,
+        db: int = 0,
+        password: str | None = None,
+        use_mock: bool = False,
+        config: RedisConfig | None = None,
+    ) -> None:
+        """Initialize Redis connection and core components.
+
+        Args:
+            host: Redis host
+            port: Redis port
+            db: Redis database number
+            password: Redis password (optional)
+            use_mock: Use in-memory mock for testing
+            config: Full RedisConfig for advanced settings (overrides other args)
+        """
+        # Use config if provided, otherwise build from individual args
+        if config is not None:
+            self._config = config
+        else:
+            # Check environment variable for Redis enablement (default: disabled)
+            redis_enabled = os.getenv("REDIS_ENABLED", "false").lower() in (
+                "true",
+                "1",
+                "yes",
+            )
+
+            # Use environment variables for configuration if available
+            env_host = os.getenv("REDIS_HOST", host)
+            env_port = int(os.getenv("REDIS_PORT", str(port)))
+            env_db = int(os.getenv("REDIS_DB", str(db)))
+            env_password = os.getenv("REDIS_PASSWORD", password)
+
+            # If Redis is not enabled via env var, force mock mode
+            if not redis_enabled and not use_mock:
+                use_mock = True
+                logger.info(
+                    "redis_disabled_via_env",
+                    message="Redis not enabled in environment, using mock mode",
+                )
+
+            self._config = RedisConfig(
+                host=env_host,
+                port=env_port,
+                db=env_db,
+                password=env_password if env_password else None,
+                use_mock=use_mock,
+            )
+
+        self.use_mock = self._config.use_mock or not REDIS_AVAILABLE
+
+        # Initialize metrics
+        self._metrics = RedisMetrics()
+
+        # Mock storage for testing
+        self._mock_storage: dict[str, tuple[Any, float | None]] = {}
+        self._mock_lists: dict[str, list[str]] = {}
+        self._mock_sorted_sets: dict[str, list[tuple[float, str]]] = {}
+        self._mock_streams: dict[str, list[tuple[str, dict]]] = {}
+
+        # Create client
+        if self.use_mock:
+            self._client = None
+        else:
+            self._client = self._create_client_with_retry()
+
+    @property
+    def client(self) -> Any:
+        """Get the Redis client instance.
+
+        Returns:
+            Redis client instance or None if using mock mode
+
+        Example:
+            >>> memory = BaseOperations(use_mock=True)
+            >>> memory.client is None
+            True
+        """
+        return self._client
+
+    @property
+    def metrics(self) -> RedisMetrics:
+        """Get Redis metrics instance.
+
+        Returns:
+            RedisMetrics instance with connection and operation statistics
+
+        Example:
+            >>> base = BaseOperations(use_mock=True)
+            >>> base.metrics.retries_total
+            0
+        """
+        return self._metrics
+
+    def _create_client_with_retry(self) -> Any:
+        """Create Redis client with exponential backoff retry.
+
+        Returns:
+            Connected Redis client
+
+        Raises:
+            ConnectionError: If all retry attempts fail
+        """
+        max_attempts = self._config.retry_max_attempts
+        base_delay = self._config.retry_base_delay
+        max_delay = self._config.retry_max_delay
+
+        last_error: Exception | None = None
+
+        for attempt in range(max_attempts):
+            try:
+                client = redis.Redis(**self._config.to_redis_kwargs())
+                # Test connection
+                client.ping()
+                logger.info(
+                    "redis_connected",
+                    host=self._config.host,
+                    port=self._config.port,
+                    attempt=attempt + 1,
+                )
+                return client
+            except (RedisConnectionError, RedisTimeoutError) as e:
+                last_error = e
+                self._metrics.retries_total += 1
+
+                if attempt < max_attempts - 1:
+                    delay = min(base_delay * (2**attempt), max_delay)
+                    logger.warning(
+                        "redis_connection_retry",
+                        attempt=attempt + 1,
+                        max_attempts=max_attempts,
+                        delay=delay,
+                        error=str(e),
+                    )
+                    time.sleep(delay)
+
+        # All retries failed
+        logger.error(
+            "redis_connection_failed",
+            max_attempts=max_attempts,
+            error=str(last_error),
+        )
+        raise last_error if last_error else ConnectionError("Failed to connect to Redis")
+
+    def _execute_with_retry(
+        self, operation: Callable[[], Any], op_name: str = "operation"
+    ) -> Any:
+        """Execute a Redis operation with retry logic.
+
+        Args:
+            operation: Callable that performs the Redis operation
+            op_name: Name of operation for logging/metrics
+
+        Returns:
+            Result of the operation
+
+        Raises:
+            ConnectionError: If all retry attempts fail
+        """
+        start_time = time.perf_counter()
+        max_attempts = self._config.retry_max_attempts
+        base_delay = self._config.retry_base_delay
+        max_delay = self._config.retry_max_delay
+
+        last_error: Exception | None = None
+
+        for attempt in range(max_attempts):
+            try:
+                result = operation()
+                latency_ms = (time.perf_counter() - start_time) * 1000
+                self._metrics.record_operation(op_name, latency_ms, success=True)
+                return result
+            except (RedisConnectionError, RedisTimeoutError) as e:
+                last_error = e
+                self._metrics.retries_total += 1
+
+                if attempt < max_attempts - 1:
+                    delay = min(base_delay * (2**attempt), max_delay)
+                    logger.warning(
+                        "redis_operation_retry",
+                        operation=op_name,
+                        attempt=attempt + 1,
+                        delay=delay,
+                    )
+                    time.sleep(delay)
+
+        latency_ms = (time.perf_counter() - start_time) * 1000
+        self._metrics.record_operation(op_name, latency_ms, success=False)
+        raise last_error if last_error else ConnectionError("Redis operation failed")
+
+    def _get(self, key: str) -> str | None:
+        """Get value from Redis or mock storage.
+
+        Args:
+            key: Key to retrieve
+
+        Returns:
+            Value as string, or None if not found
+        """
+        # Mock mode path
+        if self.use_mock:
+            if key in self._mock_storage:
+                value, expires = self._mock_storage[key]
+                if expires is None or datetime.now().timestamp() < expires:
+                    return str(value) if value is not None else None
+                del self._mock_storage[key]
+            return None
+
+        # Real Redis path
+        if self._client is None:
+            return None
+
+        result = self._client.get(key)
+        return str(result) if result else None
+
+    def _set(self, key: str, value: str, ttl: int | None = None) -> bool:
+        """Set value in Redis or mock storage.
+
+        Args:
+            key: Key to set
+            value: Value to store
+            ttl: Time-to-live in seconds (optional)
+
+        Returns:
+            True if successful
+        """
+        # Mock mode path
+        if self.use_mock:
+            expires = datetime.now().timestamp() + ttl if ttl else None
+            self._mock_storage[key] = (value, expires)
+            return True
+
+        # Real Redis path
+        if self._client is None:
+            return False
+
+        # Set in Redis
+        if ttl:
+            self._client.setex(key, ttl, value)
+        else:
+            result = self._client.set(key, value)
+            if not result:
+                return False
+
+        return True
+
+    def _delete(self, key: str) -> bool:
+        """Delete key from Redis or mock storage.
+
+        Args:
+            key: Key to delete
+
+        Returns:
+            True if key was deleted
+        """
+        # Mock mode path
+        if self.use_mock:
+            if key in self._mock_storage:
+                del self._mock_storage[key]
+                return True
+            return False
+
+        # Real Redis path
+        if self._client is None:
+            return False
+
+        return bool(self._client.delete(key) > 0)
+
+    def _keys(self, pattern: str) -> list[str]:
+        """Get keys matching pattern.
+
+        Args:
+            pattern: Glob-style pattern to match
+
+        Returns:
+            List of matching keys
+        """
+        if self.use_mock:
+            import fnmatch
+
+            # Use list comp for small result sets (typical <1000 keys)
+            return [k for k in self._mock_storage.keys() if fnmatch.fnmatch(k, pattern)]
+
+        if self._client is None:
+            return []
+
+        keys = self._client.keys(pattern)
+        # Convert bytes to strings - needed for API return type
+        return [k.decode() if isinstance(k, bytes) else str(k) for k in keys]
+
+    def ping(self) -> bool:
+        """Check Redis connection health.
+
+        Returns:
+            True if connected and responsive
+        """
+        if self.use_mock:
+            return True
+        if self._client is None:
+            return False
+        try:
+            return bool(self._client.ping())
+        except Exception:  # noqa: BLE001
+            # INTENTIONAL: Health check should not raise, just return False
+            return False
+
+    def get_stats(self) -> dict:
+        """Get memory statistics.
+
+        Returns:
+            Dict with memory stats including mode, key counts by prefix
+        """
+        if self.use_mock:
+            # Use generator expressions for memory-efficient counting
+            return {
+                "mode": "mock",
+                "total_keys": len(self._mock_storage),
+                "working_keys": sum(
+                    1 for k in self._mock_storage if k.startswith(self.PREFIX_WORKING)
+                ),
+                "staged_keys": sum(
+                    1 for k in self._mock_storage if k.startswith(self.PREFIX_STAGED)
+                ),
+                "conflict_keys": sum(
+                    1 for k in self._mock_storage if k.startswith(self.PREFIX_CONFLICT)
+                ),
+            }
+
+        if self._client is None:
+            return {"mode": "disconnected", "error": "No Redis client"}
+
+        info = self._client.info("memory")
+        return {
+            "mode": "redis",
+            "used_memory": info.get("used_memory_human"),
+            "peak_memory": info.get("used_memory_peak_human"),
+            "total_keys": self._client.dbsize(),
+            "working_keys": len(self._keys(f"{self.PREFIX_WORKING}*")),
+            "staged_keys": len(self._keys(f"{self.PREFIX_STAGED}*")),
+            "conflict_keys": len(self._keys(f"{self.PREFIX_CONFLICT}*")),
+        }
+
+    def get_metrics(self) -> dict:
+        """Get operation metrics for observability.
+
+        Returns:
+            Dict with operation counts, latencies, and success rates
+        """
+        return self._metrics.to_dict()
+
+    def reset_metrics(self) -> None:
+        """Reset all metrics to zero."""
+        self._metrics = RedisMetrics()
+
+    def close(self) -> None:
+        """Close Redis connection and cleanup resources."""
+        if self._client:
+            self._client.close()
+            self._client = None
+        logger.info("redis_connection_closed")

attune/memory/short_term/batch.py

@@ -0,0 +1,219 @@
+"""Batch operations for efficient bulk processing.
+
+This module provides efficient batch operations using Redis pipelines:
+- Batch stash: Store multiple items in single round-trip
+- Batch retrieve: Get multiple items in single round-trip
+
+Benefits:
+- Reduces network round-trips
+- Atomic execution (all or nothing)
+- Better throughput for bulk operations
+
+Classes:
+    BatchOperations: Bulk stash/retrieve with Redis pipelines
+
+Example:
+    >>> from attune.memory.short_term.batch import BatchOperations
+    >>> from attune.memory.types import AgentCredentials, AccessTier
+    >>> batch_ops = BatchOperations(base_ops)
+    >>> creds = AgentCredentials("agent_1", AccessTier.CONTRIBUTOR)
+    >>> items = [("key1", {"a": 1}), ("key2", {"b": 2})]
+    >>> count = batch_ops.stash_batch(items, creds)
+    >>> data = batch_ops.retrieve_batch(["key1", "key2"], creds)
+
+Copyright 2025 Smart-AI-Memory
+Licensed under Fair Source License 0.9
+"""
+
+from __future__ import annotations
+
+import json
+import time
+from datetime import datetime
+from typing import TYPE_CHECKING, Any
+
+import structlog
+
+from attune.memory.types import (
+    AgentCredentials,
+    TTLStrategy,
+)
+
+if TYPE_CHECKING:
+    from attune.memory.short_term.base import BaseOperations
+
+logger = structlog.get_logger(__name__)
+
+
+class BatchOperations:
+    """Batch operations using Redis pipelines.
+
+    Provides efficient bulk stash/retrieve operations that reduce
+    network round-trips by batching multiple operations into a
+    single Redis pipeline execution.
+
+    The class is designed to be composed with BaseOperations
+    for dependency injection and access to Redis client.
+
+    Attributes:
+        PREFIX_WORKING: Key prefix for working memory namespace
+
+    Example:
+        >>> batch_ops = BatchOperations(base_ops)
+        >>> creds = AgentCredentials("agent_1", AccessTier.CONTRIBUTOR)
+        >>> items = [("analysis", {"score": 95}), ("summary", {"text": "..."})]
+        >>> count = batch_ops.stash_batch(items, creds)
+        2
+        >>> batch_ops.retrieve_batch(["analysis", "summary"], creds)
+        {'analysis': {'score': 95}, 'summary': {'text': '...'}}
+    """
+
+    PREFIX_WORKING = "empathy:working:"
+
+    def __init__(self, base: BaseOperations) -> None:
+        """Initialize batch operations.
+
+        Args:
+            base: BaseOperations instance for storage access
+        """
+        self._base = base
+
+    def stash_batch(
+        self,
+        items: list[tuple[str, Any]],
+        credentials: AgentCredentials,
+        ttl: TTLStrategy = TTLStrategy.WORKING_RESULTS,
+    ) -> int:
+        """Stash multiple items in a single operation.
+
+        Uses Redis pipeline for efficiency (reduces network round-trips).
+
+        Args:
+            items: List of (key, data) tuples
+            credentials: Agent credentials
+            ttl: Time-to-live strategy (applied to all items)
+
+        Returns:
+            Number of items successfully stashed
+
+        Raises:
+            TypeError: If items is not a list
+            PermissionError: If credentials lack write access
+
+        Example:
+            >>> items = [("key1", {"a": 1}), ("key2", {"b": 2})]
+            >>> count = batch_ops.stash_batch(items, creds)
+            2
+        """
+        # Pattern 5: Type validation
+        if not isinstance(items, list):
+            raise TypeError(f"items must be list, got {type(items).__name__}")
+
+        if not credentials.can_stage():
+            raise PermissionError(
+                f"Agent {credentials.agent_id} cannot write to memory. "
+                "Requires CONTRIBUTOR tier or higher.",
+            )
+
+        if not items:
+            return 0
+
+        start_time = time.perf_counter()
+
+        # Handle mock storage mode
+        if self._base.use_mock:
+            count = 0
+            for key, data in items:
+                full_key = f"{self.PREFIX_WORKING}{credentials.agent_id}:{key}"
+                payload = {
+                    "data": data,
+                    "agent_id": credentials.agent_id,
+                    "stashed_at": datetime.now().isoformat(),
+                }
+                expires = datetime.now().timestamp() + ttl.value
+                self._base._mock_storage[full_key] = (json.dumps(payload), expires)
+                count += 1
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            self._base._metrics.record_operation("stash_batch", latency_ms)
+            return count
+
+        # Handle real Redis client
+        if self._base._client is None:
+            return 0
+
+        pipe = self._base._client.pipeline()
+        for key, data in items:
+            full_key = f"{self.PREFIX_WORKING}{credentials.agent_id}:{key}"
+            payload = {
+                "data": data,
+                "agent_id": credentials.agent_id,
+                "stashed_at": datetime.now().isoformat(),
+            }
+            pipe.setex(full_key, ttl.value, json.dumps(payload))
+
+        results = pipe.execute()
+        count = sum(1 for r in results if r)
+        latency_ms = (time.perf_counter() - start_time) * 1000
+        self._base._metrics.record_operation("stash_batch", latency_ms)
+
+        logger.info("batch_stash_complete", count=count, total=len(items))
+        return count
+
+    def retrieve_batch(
+        self,
+        keys: list[str],
+        credentials: AgentCredentials,
+        agent_id: str | None = None,
+    ) -> dict[str, Any]:
+        """Retrieve multiple items in a single operation.
+
+        Uses Redis MGET for efficiency (single round-trip for all keys).
+
+        Args:
+            keys: List of keys to retrieve
+            credentials: Agent credentials
+            agent_id: Owner agent ID (defaults to credentials agent)
+
+        Returns:
+            Dict mapping key to data (missing keys omitted)
+
+        Example:
+            >>> data = batch_ops.retrieve_batch(["key1", "key2"], creds)
+            >>> print(data["key1"])
+            {'a': 1}
+        """
+        if not keys:
+            return {}
+
+        start_time = time.perf_counter()
+        owner = agent_id or credentials.agent_id
+        results: dict[str, Any] = {}
+
+        # Handle mock storage mode
+        if self._base.use_mock:
+            for key in keys:
+                full_key = f"{self.PREFIX_WORKING}{owner}:{key}"
+                if full_key in self._base._mock_storage:
+                    value, expires = self._base._mock_storage[full_key]
+                    if expires is None or datetime.now().timestamp() < expires:
+                        payload = json.loads(str(value))
+                        results[key] = payload.get("data")
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            self._base._metrics.record_operation("retrieve_batch", latency_ms)
+            return results
+
+        # Handle real Redis client
+        if self._base._client is None:
+            return {}
+
+        full_keys = [f"{self.PREFIX_WORKING}{owner}:{key}" for key in keys]
+        values = self._base._client.mget(full_keys)
+
+        for key, value in zip(keys, values, strict=False):
+            if value:
+                payload = json.loads(str(value))
+                results[key] = payload.get("data")
+
+        latency_ms = (time.perf_counter() - start_time) * 1000
+        self._base._metrics.record_operation("retrieve_batch", latency_ms)
+        return results