aury-agent 0.0.4__py3-none-any.whl

aury/agents/core/services/message.py

@@ -0,0 +1,53 @@
+"""Message service protocol."""
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Any, Protocol, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..types.block import BlockEvent, PersistedBlock
+    from ..types.message import Message
+
+
+class MessageService(Protocol):
+    """Protocol for message/block management.
+    
+    Handles message and block persistence.
+    """
+    
+    @abstractmethod
+    async def append(
+        self,
+        session_id: str,
+        invocation_id: str,
+        block: "BlockEvent",
+    ) -> "PersistedBlock":
+        """Append a block to the message history."""
+        ...
+    
+    @abstractmethod
+    async def list_blocks(
+        self,
+        session_id: str,
+        invocation_id: str | None = None,
+        limit: int = 100,
+    ) -> list["PersistedBlock"]:
+        """List blocks for a session/invocation."""
+        ...
+    
+    @abstractmethod
+    async def update_block(
+        self,
+        block_id: str,
+        data: dict[str, Any],
+    ) -> None:
+        """Update a block's data."""
+        ...
+    
+    @abstractmethod
+    async def get_block(self, block_id: str) -> "PersistedBlock | None":
+        """Get a specific block."""
+        ...
+
+
+__all__ = ["MessageService"]

aury/agents/core/services/session.py

@@ -0,0 +1,53 @@
+"""Session service protocol."""
+from __future__ import annotations
+
+from abc import abstractmethod
+from typing import Protocol, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..types.session import Session, ControlFrame
+
+
+class SessionService(Protocol):
+    """Protocol for session management.
+    
+    Handles session CRUD and control stack operations.
+    """
+    
+    @abstractmethod
+    async def create(self, root_agent_id: str, **kwargs) -> "Session":
+        """Create a new session."""
+        ...
+    
+    @abstractmethod
+    async def get(self, session_id: str) -> "Session | None":
+        """Get session by ID."""
+        ...
+    
+    @abstractmethod
+    async def update(self, session: "Session") -> None:
+        """Update session."""
+        ...
+    
+    @abstractmethod
+    async def delete(self, session_id: str) -> bool:
+        """Delete session."""
+        ...
+    
+    @abstractmethod
+    async def list(self, limit: int = 100, offset: int = 0) -> list["Session"]:
+        """List sessions."""
+        ...
+    
+    @abstractmethod
+    async def push_control(self, session_id: str, frame: "ControlFrame") -> None:
+        """Push control frame to session's control stack."""
+        ...
+    
+    @abstractmethod
+    async def pop_control(self, session_id: str) -> "ControlFrame | None":
+        """Pop control frame from session's control stack."""
+        ...
+
+
+__all__ = ["SessionService"]

aury/agents/core/signals.py

