aury-agent 0.0.4__py3-none-any.whl

aury/agents/messages/store.py

@@ -0,0 +1,118 @@
+"""Message store protocol and implementations.
+
+Note: For production use, prefer MessageBackend (backends/message/).
+This module provides a simple protocol and in-memory implementation for testing.
+"""
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from .types import Message
+
+
+@runtime_checkable
+class MessageStore(Protocol):
+    """Protocol for message storage.
+    
+    Note: For production, use MessageBackend instead.
+    This protocol is kept for backward compatibility.
+    """
+    
+    async def add(
+        self,
+        session_id: str,
+        message: Message,
+        namespace: str | None = None,
+    ) -> None:
+        """Add a message to session history."""
+        ...
+    
+    async def get_all(
+        self,
+        session_id: str,
+        namespace: str | None = None,
+    ) -> list[Message]:
+        """Get all messages for a session."""
+        ...
+    
+    async def get_recent(
+        self,
+        session_id: str,
+        limit: int,
+        namespace: str | None = None,
+    ) -> list[Message]:
+        """Get recent messages for a session."""
+        ...
+    
+    async def delete_by_invocation(
+        self,
+        session_id: str,
+        invocation_id: str,
+        namespace: str | None = None,
+    ) -> int:
+        """Delete messages by invocation ID."""
+        ...
+
+
+class InMemoryMessageStore:
+    """In-memory message store for testing."""
+    
+    def __init__(self) -> None:
+        # Key format: "session_id" or "session_id:namespace"
+        self._messages: dict[str, list[Message]] = {}
+    
+    def _make_key(self, session_id: str, namespace: str | None) -> str:
+        if namespace:
+            return f"{session_id}:{namespace}"
+        return session_id
+    
+    async def add(
+        self,
+        session_id: str,
+        message: Message,
+        namespace: str | None = None,
+    ) -> None:
+        key = self._make_key(session_id, namespace)
+        if key not in self._messages:
+            self._messages[key] = []
+        self._messages[key].append(message)
+    
+    async def get_all(
+        self,
+        session_id: str,
+        namespace: str | None = None,
+    ) -> list[Message]:
+        key = self._make_key(session_id, namespace)
+        return self._messages.get(key, []).copy()
+    
+    async def get_recent(
+        self,
+        session_id: str,
+        limit: int,
+        namespace: str | None = None,
+    ) -> list[Message]:
+        key = self._make_key(session_id, namespace)
+        messages = self._messages.get(key, [])
+        return messages[-limit:] if limit else messages.copy()
+    
+    async def delete_by_invocation(
+        self,
+        session_id: str,
+        invocation_id: str,
+        namespace: str | None = None,
+    ) -> int:
+        key = self._make_key(session_id, namespace)
+        if key not in self._messages:
+            return 0
+        
+        original = self._messages[key]
+        self._messages[key] = [
+            m for m in original if m.invocation_id != invocation_id
+        ]
+        return len(original) - len(self._messages[key])
+
+
+__all__ = [
+    "MessageStore",
+    "InMemoryMessageStore",
+]

aury/agents/messages/types.py

@@ -0,0 +1,88 @@
+"""Message type definitions.
+
+Core types for the message system:
+- Message: A single message in conversation history
+- MessageRole: User/Assistant/Tool/System
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+
+class MessageRole(Enum):
+    """Message role types."""
+    USER = "user"
+    ASSISTANT = "assistant"
+    TOOL = "tool"
+    SYSTEM = "system"
+
+
+@dataclass
+class Message:
+    """A message in conversation history.
+    
+    Attributes:
+        role: Message role (user/assistant/tool/system)
+        content: Message content (string or content parts)
+        invocation_id: Which invocation this message belongs to
+        tool_call_id: Tool call ID (for tool messages)
+        created_at: When the message was created
+        metadata: Additional metadata
+    """
+    role: str
+    content: str | list[dict[str, Any]]
+    invocation_id: str = ""
+    tool_call_id: str | None = None
+    created_at: datetime = field(default_factory=datetime.now)
+    metadata: dict[str, Any] = field(default_factory=dict)
+    
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dict for storage."""
+        result = {
+            "role": self.role,
+            "content": self.content,
+            "invocation_id": self.invocation_id,
+            "created_at": self.created_at.isoformat(),
+        }
+        if self.tool_call_id:
+            result["tool_call_id"] = self.tool_call_id
+        if self.metadata:
+            result["metadata"] = self.metadata
+        return result
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "Message":
+        """Create from dict."""
+        created_at = data.get("created_at")
+        if isinstance(created_at, str):
+            created_at = datetime.fromisoformat(created_at)
+        elif created_at is None:
+            created_at = datetime.now()
+        
+        return cls(
+            role=data["role"],
+            content=data["content"],
+            invocation_id=data.get("invocation_id", ""),
+            tool_call_id=data.get("tool_call_id"),
+            created_at=created_at,
+            metadata=data.get("metadata", {}),
+        )
+    
+    def to_llm_format(self) -> dict[str, Any]:
+        """Convert to LLM message format."""
+        result: dict[str, Any] = {
+            "role": self.role,
+            "content": self.content,
+        }
+        if self.tool_call_id:
+            result["tool_call_id"] = self.tool_call_id
+        return result
+
+
+__all__ = [
+    "MessageRole",
+    "Message",
+]

