aury-agent 0.0.4__py3-none-any.whl

aury/agents/middleware/raw_message.py

@@ -0,0 +1,153 @@
+"""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],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Store complete message to RawMessageStore.
+        
+        Args:
+            message: Complete message dict with 'role', 'content', etc.
+            context: Execution context with 'invocation_id', etc.
+            
+        Returns:
+            The message with added 'raw_msg_id' field
+        """
+        invocation_id = context.get("invocation_id", "")
+        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,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Clean up raw messages when invocation completes.
+        
+        Only cleans up if persist_raw is False.
+        """
+        if self.persist_raw:
+            return HookResult.proceed()
+        
+        invocation_id = context.get("invocation_id", "")
+        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"]

aury/agents/middleware/truncation.py

@@ -0,0 +1,139 @@
+"""MessageTruncationMiddleware - truncates large message content before persistence.
+
+This middleware intercepts on_message_save and truncates content
+that exceeds configured limits, ensuring historical messages don't
+consume excessive storage.
+
+Current invocation messages remain complete in State (for recovery),
+only persisted messages get truncated.
+"""
+from __future__ import annotations
+
+from typing import Any
+
+from .base import BaseMiddleware
+
+
+class MessageTruncationMiddleware(BaseMiddleware):
+    """Truncates large message content before persistence.
+    
+    Use this middleware BEFORE MessageBackendMiddleware in the chain:
+    
+        middleware = MiddlewareChain([
+            MessageTruncationMiddleware(max_content_length=2000),
+            MessageBackendMiddleware(),
+        ])
+    
+    Features:
+    - Truncates string content exceeding max_content_length
+    - Truncates individual items in list content (tool_use, tool_result)
+    - Adds "[truncated]" marker to truncated content
+    - Configurable truncation threshold
+    """
+    
+    def __init__(
+        self,
+        max_content_length: int = 2000,
+        truncate_marker: str = "... [truncated]",
+    ):
+        """Initialize MessageTruncationMiddleware.
+        
+        Args:
+            max_content_length: Maximum content length before truncation
+            truncate_marker: Marker appended to truncated content
+        """
+        self.max_content_length = max_content_length
+        self.truncate_marker = truncate_marker
+    
+    async def on_message_save(
+        self,
+        message: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Truncate message content before saving.
+        
+        Args:
+            message: Message dict with 'role', 'content', etc.
+            context: Execution context
+            
+        Returns:
+            Modified message with truncated content
+        """
+        content = message.get("content")
+        if content is None:
+            return message
+        
+        # Create a copy to avoid mutating original
+        truncated_message = message.copy()
+        truncated_message["content"] = self._truncate_content(content)
+        
+        return truncated_message
+    
+    def _truncate_content(self, content: Any) -> Any:
+        """Truncate content based on type.
+        
+        Args:
+            content: String or list content
+            
+        Returns:
+            Truncated content
+        """
+        if isinstance(content, str):
+            return self._truncate_string(content)
+        
+        if isinstance(content, list):
+            return self._truncate_list(content)
+        
+        # Unknown type, return as-is
+        return content
+    
+    def _truncate_string(self, text: str) -> str:
+        """Truncate string content."""
+        if len(text) <= self.max_content_length:
+            return text
+        
+        # Keep first part, add marker
+        return text[:self.max_content_length] + self.truncate_marker
+    
+    def _truncate_list(self, items: list[Any]) -> list[Any]:
+        """Truncate list content (tool_use, tool_result, etc.)."""
+        truncated_items = []
+        
+        for item in items:
+            if isinstance(item, dict):
+                truncated_item = item.copy()
+                
+                # Truncate text content
+                if "text" in truncated_item and isinstance(truncated_item["text"], str):
+                    truncated_item["text"] = self._truncate_string(truncated_item["text"])
+                
+                # Truncate tool_result content
+                if "content" in truncated_item and isinstance(truncated_item["content"], str):
+                    truncated_item["content"] = self._truncate_string(truncated_item["content"])
+                
+                # Truncate tool_use input (arguments)
+                if "input" in truncated_item and isinstance(truncated_item["input"], dict):
+                    truncated_item["input"] = self._truncate_dict(truncated_item["input"])
+                
+                truncated_items.append(truncated_item)
+            else:
+                truncated_items.append(item)
+        
+        return truncated_items
+    
+    def _truncate_dict(self, data: dict[str, Any]) -> dict[str, Any]:
+        """Truncate string values in a dict (for tool arguments)."""
+        truncated = {}
+        for key, value in data.items():
+            if isinstance(value, str):
+                truncated[key] = self._truncate_string(value)
+            elif isinstance(value, dict):
+                truncated[key] = self._truncate_dict(value)
+            elif isinstance(value, list):
+                truncated[key] = self._truncate_list(value)
+            else:
+                truncated[key] = value
+        return truncated
+
+
+__all__ = ["MessageTruncationMiddleware"]