@@ -0,0 +1,109 @@
+"""Control flow signals for agent execution.
+
+These are NOT exceptions - they are control flow signals that inherit from
+BaseException to avoid being caught by generic `except Exception` handlers.
+
+Similar to KeyboardInterrupt and SystemExit, these signals control execution
+flow rather than indicate errors.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+
+class SuspendSignal(BaseException):
+    """Base class for suspension signals.
+    
+    Signals that execution should be suspended (not failed).
+    Inherits from BaseException so `except Exception` won't catch it.
+    
+    Usage:
+        try:
+            await agent.run()
+        except SuspendSignal as s:
+            # Handle suspension (HITL, pause, etc.)
+            handle_suspend(s)
+        except Exception as e:
+            # Handle actual errors
+            handle_error(e)
+    """
+    pass
+
+
+@dataclass
+class HITLSuspend(SuspendSignal):
+    """Signal for Human-in-the-Loop suspension.
+    
+    Raised when agent needs human input to continue.
+    Contains all information needed to:
+    1. Display the request to user
+    2. Resume execution after user responds
+    
+    Attributes:
+        request_id: Unique ID for matching response
+        request_type: Type of request (ask_user, confirm, form, etc.)
+        message: Display message for user
+        options: Optional list of choices
+        node_id: Workflow node ID if triggered from workflow
+        tool_name: Tool name if triggered from tool
+        block_id: Associated UI block ID
+        metadata: Additional context
+    """
+    request_id: str
+    request_type: str = "ask_user"
+    message: str | None = None
+    options: list[str] | None = None
+    node_id: str | None = None
+    tool_name: str | None = None
+    block_id: str | None = None
+    metadata: dict[str, Any] = field(default_factory=dict)
+    
+    def __post_init__(self):
+        # Initialize BaseException with a message
+        super().__init__(f"HITL suspend: {self.request_type} ({self.request_id})")
+    
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for storage/transmission."""
+        return {
+            "request_id": self.request_id,
+            "request_type": self.request_type,
+            "message": self.message,
+            "options": self.options,
+            "node_id": self.node_id,
+            "tool_name": self.tool_name,
+            "block_id": self.block_id,
+            "metadata": self.metadata,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "HITLSuspend":
+        """Create from dictionary."""
+        return cls(
+            request_id=data["request_id"],
+            request_type=data.get("request_type", "ask_user"),
+            message=data.get("message"),
+            options=data.get("options"),
+            node_id=data.get("node_id"),
+            tool_name=data.get("tool_name"),
+            block_id=data.get("block_id"),
+            metadata=data.get("metadata", {}),
+        )
+
+
+class PauseSuspend(SuspendSignal):
+    """Signal for manual pause (user-initiated).
+    
+    Raised when user requests to pause execution.
+    """
+    
+    def __init__(self, reason: str = "User requested pause"):
+        self.reason = reason
+        super().__init__(reason)
+
+
+__all__ = [
+    "SuspendSignal",
+    "HITLSuspend",
+    "PauseSuspend",
+]

aury/agents/core/state.py

@@ -0,0 +1,363 @@
+"""State management with checkpoint support.
+
+State provides layered storage with three partitions:
+- vars: User/developer variables for prompt formatting
+- data: Data flow shared between ReactAgent and WorkflowAgent
+- execution: Execution state (step, message_ids, pending_request, etc.)
+
+Supports path-based access, pattern matching, and checkpoint/restore.
+"""
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..backends.state import StateBackend
+
+
+class State:
+    """Layered state management with checkpoint support.
+    
+    State is NOT persisted on every set() call. Instead:
+    - set() writes to in-memory buffer
+    - checkpoint() persists buffer to backend
+    - restore() loads from backend
+    
+    This allows efficient batching and recovery from interrupts.
+    
+    Example:
+        state = State(backend, session_id)
+        await state.restore()  # Load existing state
+        
+        state.set("vars.user_name", "高鑫")
+        state.set("workflow.node1_output", result)
+        
+        await state.checkpoint()  # Persist changes
+    """
+    
+    def __init__(
+        self,
+        backend: "StateBackend",
+        session_id: str,
+        *,
+        initial_data: dict[str, Any] | None = None,
+    ):
+        """Initialize state.
+        
+        Args:
+            backend: Storage backend for persistence
+            session_id: Session identifier (used as storage key)
+            initial_data: Optional initial state data
+        """
+        self._backend = backend
+        self._session_id = session_id
+        self._invocation_id: str | None = None
+        self._data: dict[str, Any] = initial_data or {
+            "vars": {},
+            "data": {},
+            "execution": {},
+        }
+        self._dirty = False
+    
+    # ========== Partition Access ==========
+    
+    @property
+    def vars(self) -> dict[str, Any]:
+        """User/developer variables (for prompt formatting).
+        
+        Example:
+            prompt = template.format(**state.vars)
+        """
+        if "vars" not in self._data:
+            self._data["vars"] = {}
+        return self._data["vars"]
+    
+    @property
+    def data(self) -> dict[str, Any]:
+        """Data flow shared between ReactAgent and WorkflowAgent.
+        
+        Used for:
+        - Workflow node outputs
+        - Tool results that need to be shared
+        - Any data passed between agents
+        
+        Example:
+            state.data["search_results"] = results
+            state.data["node1_output"] = output
+        """
+        if "data" not in self._data:
+            self._data["data"] = {}
+        return self._data["data"]
+    
+    @property
+    def execution(self) -> dict[str, Any]:
+        """Execution state for current invocation.
+        
+        Used for:
+        - step: Current step number
+        - message_ids: References to raw messages
+        - pending_request: HITL request if suspended
+        - current_node: Workflow current node
+        - completed_nodes: Workflow completed nodes
+        
+        Example:
+            state.execution["step"] = 5
+            state.execution["message_ids"] = ["msg_001", "msg_002"]
+        """
+        if "execution" not in self._data:
+            self._data["execution"] = {}
+        return self._data["execution"]
+    
+    # ========== Path Operations ==========
+    
+    def get(self, path: str, default: Any = None) -> Any:
+        """Get value by path.
+        
+        Args:
+            path: Dot-separated path (e.g., "vars.user_name")
+            default: Default value if path not found
+            
+        Returns:
+            Value at path, or default
+            
+        Example:
+            state.get("vars.user_name")     # → "高鑫"
+            state.get("workflow.node1")     # → {...}
+            state.get("agent.missing", 0)   # → 0
+        """
+        parts = path.split(".")
+        current = self._data
+        
+        for part in parts:
+            if isinstance(current, dict) and part in current:
+                current = current[part]
+            else:
+                return default
+        
+        return current
+    
+    def set(self, path: str, value: Any) -> None:
+        """Set value by path (writes to buffer, not persisted).
+        
+        Args:
+            path: Dot-separated path (e.g., "vars.user_name")
+            value: Value to set
+            
+        Example:
+            state.set("vars.user_name", "高鑫")
+            state.set("workflow.node1_output", result)
+        """
+        parts = path.split(".")
+        current = self._data
+        
+        # Navigate to parent, creating dicts as needed
+        for part in parts[:-1]:
+            if part not in current:
+                current[part] = {}
+            current = current[part]
+        
+        # Set value
+        current[parts[-1]] = value
+        self._dirty = True
+    
+    def delete(self, path: str) -> bool:
+        """Delete value by path.
+        
+        Args:
+            path: Dot-separated path
+            
+        Returns:
+            True if deleted, False if not found
+        """
+        parts = path.split(".")
+        current = self._data
+        
+        # Navigate to parent
+        for part in parts[:-1]:
+            if isinstance(current, dict) and part in current:
+                current = current[part]
+            else:
+                return False
+        
+        # Delete
+        if isinstance(current, dict) and parts[-1] in current:
+            del current[parts[-1]]
+            self._dirty = True
+            return True
+        
+        return False
+    
+    def match(self, pattern: str) -> dict[str, Any]:
+        """Get values matching prefix pattern.
+        
+        Args:
+            pattern: Prefix pattern ending with '*' (e.g., "workflow.*")
+            
+        Returns:
+            Dict of matching key-value pairs (keys without prefix)
+            
+        Example:
+            state.match("workflow.*")  # → {"node1_output": ..., "node2_output": ...}
+            state.match("vars.*")      # → {"user_name": ..., "project_name": ...}
+        """
+        if not pattern.endswith("*"):
+            # Exact match, return single value wrapped
+            val = self.get(pattern)
+            return {pattern.split(".")[-1]: val} if val is not None else {}
+        
+        # Prefix match
+        prefix = pattern[:-1]  # Remove '*'
+        if prefix.endswith("."):
+            prefix = prefix[:-1]  # Remove trailing '.'
+        
+        parent = self.get(prefix)
+        if isinstance(parent, dict):
+            return parent.copy()
+        
+        return {}
+    
+    def has(self, path: str) -> bool:
+        """Check if path exists.
+        
+        Args:
+            path: Dot-separated path
+            
+        Returns:
+            True if path exists
+        """
+        return self.get(path) is not None
+    
+    def clear(self, partition: str | None = None) -> None:
+        """Clear state.
+        
+        Args:
+            partition: If provided, clear only that partition.
+                       If None, clear all partitions.
+        """
+        if partition:
+            if partition in self._data:
+                self._data[partition] = {}
+                self._dirty = True
+        else:
+            self._data = {
+                "vars": {},
+                "data": {},
+                "execution": {},
+            }
+            self._dirty = True
+    
+    # ========== Persistence ==========
+    
+    def set_invocation_id(self, invocation_id: str) -> None:
+        """Set current invocation ID for persistence.
+        
+        State is persisted per invocation, not per session.
+        """
+        self._invocation_id = invocation_id
+    
+    @property
+    def is_dirty(self) -> bool:
+        """Check if state has unsaved changes."""
+        return self._dirty
+    
+    async def checkpoint(self) -> None:
+        """Persist current state to backend.
+        
+        State is saved per invocation_id (not session_id).
+        
+        Call at key points:
+        - After step completion
+        - After tool execution
+        - Before HITL suspend
+        """
+        if not self._dirty:
+            return
+        
+        # Use invocation_id if set, otherwise fall back to session_id
+        key = self._invocation_id or self._session_id
+        await self._backend.set("state", key, self._data)
+        self._dirty = False
+    
+    async def restore(self, invocation_id: str | None = None) -> bool:
+        """Restore state from backend.
+        
+        Args:
+            invocation_id: Specific invocation to restore. If None, uses current.
+        
+        Returns:
+            True if state was restored, False if no saved state
+        """
+        key = invocation_id or self._invocation_id or self._session_id
+        data = await self._backend.get("state", key)
+        if data:
+            self._data = data
+            # Ensure all partitions exist
+            for partition in ("vars", "data", "execution"):
+                if partition not in self._data:
+                    self._data[partition] = {}
+            self._dirty = False
+            if invocation_id:
+                self._invocation_id = invocation_id
+            return True
+        return False
+    
+    # ========== Utility ==========
+    
+    def to_dict(self) -> dict[str, Any]:
+        """Export state as dict."""
+        return self._data.copy()
+    
+    def update(self, data: dict[str, Any]) -> None:
+        """Bulk update state.
+        
+        Args:
+            data: Dict with partition keys (vars, data, execution)
+        """
+        for key in ("vars", "data", "execution"):
+            if key in data:
+                self._data[key].update(data[key])
+                self._dirty = True
+    
+    def __repr__(self) -> str:
+        dirty_marker = " (dirty)" if self._dirty else ""
+        return f"<State session={self._session_id}{dirty_marker}>"
+
+
+# Convenience type for state isolation in SubAgents
+class StateSnapshot:
+    """Immutable snapshot of state for isolation."""
+    
+    def __init__(self, data: dict[str, Any]):
+        self._data = data
+    
+    @property
+    def vars(self) -> dict[str, Any]:
+        return self._data.get("vars", {}).copy()
+    
+    @property
+    def data(self) -> dict[str, Any]:
+        return self._data.get("data", {}).copy()
+    
+    @property
+    def execution(self) -> dict[str, Any]:
+        return self._data.get("execution", {}).copy()
+    
+    def get(self, path: str, default: Any = None) -> Any:
+        parts = path.split(".")
+        current = self._data
+        for part in parts:
+            if isinstance(current, dict) and part in current:
+                current = current[part]
+            else:
+                return default
+        return current
+    
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "vars": self.vars,
+            "data": self.data,
+            "execution": self.execution,
+        }
+
+
+__all__ = ["State", "StateSnapshot"]