aury/agents/middleware/__init__.py

@@ -0,0 +1,31 @@
+"""Middleware system for request/response processing.
+
+Middleware provides hooks for intercepting and modifying:
+- LLM requests/responses
+- Agent lifecycle events
+- Tool execution
+- Sub-agent delegation
+- Message persistence
+"""
+from .types import TriggerMode, HookAction, HookResult, MiddlewareConfig
+from .base import Middleware, BaseMiddleware
+from .chain import MiddlewareChain
+from .message_container import MessageContainerMiddleware
+from .message import MessageBackendMiddleware
+from .truncation import MessageTruncationMiddleware
+from .raw_message import RawMessageMiddleware
+
+__all__ = [
+    "TriggerMode",
+    "HookAction",
+    "HookResult",
+    "MiddlewareConfig",
+    "Middleware",
+    "BaseMiddleware",
+    "MiddlewareChain",
+    # Default middlewares
+    "MessageContainerMiddleware",
+    "MessageBackendMiddleware",
+    "MessageTruncationMiddleware",
+    "RawMessageMiddleware",
+]

aury/agents/middleware/base.py

@@ -0,0 +1,341 @@
+"""Middleware protocol and base implementation."""
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable, TYPE_CHECKING
+
+from .types import HookResult, MiddlewareConfig
+
+if TYPE_CHECKING:
+    from ..core.types.tool import BaseTool, ToolResult
+
+
+@runtime_checkable
+class Middleware(Protocol):
+    """Middleware protocol for request/response processing.
+    
+    Includes both LLM request/response hooks and agent lifecycle hooks.
+    """
+    
+    @property
+    def config(self) -> MiddlewareConfig:
+        """Get middleware configuration."""
+        ...
+    
+    # ========== LLM Request/Response Hooks ==========
+    
+    async def on_request(
+        self,
+        request: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Process request before LLM call.
+        
+        Args:
+            request: The request to process
+            context: Execution context
+            
+        Returns:
+            Modified request, or None to skip further processing
+        """
+        ...
+    
+    async def on_response(
+        self,
+        response: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Process response after LLM call.
+        
+        Args:
+            response: The response to process
+            context: Execution context
+            
+        Returns:
+            Modified response, or None to skip further processing
+        """
+        ...
+    
+    async def on_error(
+        self,
+        error: Exception,
+        context: dict[str, Any],
+    ) -> Exception | None:
+        """Handle errors.
+        
+        Args:
+            error: The exception that occurred
+            context: Execution context
+            
+        Returns:
+            Modified exception, or None to suppress
+        """
+        ...
+    
+    async def on_model_stream(
+        self,
+        chunk: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Process streaming chunk (triggered by trigger_mode).
+        
+        Args:
+            chunk: The streaming chunk
+            context: Execution context
+            
+        Returns:
+            Modified chunk, or None to skip further processing
+        """
+        ...
+    
+    # ========== Agent Lifecycle Hooks ==========
+    
+    async def on_agent_start(
+        self,
+        agent_id: str,
+        input_data: Any,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Called when agent starts processing.
+        
+        Args:
+            agent_id: The agent identifier
+            input_data: Input to the agent
+            context: Execution context
+            
+        Returns:
+            HookResult controlling execution flow
+        """
+        ...
+    
+    async def on_agent_end(
+        self,
+        agent_id: str,
+        result: Any,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Called when agent completes processing.
+        
+        Args:
+            agent_id: The agent identifier
+            result: Agent's result
+            context: Execution context
+            
+        Returns:
+            HookResult (only CONTINUE/STOP meaningful here)
+        """
+        ...
+    
+    async def on_tool_call(
+        self,
+        tool: "BaseTool",
+        params: dict[str, Any],
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Called before tool execution.
+        
+        Args:
+            tool: The tool to be called
+            params: Tool parameters
+            context: Execution context
+            
+        Returns:
+            HookResult - SKIP to skip tool, RETRY to modify params
+        """
+        ...
+    
+    async def on_tool_end(
+        self,
+        tool: "BaseTool",
+        result: "ToolResult",
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Called after tool execution.
+        
+        Args:
+            tool: The tool that was called
+            result: Tool execution result
+            context: Execution context
+            
+        Returns:
+            HookResult - RETRY to re-execute tool
+        """
+        ...
+    
+    async def on_subagent_start(
+        self,
+        parent_agent_id: str,
+        child_agent_id: str,
+        mode: str,  # "embedded" or "delegated"
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Called when delegating to a sub-agent.
+        
+        Args:
+            parent_agent_id: Parent agent identifier
+            child_agent_id: Child agent identifier
+            mode: Delegation mode
+            context: Execution context
+            
+        Returns:
+            HookResult - SKIP to skip delegation
+        """
+        ...
+    
+    async def on_subagent_end(
+        self,
+        parent_agent_id: str,
+        child_agent_id: str,
+        result: Any,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Called when sub-agent completes.
+        
+        Args:
+            parent_agent_id: Parent agent identifier
+            child_agent_id: Child agent identifier
+            result: Sub-agent's result
+            context: Execution context
+            
+        Returns:
+            HookResult (for post-processing)
+        """
+        ...
+    
+    async def on_message_save(
+        self,
+        message: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Called before saving a message to history.
+        
+        Allows middlewares to transform, filter, or block messages
+        before they are persisted.
+        
+        Args:
+            message: Message dict with 'role', 'content', etc.
+            context: Execution context
+            
+        Returns:
+            Modified message, or None to skip saving
+        """
+        ...
+
+
+class BaseMiddleware:
+    """Base middleware implementation with sensible defaults.
+    
+    Subclass and override specific hooks as needed.
+    All hooks have sensible pass-through defaults.
+    """
+    
+    _config: MiddlewareConfig = MiddlewareConfig()
+    
+    @property
+    def config(self) -> MiddlewareConfig:
+        return self._config
+    
+    # ========== LLM Request/Response Hooks ==========
+    
+    async def on_request(
+        self,
+        request: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Default: pass through."""
+        return request
+    
+    async def on_response(
+        self,
+        response: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Default: pass through."""
+        return response
+    
+    async def on_error(
+        self,
+        error: Exception,
+        context: dict[str, Any],
+    ) -> Exception | None:
+        """Default: re-raise error."""
+        return error
+    
+    async def on_model_stream(
+        self,
+        chunk: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Default: pass through."""
+        return chunk
+    
+    # ========== Agent Lifecycle Hooks ==========
+    
+    async def on_agent_start(
+        self,
+        agent_id: str,
+        input_data: Any,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Default: continue."""
+        return HookResult.proceed()
+    
+    async def on_agent_end(
+        self,
+        agent_id: str,
+        result: Any,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Default: continue."""
+        return HookResult.proceed()
+    
+    async def on_tool_call(
+        self,
+        tool: "BaseTool",
+        params: dict[str, Any],
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Default: continue."""
+        return HookResult.proceed()
+    
+    async def on_tool_end(
+        self,
+        tool: "BaseTool",
+        result: "ToolResult",
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Default: continue."""
+        return HookResult.proceed()
+    
+    async def on_subagent_start(
+        self,
+        parent_agent_id: str,
+        child_agent_id: str,
+        mode: str,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Default: continue."""
+        return HookResult.proceed()
+    
+    async def on_subagent_end(
+        self,
+        parent_agent_id: str,
+        child_agent_id: str,
+        result: Any,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Default: continue."""
+        return HookResult.proceed()
+    
+    async def on_message_save(
+        self,
+        message: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Default: pass through."""
+        return message
+
+
+__all__ = [
+    "Middleware",
+    "BaseMiddleware",
+]