aury/agents/middleware/types.py

@@ -0,0 +1,81 @@
+"""Middleware types and data classes."""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+
+class TriggerMode(Enum):
+    """Middleware trigger mode for streaming."""
+    EVERY_TOKEN = "every_token"  # Trigger on every token
+    EVERY_N_TOKENS = "every_n_tokens"  # Trigger every N tokens
+    ON_BOUNDARY = "on_boundary"  # Trigger on sentence/paragraph boundaries
+
+
+class HookAction(Enum):
+    """Control flow action returned by lifecycle hooks.
+    
+    CONTINUE: Proceed with normal execution
+    SKIP: Skip the current operation (tool call, etc.)
+    RETRY: Retry the current operation (with modified params)
+    STOP: Stop the agent execution entirely
+    """
+    CONTINUE = "continue"
+    SKIP = "skip"
+    RETRY = "retry"
+    STOP = "stop"
+
+
+@dataclass
+class HookResult:
+    """Result from a lifecycle hook.
+    
+    Attributes:
+        action: Control flow action
+        modified_data: Modified data (for RETRY action)
+        message: Optional message explaining the action
+        metadata: Additional metadata
+    """
+    action: HookAction = HookAction.CONTINUE
+    modified_data: Any = None
+    message: str | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+    
+    @classmethod
+    def proceed(cls) -> "HookResult":
+        """Continue with normal execution."""
+        return cls(action=HookAction.CONTINUE)
+    
+    @classmethod
+    def skip(cls, message: str | None = None) -> "HookResult":
+        """Skip the current operation."""
+        return cls(action=HookAction.SKIP, message=message)
+    
+    @classmethod
+    def retry(cls, modified_data: Any = None, message: str | None = None) -> "HookResult":
+        """Retry with optional modified data."""
+        return cls(action=HookAction.RETRY, modified_data=modified_data, message=message)
+    
+    @classmethod
+    def stop(cls, message: str | None = None) -> "HookResult":
+        """Stop agent execution."""
+        return cls(action=HookAction.STOP, message=message)
+
+
+@dataclass
+class MiddlewareConfig:
+    """Middleware configuration."""
+    trigger_mode: TriggerMode = TriggerMode.EVERY_TOKEN
+    trigger_n: int = 10  # For EVERY_N_TOKENS mode
+    async_mode: bool = True  # Execute asynchronously
+    priority: int = 0  # Lower priority runs first
+    inherit: bool = True  # Whether to pass to sub-agents (can be overridden at use() time)
+
+
+__all__ = [
+    "TriggerMode",
+    "HookAction",
+    "HookResult",
+    "MiddlewareConfig",
+]

aury/agents/plugin.py

