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

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"]