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

attune/memory/short_term/pubsub.py

@@ -0,0 +1,286 @@
+"""Pub/Sub messaging for real-time agent communication.
+
+This module provides publish/subscribe messaging:
+- Publish: Send messages to channels
+- Subscribe: Register handlers for channels
+- Unsubscribe: Remove handlers
+- Background listener thread management
+
+Key Prefix: PREFIX_PUBSUB = "pubsub:"
+
+Classes:
+    PubSubManager: Real-time publish/subscribe operations
+
+Example:
+    >>> from attune.memory.short_term.pubsub import PubSubManager
+    >>> from attune.memory.types import AgentCredentials, AccessTier
+    >>> pubsub = PubSubManager(base_ops)
+    >>> creds = AgentCredentials("agent_1", AccessTier.CONTRIBUTOR)
+    >>> def handler(msg): print(msg)
+    >>> pubsub.subscribe("signals", handler)
+    >>> pubsub.publish("signals", {"event": "done"}, creds)
+
+Copyright 2025 Smart-AI-Memory
+Licensed under Fair Source License 0.9
+"""
+
+from __future__ import annotations
+
+import json
+import threading
+import time
+from collections.abc import Callable
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+import structlog
+
+from attune.memory.types import (
+    AgentCredentials,
+)
+
+if TYPE_CHECKING:
+    from redis.client import PubSub
+
+    from attune.memory.short_term.base import BaseOperations
+
+logger = structlog.get_logger(__name__)
+
+
+class PubSubManager:
+    """Real-time publish/subscribe operations.
+
+    Provides channel-based messaging for agent communication.
+    Uses Redis Pub/Sub for real-time message delivery with
+    background listener threads.
+
+    The class manages its own state for subscriptions and
+    background threads, composed with BaseOperations for
+    Redis client access.
+
+    Attributes:
+        PREFIX_PUBSUB: Key prefix for pubsub channels
+
+    Example:
+        >>> pubsub = PubSubManager(base_ops)
+        >>> def on_signal(msg):
+        ...     print(f"Signal: {msg['data']}")
+        >>> pubsub.subscribe("agent_signals", on_signal)
+        >>> pubsub.publish("agent_signals", {"type": "heartbeat"}, creds)
+    """
+
+    PREFIX_PUBSUB = "pubsub:"
+
+    def __init__(self, base: BaseOperations) -> None:
+        """Initialize pub/sub manager.
+
+        Args:
+            base: BaseOperations instance for Redis client access
+        """
+        self._base = base
+        self._pubsub: PubSub | None = None
+        self._pubsub_thread: threading.Thread | None = None
+        self._subscriptions: dict[str, list[Callable[[dict], None]]] = {}
+        self._pubsub_running: bool = False
+        self._mock_pubsub_handlers: dict[str, list[Callable[[dict], None]]] = {}
+
+    def publish(
+        self,
+        channel: str,
+        message: dict,
+        credentials: AgentCredentials,
+    ) -> int:
+        """Publish a message to a channel for real-time notifications.
+
+        Args:
+            channel: Channel name (will be prefixed)
+            message: Message payload (dict)
+            credentials: Agent credentials (must be CONTRIBUTOR+)
+
+        Returns:
+            Number of subscribers that received the message
+
+        Raises:
+            PermissionError: If credentials lack publish access
+
+        Example:
+            >>> pubsub.publish(
+            ...     "agent_signals",
+            ...     {"event": "task_complete", "task_id": "123"},
+            ...     creds
+            ... )
+            2
+        """
+        if not credentials.can_stage():
+            raise PermissionError(
+                f"Agent {credentials.agent_id} cannot publish. "
+                "Requires CONTRIBUTOR tier or higher.",
+            )
+
+        start_time = time.perf_counter()
+        full_channel = f"{self.PREFIX_PUBSUB}{channel}"
+
+        payload = {
+            "channel": channel,
+            "from_agent": credentials.agent_id,
+            "timestamp": datetime.now().isoformat(),
+            "data": message,
+        }
+
+        # Handle mock mode
+        if self._base.use_mock:
+            handlers = self._mock_pubsub_handlers.get(full_channel, [])
+            for handler in handlers:
+                try:
+                    handler(payload)
+                except Exception as e:
+                    logger.warning("pubsub_handler_error", channel=channel, error=str(e))
+            latency_ms = (time.perf_counter() - start_time) * 1000
+            self._base._metrics.record_operation("publish", latency_ms)
+            return len(handlers)
+
+        # Handle real Redis client
+        if self._base._client is None:
+            return 0
+
+        count = self._base._client.publish(full_channel, json.dumps(payload))
+        latency_ms = (time.perf_counter() - start_time) * 1000
+        self._base._metrics.record_operation("publish", latency_ms)
+
+        logger.debug("pubsub_published", channel=channel, subscribers=count)
+        return int(count)
+
+    def subscribe(
+        self,
+        channel: str,
+        handler: Callable[[dict], None],
+        credentials: AgentCredentials | None = None,
+    ) -> bool:
+        """Subscribe to a channel for real-time notifications.
+
+        Args:
+            channel: Channel name to subscribe to
+            handler: Callback function receiving message dict
+            credentials: Optional credentials (any tier can subscribe)
+
+        Returns:
+            True if subscribed successfully
+
+        Example:
+            >>> def on_message(msg):
+            ...     print(f"Received: {msg['data']}")
+            >>> pubsub.subscribe("agent_signals", on_message)
+            True
+        """
+        full_channel = f"{self.PREFIX_PUBSUB}{channel}"
+
+        # Handle mock mode
+        if self._base.use_mock:
+            if full_channel not in self._mock_pubsub_handlers:
+                self._mock_pubsub_handlers[full_channel] = []
+            self._mock_pubsub_handlers[full_channel].append(handler)
+            logger.info("pubsub_subscribed_mock", channel=channel)
+            return True
+
+        # Handle real Redis client
+        if self._base._client is None:
+            return False
+
+        # Store handler
+        if full_channel not in self._subscriptions:
+            self._subscriptions[full_channel] = []
+        self._subscriptions[full_channel].append(handler)
+
+        # Create pubsub if needed
+        if self._pubsub is None:
+            self._pubsub = self._base._client.pubsub()
+
+        # Subscribe with internal handler
+        self._pubsub.subscribe(**{full_channel: self._pubsub_message_handler})
+
+        # Start listener thread if not running
+        if not self._pubsub_running:
+            self._pubsub_running = True
+            self._pubsub_thread = threading.Thread(
+                target=self._pubsub_listener,
+                daemon=True,
+                name="redis-pubsub-listener",
+            )
+            self._pubsub_thread.start()
+
+        logger.info("pubsub_subscribed", channel=channel)
+        return True
+
+    def _pubsub_message_handler(self, message: dict) -> None:
+        """Internal handler for pubsub messages."""
+        if message["type"] != "message":
+            return
+
+        channel = message["channel"]
+        if isinstance(channel, bytes):
+            channel = channel.decode()
+
+        try:
+            payload = json.loads(message["data"])
+        except json.JSONDecodeError:
+            payload = {"raw": message["data"]}
+
+        handlers = self._subscriptions.get(channel, [])
+        for handler in handlers:
+            try:
+                handler(payload)
+            except Exception as e:
+                logger.warning("pubsub_handler_error", channel=channel, error=str(e))
+
+    def _pubsub_listener(self) -> None:
+        """Background thread for listening to pubsub messages."""
+        while self._pubsub_running and self._pubsub:
+            try:
+                self._pubsub.get_message(ignore_subscribe_messages=True, timeout=1.0)
+            except Exception as e:
+                logger.warning("pubsub_listener_error", error=str(e))
+                time.sleep(1)
+
+    def unsubscribe(self, channel: str) -> bool:
+        """Unsubscribe from a channel.
+
+        Args:
+            channel: Channel name to unsubscribe from
+
+        Returns:
+            True if unsubscribed successfully
+
+        Example:
+            >>> pubsub.unsubscribe("agent_signals")
+            True
+        """
+        full_channel = f"{self.PREFIX_PUBSUB}{channel}"
+
+        # Handle mock mode
+        if self._base.use_mock:
+            self._mock_pubsub_handlers.pop(full_channel, None)
+            return True
+
+        # Handle real Redis client
+        if self._pubsub is None:
+            return False
+
+        self._pubsub.unsubscribe(full_channel)
+        self._subscriptions.pop(full_channel, None)
+        return True
+
+    def close(self) -> None:
+        """Close pubsub connection and stop listener thread.
+
+        Should be called when the manager is no longer needed
+        to clean up resources and stop background threads.
+
+        Example:
+            >>> pubsub.close()
+        """
+        self._pubsub_running = False
+        if self._pubsub:
+            self._pubsub.close()
+            self._pubsub = None
+        self._subscriptions.clear()
+        self._mock_pubsub_handlers.clear()