@@ -0,0 +1,162 @@
+"""Plugin system for packaging providers, middlewares and tools.
+
+The Plugin class provides an optional assembly layer that packages:
+- ContextProviders (context sources)
+- Middlewares (execution hooks)
+- Tools, Skills, SubAgents (direct contributions)
+
+Design Goals:
+1. Simple case: Use plugins for packaged functionality
+2. Flexible case: Use providers/middlewares directly
+3. Mixed case: Combine plugins with custom providers/middlewares
+
+Example Usage:
+    # Simple: use plugins
+    agent = ReactAgent.create(
+        llm=llm,
+        plugins=[CodePlugin()],
+    )
+
+    # Flexible: use providers/middlewares directly
+    agent = ReactAgent.create(
+        llm=llm,
+        context_providers=[MyCustomProvider()],
+        middlewares=[MyMiddleware()],
+    )
+
+    # Mixed + custom
+    agent = ReactAgent.create(
+        llm=llm,
+        plugins=[CodePlugin()],
+        context_providers=[MyRAGProvider()],
+        subagents=[researcher_config],
+        delegate_tool_class=MyDelegateTool,
+    )
+
+Inheritance Rules:
+- Plugin.inherit: Whether this plugin should be inherited by sub-agents
+- AgentConfig.inherit_plugins: Whether to inherit plugins from parent
+- Both must be True for inheritance to occur
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .core.types.tool import BaseTool
+    from .skill import Skill
+    from .backends.subagent import AgentConfig as SubAgentConfig
+    from .context_providers.base import ContextProvider
+    from .middleware.base import Middleware
+
+
+# =============================================================================
+# Plugin class
+# =============================================================================
+
+# @dataclass
+# class Plugin:
+#     """Plugin packages providers, middlewares, tools, skills, and subagents.
+#     
+#     Attributes:
+#         name: Plugin identifier
+#         priority: Execution order (lower runs first, default 100)
+#         inherit: Whether this plugin should be inherited by sub-agents
+#         tools: Tools provided by this plugin
+#         skills: Skills provided by this plugin
+#         subagents: SubAgent configs provided by this plugin
+#         providers: ContextProviders provided by this plugin
+#         middlewares: Middlewares provided by this plugin
+#     
+#     Priority Guidelines:
+#         0-50: Core/framework plugins (run first)
+#         50-100: Standard plugins
+#         100+: User plugins (run last)
+#     
+#     Example:
+#         @dataclass
+#         class CodePlugin(Plugin):
+#             name: str = "code"
+#             priority: int = 50
+#             inherit: bool = True
+#             
+#             def __post_init__(self):
+#                 self.tools = [
+#                     FileReadTool(),
+#                     FileWriteTool(),
+#                     ShellTool(),
+#                 ]
+#                 self.providers = [
+#                     ProjectContextProvider(),
+#                 ]
+#                 self.middlewares = [
+#                     CodeReviewMiddleware(),
+#                 ]
+#     """
+#     
+#     name: str
+#     priority: int = 100
+#     inherit: bool = False
+#     
+#     tools: list["BaseTool"] = field(default_factory=list)
+#     skills: list["Skill"] = field(default_factory=list)
+#     subagents: list["SubAgentConfig"] = field(default_factory=list)
+#     providers: list["ContextProvider"] = field(default_factory=list)
+#     middlewares: list["Middleware"] = field(default_factory=list)
+
+
+# =============================================================================
+# PluginChain for managing multiple plugins
+# =============================================================================
+
+# class PluginChain:
+#     """Chain of plugins sorted by priority.
+#     
+#     Provides methods to:
+#     - Collect all tools from plugins
+#     - Collect all providers from plugins
+#     - Collect all middlewares from plugins
+#     - Get inheritable plugins for sub-agents
+#     """
+#     
+#     def __init__(self, plugins: list[Plugin] | None = None):
+#         self._plugins: list[Plugin] = []
+#         if plugins:
+#             for p in plugins:
+#                 self.add(p)
+#     
+#     def add(self, plugin: Plugin) -> "PluginChain":
+#         """Add plugin and maintain sorted order by priority."""
+#         self._plugins.append(plugin)
+#         self._plugins.sort(key=lambda p: p.priority)
+#         return self
+#     
+#     def collect_tools(self) -> list["BaseTool"]:
+#         """Collect all tools from all plugins."""
+#         tools = []
+#         for p in self._plugins:
+#             tools.extend(p.tools)
+#         return tools
+#     
+#     def collect_providers(self) -> list["ContextProvider"]:
+#         """Collect all providers from all plugins."""
+#         providers = []
+#         for p in self._plugins:
+#             providers.extend(p.providers)
+#         return providers
+#     
+#     def collect_middlewares(self) -> list["Middleware"]:
+#         """Collect all middlewares from all plugins."""
+#         middlewares = []
+#         for p in self._plugins:
+#             middlewares.extend(p.middlewares)
+#         return middlewares
+#     
+#     def get_inheritable(self) -> list[Plugin]:
+#         """Get plugins that should be inherited by sub-agents."""
+#         return [p for p in self._plugins if p.inherit]
+
+
+__all__: list[str] = []  # Classes are commented out

aury/agents/react/__init__.py

@@ -0,0 +1,4 @@
+"""React Agent implementation."""
+from .agent import ReactAgent
+
+__all__ = ["ReactAgent"]