aury-agent 0.0.9__py3-none-any.whl → 0.0.11__py3-none-any.whl

aury/agents/messages/raw_store.py

@@ -1,224 +0,0 @@
-"""Raw message store for complete message storage.
-
-RawMessageStore stores complete, untruncated messages for:
-- HITL recovery (restore exact state)
-- Full-context recall (when truncated history is insufficient)
-- Audit/debugging
-
-Messages are stored per invocation and can be cleaned up after invocation completes.
-"""
-from __future__ import annotations
-
-from typing import Any, Protocol, runtime_checkable
-from datetime import datetime
-
-from ..core.types.session import generate_id
-
-
-@runtime_checkable
-class RawMessageStore(Protocol):
-    """Protocol for raw message storage.
-    
-    Stores complete messages per invocation.
-    """
-    
-    async def add(
-        self,
-        invocation_id: str,
-        message: dict[str, Any],
-    ) -> str:
-        """Add a message and return its ID.
-        
-        Args:
-            invocation_id: Invocation this message belongs to
-            message: Complete message dict
-            
-        Returns:
-            Generated message ID
-        """
-        ...
-    
-    async def get(self, msg_id: str) -> dict[str, Any] | None:
-        """Get a single message by ID.
-        
-        Args:
-            msg_id: Message ID
-            
-        Returns:
-            Message dict or None if not found
-        """
-        ...
-    
-    async def get_many(self, msg_ids: list[str]) -> list[dict[str, Any]]:
-        """Get multiple messages by IDs.
-        
-        Args:
-            msg_ids: List of message IDs
-            
-        Returns:
-            List of message dicts (in same order, None for missing)
-        """
-        ...
-    
-    async def get_by_invocation(self, invocation_id: str) -> list[dict[str, Any]]:
-        """Get all messages for an invocation.
-        
-        Args:
-            invocation_id: Invocation ID
-            
-        Returns:
-            List of messages in chronological order
-        """
-        ...
-    
-    async def delete_by_invocation(self, invocation_id: str) -> int:
-        """Delete all messages for an invocation.
-        
-        Args:
-            invocation_id: Invocation ID
-            
-        Returns:
-            Number of messages deleted
-        """
-        ...
-
-
-class StateBackendRawMessageStore:
-    """RawMessageStore implementation backed by StateBackend."""
-    
-    def __init__(self, backend: Any):  # StateBackend
-        """Initialize with StateBackend.
-        
-        Args:
-            backend: StateBackend instance
-        """
-        self._backend = backend
-    
-    async def add(
-        self,
-        invocation_id: str,
-        message: dict[str, Any],
-    ) -> str:
-        """Add a message using StateBackend."""
-        msg_id = generate_id("rmsg")
-        
-        # Store message with metadata
-        await self._backend.set("raw_messages", msg_id, {
-            "id": msg_id,
-            "invocation_id": invocation_id,
-            "message": message,
-            "created_at": datetime.now().isoformat(),
-        })
-        
-        # Add to invocation's message list
-        inv_key = f"raw_msg_ids:{invocation_id}"
-        msg_ids = await self._backend.get("raw_messages", inv_key) or []
-        msg_ids.append(msg_id)
-        await self._backend.set("raw_messages", inv_key, msg_ids)
-        
-        return msg_id
-    
-    async def get(self, msg_id: str) -> dict[str, Any] | None:
-        """Get a single message."""
-        data = await self._backend.get("raw_messages", msg_id)
-        if data:
-            return data.get("message")
-        return None
-    
-    async def get_many(self, msg_ids: list[str]) -> list[dict[str, Any]]:
-        """Get multiple messages."""
-        results = []
-        for msg_id in msg_ids:
-            msg = await self.get(msg_id)
-            if msg:
-                results.append(msg)
-        return results
-    
-    async def get_by_invocation(self, invocation_id: str) -> list[dict[str, Any]]:
-        """Get all messages for an invocation."""
-        inv_key = f"raw_msg_ids:{invocation_id}"
-        msg_ids = await self._backend.get("raw_messages", inv_key) or []
-        return await self.get_many(msg_ids)
-    
-    async def delete_by_invocation(self, invocation_id: str) -> int:
-        """Delete all messages for an invocation."""
-        inv_key = f"raw_msg_ids:{invocation_id}"
-        msg_ids = await self._backend.get("raw_messages", inv_key) or []
-        
-        # Delete each message
-        for msg_id in msg_ids:
-            await self._backend.remove("raw_messages", msg_id)
-        
-        # Delete the index
-        await self._backend.remove("raw_messages", inv_key)
-        
-        return len(msg_ids)
-
-
-class InMemoryRawMessageStore:
-    """In-memory raw message store for testing."""
-    
-    def __init__(self) -> None:
-        self._messages: dict[str, dict[str, Any]] = {}  # msg_id -> message data
-        self._invocation_index: dict[str, list[str]] = {}  # inv_id -> [msg_ids]
-    
-    async def add(
-        self,
-        invocation_id: str,
-        message: dict[str, Any],
-    ) -> str:
-        """Add a message."""
-        msg_id = generate_id("rmsg")
-        
-        self._messages[msg_id] = {
-            "id": msg_id,
-            "invocation_id": invocation_id,
-            "message": message,
-            "created_at": datetime.now().isoformat(),
-        }
-        
-        if invocation_id not in self._invocation_index:
-            self._invocation_index[invocation_id] = []
-        self._invocation_index[invocation_id].append(msg_id)
-        
-        return msg_id
-    
-    async def get(self, msg_id: str) -> dict[str, Any] | None:
-        """Get a single message."""
-        data = self._messages.get(msg_id)
-        if data:
-            return data.get("message")
-        return None
-    
-    async def get_many(self, msg_ids: list[str]) -> list[dict[str, Any]]:
-        """Get multiple messages."""
-        results = []
-        for msg_id in msg_ids:
-            msg = await self.get(msg_id)
-            if msg:
-                results.append(msg)
-        return results
-    
-    async def get_by_invocation(self, invocation_id: str) -> list[dict[str, Any]]:
-        """Get all messages for an invocation."""
-        msg_ids = self._invocation_index.get(invocation_id, [])
-        return await self.get_many(msg_ids)
-    
-    async def delete_by_invocation(self, invocation_id: str) -> int:
-        """Delete all messages for an invocation."""
-        msg_ids = self._invocation_index.pop(invocation_id, [])
-        for msg_id in msg_ids:
-            self._messages.pop(msg_id, None)
-        return len(msg_ids)
-    
-    def clear(self) -> None:
-        """Clear all messages (for testing)."""
-        self._messages.clear()
-        self._invocation_index.clear()
-
-
-__all__ = [
-    "RawMessageStore",
-    "StateBackendRawMessageStore",
-    "InMemoryRawMessageStore",
-]