attune/memory/short_term/queues.py

@@ -0,0 +1,244 @@
+"""Task queue operations using Redis lists.
+
+This module provides queue operations for task processing:
+- Push: Add task to queue (LPUSH/RPUSH)
+- Pop: Remove and return task (LPOP/BLPOP)
+- Length: Get queue size
+- Peek: View task without removing
+
+Key Prefix: PREFIX_QUEUE = "queue:"
+
+Use Cases:
+- Background job queues
+- Task distribution
+- Work stealing patterns
+
+Classes:
+    QueueManager: Redis list operations for task queues
+
+Example:
+    >>> from attune.memory.short_term.queues import QueueManager
+    >>> from attune.memory.types import AgentCredentials, AccessTier
+    >>> queues = QueueManager(base_ops)
+    >>> creds = AgentCredentials("agent_1", AccessTier.CONTRIBUTOR)
+    >>> queues.push("tasks", {"type": "analyze", "file": "main.py"}, creds)
+    >>> task = queues.pop("tasks", creds)
+
+Copyright 2025 Smart-AI-Memory
+Licensed under Fair Source License 0.9
+"""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+import structlog
+
+from attune.memory.types import (
+    AgentCredentials,
+)
+
+if TYPE_CHECKING:
+    from attune.memory.short_term.base import BaseOperations
+
+logger = structlog.get_logger(__name__)
+
+
+class QueueManager:
+    """Redis list operations for task queues.
+
+    Provides FIFO queue operations using Redis lists for task
+    distribution and background job processing.
+
+    The class manages its own mock list storage for testing,
+    composed with BaseOperations for Redis client access.
+
+    Attributes:
+        PREFIX_QUEUE: Key prefix for queue names
+
+    Example:
+        >>> queues = QueueManager(base_ops)
+        >>> creds = AgentCredentials("agent_1", AccessTier.CONTRIBUTOR)
+        >>> queues.push("analysis_tasks", {"file": "main.py"}, creds)
+        >>> task = queues.pop("analysis_tasks", creds, timeout=5)
+    """
+
+    PREFIX_QUEUE = "queue:"
+
+    def __init__(self, base: BaseOperations) -> None:
+        """Initialize queue manager.
+
+        Args:
+            base: BaseOperations instance for Redis client access
+        """
+        self._base = base
+        self._mock_lists: dict[str, list[str]] = {}
+
+    def push(
+        self,
+        queue_name: str,
+        task: dict,
+        credentials: AgentCredentials,
+        priority: bool = False,
+    ) -> int:
+        """Push a task to a queue.
+
+        Args:
+            queue_name: Name of the queue
+            task: Task data
+            credentials: Agent credentials (must be CONTRIBUTOR+)
+            priority: If True, push to front (high priority)
+
+        Returns:
+            New queue length
+
+        Raises:
+            PermissionError: If credentials lack write access
+
+        Example:
+            >>> task = {"type": "analyze", "file": "main.py"}
+            >>> queues.push("agent_tasks", task, creds)
+            1
+        """
+        if not credentials.can_stage():
+            raise PermissionError(
+                f"Agent {credentials.agent_id} cannot push to queue. "
+                "Requires CONTRIBUTOR tier or higher.",
+            )
+
+        full_queue = f"{self.PREFIX_QUEUE}{queue_name}"
+        payload = json.dumps(
+            {
+                "task": task,
+                "queued_by": credentials.agent_id,
+                "queued_at": datetime.now().isoformat(),
+            },
+        )
+
+        # Handle mock mode
+        if self._base.use_mock:
+            if full_queue not in self._mock_lists:
+                self._mock_lists[full_queue] = []
+            if priority:
+                self._mock_lists[full_queue].insert(0, payload)
+            else:
+                self._mock_lists[full_queue].append(payload)
+            return len(self._mock_lists[full_queue])
+
+        # Handle real Redis client
+        if self._base._client is None:
+            return 0
+
+        if priority:
+            return int(self._base._client.lpush(full_queue, payload))
+        return int(self._base._client.rpush(full_queue, payload))
+
+    def pop(
+        self,
+        queue_name: str,
+        credentials: AgentCredentials,
+        timeout: int = 0,
+    ) -> dict | None:
+        """Pop a task from a queue.
+
+        Args:
+            queue_name: Name of the queue
+            credentials: Agent credentials
+            timeout: Seconds to block waiting (0 = no block)
+
+        Returns:
+            Task data or None if queue empty
+
+        Example:
+            >>> task = queues.pop("agent_tasks", creds, timeout=5)
+            >>> if task:
+            ...     process(task["task"])
+        """
+        full_queue = f"{self.PREFIX_QUEUE}{queue_name}"
+
+        # Handle mock mode
+        if self._base.use_mock:
+            if full_queue not in self._mock_lists or not self._mock_lists[full_queue]:
+                return None
+            payload = self._mock_lists[full_queue].pop(0)
+            data: dict = json.loads(payload)
+            return data
+
+        # Handle real Redis client
+        if self._base._client is None:
+            return None
+
+        if timeout > 0:
+            result = self._base._client.blpop(full_queue, timeout=timeout)
+            if result:
+                data = json.loads(result[1])
+                return data
+            return None
+
+        result = self._base._client.lpop(full_queue)
+        if result:
+            data = json.loads(result)
+            return data
+        return None
+
+    def length(self, queue_name: str) -> int:
+        """Get the length of a queue.
+
+        Args:
+            queue_name: Name of the queue
+
+        Returns:
+            Number of items in the queue
+
+        Example:
+            >>> count = queues.length("agent_tasks")
+            >>> print(f"Tasks pending: {count}")
+        """
+        full_queue = f"{self.PREFIX_QUEUE}{queue_name}"
+
+        # Handle mock mode
+        if self._base.use_mock:
+            return len(self._mock_lists.get(full_queue, []))
+
+        # Handle real Redis client
+        if self._base._client is None:
+            return 0
+
+        return int(self._base._client.llen(full_queue))
+
+    def peek(
+        self,
+        queue_name: str,
+        credentials: AgentCredentials,
+        count: int = 1,
+    ) -> list[dict]:
+        """Peek at tasks in a queue without removing them.
+
+        Args:
+            queue_name: Name of the queue
+            credentials: Agent credentials
+            count: Number of items to peek
+
+        Returns:
+            List of task data
+
+        Example:
+            >>> tasks = queues.peek("agent_tasks", creds, count=5)
+            >>> for task in tasks:
+            ...     print(task["task"]["type"])
+        """
+        full_queue = f"{self.PREFIX_QUEUE}{queue_name}"
+
+        # Handle mock mode
+        if self._base.use_mock:
+            items = self._mock_lists.get(full_queue, [])[:count]
+            return [json.loads(item) for item in items]
+
+        # Handle real Redis client
+        if self._base._client is None:
+            return []
+
+        items = self._base._client.lrange(full_queue, 0, count - 1)
+        return [json.loads(item) for item in items]