aury/agents/core/types/__init__.py

@@ -0,0 +1,107 @@
+"""Type definitions for Aury Agent Framework."""
+from .session import (
+    generate_id,
+    InvocationState,
+    InvocationMode,
+    ControlFrame,
+    Session,
+    Invocation,
+)
+from .subagent import (
+    SubAgentMode,
+    SubAgentInput,
+    SubAgentMetadata,
+    SubAgentResult,
+)
+from .block import (
+    BlockKind,
+    BlockOp,
+    Persistence,
+    ActorInfo,
+    BlockEvent,
+    PersistedBlock,
+    BlockHandle,
+    BlockAggregator,
+    BlockMerger,
+    register_merger,
+    get_merger,
+    # Helper functions
+    text_block,
+    text_delta,
+    thinking_block,
+    thinking_delta,
+    tool_use_block,
+    tool_use_patch,
+    tool_result_block,
+    error_block,
+)
+from .message import (
+    MessageRole,
+    Message,
+    PromptInput,
+)
+from .tool import (
+    ToolInfo,
+    ToolContext,
+    ToolResult,
+    ToolInvocationState,
+    ToolInvocation,
+    BaseTool,
+    ToolConfig,
+)
+from .action import (
+    ActionType,
+    ActionEvent,
+    ActionCollector,
+)
+
+__all__ = [
+    # Session
+    "generate_id",
+    "InvocationState",
+    "InvocationMode",
+    "ControlFrame",
+    "Session",
+    "Invocation",
+    # SubAgent
+    "SubAgentMode",
+    "SubAgentInput",
+    "SubAgentMetadata",
+    "SubAgentResult",
+    # Block
+    "BlockKind",
+    "BlockOp",
+    "Persistence",
+    "ActorInfo",
+    "BlockEvent",
+    "PersistedBlock",
+    "BlockHandle",
+    "BlockAggregator",
+    "BlockMerger",
+    "register_merger",
+    "get_merger",
+    "text_block",
+    "text_delta",
+    "thinking_block",
+    "thinking_delta",
+    "tool_use_block",
+    "tool_use_patch",
+    "tool_result_block",
+    "error_block",
+    # Message
+    "MessageRole",
+    "Message",
+    "PromptInput",
+    # Tool
+    "ToolInfo",
+    "ToolContext",
+    "ToolResult",
+    "ToolInvocationState",
+    "ToolInvocation",
+    "BaseTool",
+    "ToolConfig",
+    # Action
+    "ActionType",
+    "ActionEvent",
+    "ActionCollector",
+]