aury/agents/middleware/raw_message.py

@@ -1,154 +0,0 @@
-"""RawMessageMiddleware - stores complete messages for HITL recovery.
-
-This middleware stores complete, untruncated messages to RawMessageStore.
-Works alongside MessageBackendMiddleware which stores truncated messages.
-
-Usage:
-    raw_store = InMemoryRawMessageStore()
-    raw_middleware = RawMessageMiddleware(raw_store, persist_raw=False)
-    
-    agent = ReactAgent.create(
-        llm=llm,
-        middlewares=[raw_middleware, MessageBackendMiddleware()],
-    )
-"""
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Any
-
-from .base import BaseMiddleware
-from .types import HookResult
-
-if TYPE_CHECKING:
-    from ..messages.raw_store import RawMessageStore
-    from ..core.state import State
-
-
-class RawMessageMiddleware(BaseMiddleware):
-    """Middleware that stores complete messages to RawMessageStore.
-    
-    Stores untruncated messages for:
-    - HITL recovery (restore exact conversation state)
-    - Full-context recall (when truncated history is insufficient)
-    
-    Messages are stored per invocation and can be cleaned up when
-    the invocation completes (controlled by persist_raw).
-    """
-    
-    def __init__(
-        self,
-        raw_store: "RawMessageStore",
-        persist_raw: bool = False,
-        state: "State | None" = None,
-    ):
-        """Initialize with RawMessageStore.
-        
-        Args:
-            raw_store: RawMessageStore for storing complete messages
-            persist_raw: Whether to keep messages after invocation completes.
-                        False = clean up after invocation (default)
-                        True = keep forever (for audit/recall)
-            state: State instance for storing message_ids in execution namespace.
-                  If provided, message IDs are automatically added to
-                  state.execution["message_ids"].
-        """
-        self.raw_store = raw_store
-        self.persist_raw = persist_raw
-        self.state = state
-        
-        # Track message IDs per invocation (for cleanup)
-        self._invocation_msg_ids: dict[str, list[str]] = {}
-    
-    def set_state(self, state: "State") -> None:
-        """Set state instance (can be set after construction).
-        
-        Args:
-            state: State instance for storing message_ids
-        """
-        self.state = state
-    
-    async def on_message_save(
-        self,
-        message: dict[str, Any],
-    ) -> dict[str, Any] | None:
-        """Store complete message to RawMessageStore.
-        
-        Args:
-            message: Complete message dict with 'role', 'content', etc.
-            
-        Returns:
-            The message with added 'raw_msg_id' field
-        """
-        from ..core.context import get_current_ctx_or_none
-        ctx = get_current_ctx_or_none()
-        invocation_id = ctx.invocation_id if ctx else ""
-        if not invocation_id:
-            return message
-        
-        # Store to raw store
-        msg_id = await self.raw_store.add(invocation_id, message)
-        
-        # Track for cleanup
-        if invocation_id not in self._invocation_msg_ids:
-            self._invocation_msg_ids[invocation_id] = []
-        self._invocation_msg_ids[invocation_id].append(msg_id)
-        
-        # Add to state.execution["message_ids"] if state is available
-        if self.state:
-            message_ids = self.state.execution.get("message_ids", [])
-            message_ids.append(msg_id)
-            self.state.execution["message_ids"] = message_ids
-        
-        # Add msg_id to message for downstream middlewares
-        message["raw_msg_id"] = msg_id
-        
-        return message
-    
-    async def on_agent_end(
-        self,
-        agent_id: str,
-        result: Any,
-    ) -> HookResult:
-        """Clean up raw messages when invocation completes.
-        
-        Only cleans up if persist_raw is False.
-        """
-        if self.persist_raw:
-            return HookResult.proceed()
-        
-        from ..core.context import get_current_ctx_or_none
-        ctx = get_current_ctx_or_none()
-        invocation_id = ctx.invocation_id if ctx else ""
-        if invocation_id:
-            await self._cleanup_invocation(invocation_id)
-        
-        return HookResult.proceed()
-    
-    async def _cleanup_invocation(self, invocation_id: str) -> int:
-        """Clean up raw messages for an invocation.
-        
-        Args:
-            invocation_id: Invocation ID to clean up
-            
-        Returns:
-            Number of messages deleted
-        """
-        # Remove from tracking
-        self._invocation_msg_ids.pop(invocation_id, None)
-        
-        # Delete from store
-        return await self.raw_store.delete_by_invocation(invocation_id)
-    
-    def get_message_ids(self, invocation_id: str) -> list[str]:
-        """Get tracked message IDs for an invocation.
-        
-        Args:
-            invocation_id: Invocation ID
-            
-        Returns:
-            List of message IDs
-        """
-        return self._invocation_msg_ids.get(invocation_id, []).copy()
-
-
-__all__ = ["RawMessageMiddleware"]