aury-agent 0.0.4__py3-none-any.whl

aury/agents/core/types/session.py

@@ -0,0 +1,257 @@
+"""Session and Invocation data structures."""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any
+from uuid import uuid4
+
+
+def generate_id(prefix: str = "") -> str:
+    """Generate a unique ID with optional prefix."""
+    uid = uuid4().hex[:12]
+    return f"{prefix}_{uid}" if prefix else uid
+
+
+class InvocationState(Enum):
+    """Invocation execution state."""
+    PENDING = "pending"
+    RUNNING = "running"
+    SUSPENDED = "suspended"  # HITL waiting
+    PAUSED = "paused"  # User paused
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+    ABORTED = "aborted"  # User stopped
+    SWITCHED = "switched"  # User switched agent
+
+
+class InvocationMode(Enum):
+    """Invocation mode."""
+    ROOT = "root"  # Root invocation
+    DELEGATED = "delegated"  # Delegated sub-agent
+    # Note: EMBEDDED doesn't create new Invocation, uses parent's
+
+
+@dataclass
+class ControlFrame:
+    """Control frame for tracking delegated agent in session."""
+    agent_id: str
+    invocation_id: str
+    entered_at: datetime = field(default_factory=datetime.now)
+    parent_invocation_id: str | None = None
+    
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "agent_id": self.agent_id,
+            "invocation_id": self.invocation_id,
+            "entered_at": self.entered_at.isoformat(),
+            "parent_invocation_id": self.parent_invocation_id,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ControlFrame":
+        return cls(
+            agent_id=data["agent_id"],
+            invocation_id=data["invocation_id"],
+            entered_at=datetime.fromisoformat(data["entered_at"]),
+            parent_invocation_id=data.get("parent_invocation_id"),
+        )
+
+
+@dataclass
+class Session:
+    """Session - container for conversations.
+    
+    A Session represents a conversation thread that can contain multiple
+    invocations (turns). Sessions can be nested (parent_id) for sub-agent scenarios.
+    """
+    id: str = field(default_factory=lambda: generate_id("sess"))
+    root_agent_id: str = ""  # Root agent for this session
+    parent_id: str | None = None  # For forked sessions
+    created_at: datetime = field(default_factory=datetime.now)
+    updated_at: datetime = field(default_factory=datetime.now)
+    metadata: dict[str, Any] = field(default_factory=dict)  # title, etc.
+    is_active: bool = True
+    
+    # Control stack for DELEGATED mode
+    control_stack: list[ControlFrame] = field(default_factory=list)
+    
+    # Revert state (if currently in reverted state)
+    revert: dict[str, Any] | None = None
+    
+    @property
+    def active_agent_id(self) -> str:
+        """Get the currently active agent (top of control stack or root)."""
+        if self.control_stack:
+            return self.control_stack[-1].agent_id
+        return self.root_agent_id
+    
+    def push_control(self, frame: ControlFrame) -> None:
+        """Push control frame when delegating to sub-agent."""
+        self.control_stack.append(frame)
+        self.updated_at = datetime.now()
+    
+    def pop_control(self) -> ControlFrame | None:
+        """Pop control frame when sub-agent returns control."""
+        if self.control_stack:
+            self.updated_at = datetime.now()
+            return self.control_stack.pop()
+        return None
+    
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "id": self.id,
+            "root_agent_id": self.root_agent_id,
+            "parent_id": self.parent_id,
+            "created_at": self.created_at.isoformat(),
+            "updated_at": self.updated_at.isoformat(),
+            "metadata": self.metadata,
+            "is_active": self.is_active,
+            "control_stack": [f.to_dict() for f in self.control_stack],
+            "revert": self.revert,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "Session":
+        """Create from dictionary."""
+        return cls(
+            id=data["id"],
+            root_agent_id=data.get("root_agent_id", ""),
+            parent_id=data.get("parent_id"),
+            created_at=datetime.fromisoformat(data["created_at"]),
+            updated_at=datetime.fromisoformat(data["updated_at"]),
+            metadata=data.get("metadata", {}),
+            is_active=data.get("is_active", True),
+            control_stack=[ControlFrame.from_dict(f) for f in data.get("control_stack", [])],
+            revert=data.get("revert"),
+        )
+
+
+@dataclass
+class Invocation:
+    """A single invocation (turn) within a session.
+    
+    An Invocation represents one user input and the agent's response,
+    including all tool calls and intermediate steps.
+    """
+    id: str = field(default_factory=lambda: generate_id("inv"))
+    session_id: str = ""
+    agent_id: str = ""  # Which agent executes this invocation
+    mode: InvocationMode = InvocationMode.ROOT  # ROOT or DELEGATED
+    state: InvocationState = InvocationState.PENDING
+    
+    # Relationship (tree structure)
+    parent_invocation_id: str | None = None  # Parent invocation (for DELEGATED)
+    
+    # Timestamps
+    created_at: datetime = field(default_factory=datetime.now)
+    started_at: datetime | None = None
+    finished_at: datetime | None = None
+    
+    # Agent state for resumption
+    agent_state: dict[str, Any] | None = None
+    pending_tool_ids: list[str] = field(default_factory=list)
+    
+    # Execution info
+    step_count: int = 0
+    snapshot_id: str | None = None  # Pre-execution snapshot for revert
+    
+    # Error info
+    error: str | None = None
+    
+    # Branch (for state isolation)
+    branch: str | None = None
+    
+    metadata: dict[str, Any] = field(default_factory=dict)
+    
+    def mark_started(self) -> None:
+        """Mark invocation as started."""
+        self.state = InvocationState.RUNNING
+        self.started_at = datetime.now()
+    
+    def mark_completed(self) -> None:
+        """Mark invocation as completed."""
+        self.state = InvocationState.COMPLETED
+        self.finished_at = datetime.now()
+    
+    def mark_failed(self, error: str) -> None:
+        """Mark invocation as failed."""
+        self.state = InvocationState.FAILED
+        self.error = error
+        self.finished_at = datetime.now()
+    
+    def mark_cancelled(self) -> None:
+        """Mark invocation as cancelled."""
+        self.state = InvocationState.CANCELLED
+        self.finished_at = datetime.now()
+    
+    def mark_aborted(self) -> None:
+        """Mark invocation as aborted by user."""
+        self.state = InvocationState.ABORTED
+        self.finished_at = datetime.now()
+    
+    def mark_switched(self) -> None:
+        """Mark invocation as ended due to agent switch."""
+        self.state = InvocationState.SWITCHED
+        self.finished_at = datetime.now()
+    
+    def mark_paused(self) -> None:
+        """Mark invocation as paused."""
+        self.state = InvocationState.PAUSED
+    
+    def mark_suspended(self) -> None:
+        """Mark invocation as suspended (HITL)."""
+        self.state = InvocationState.SUSPENDED
+    
+    @property
+    def duration_ms(self) -> int | None:
+        """Get execution duration in milliseconds."""
+        if self.started_at and self.finished_at:
+            return int((self.finished_at - self.started_at).total_seconds() * 1000)
+        return None
+    
+    def to_dict(self) -> dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {
+            "id": self.id,
+            "session_id": self.session_id,
+            "agent_id": self.agent_id,
+            "mode": self.mode.value,
+            "state": self.state.value,
+            "parent_invocation_id": self.parent_invocation_id,
+            "created_at": self.created_at.isoformat(),
+            "started_at": self.started_at.isoformat() if self.started_at else None,
+            "finished_at": self.finished_at.isoformat() if self.finished_at else None,
+            "agent_state": self.agent_state,
+            "pending_tool_ids": self.pending_tool_ids,
+            "step_count": self.step_count,
+            "snapshot_id": self.snapshot_id,
+            "error": self.error,
+            "branch": self.branch,
+            "metadata": self.metadata,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "Invocation":
+        """Create from dictionary."""
+        return cls(
+            id=data["id"],
+            session_id=data["session_id"],
+            agent_id=data.get("agent_id", ""),
+            mode=InvocationMode(data.get("mode", "root")),
+            state=InvocationState(data["state"]),
+            parent_invocation_id=data.get("parent_invocation_id"),
+            created_at=datetime.fromisoformat(data["created_at"]),
+            started_at=datetime.fromisoformat(data["started_at"]) if data.get("started_at") else None,
+            finished_at=datetime.fromisoformat(data["finished_at"]) if data.get("finished_at") else None,
+            agent_state=data.get("agent_state"),
+            pending_tool_ids=data.get("pending_tool_ids", []),
+            step_count=data.get("step_count", 0),
+            snapshot_id=data.get("snapshot_id"),
+            error=data.get("error"),
+            branch=data.get("branch"),
+            metadata=data.get("metadata", {}),
+        )

aury/agents/core/types/subagent.py

@@ -0,0 +1,154 @@
+"""SubAgent input/output types for agent delegation."""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any, Literal, TypedDict, NotRequired, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .message import Message
+
+
+class SubAgentMode(Enum):
+    """SubAgent execution mode."""
+    EMBEDDED = "embedded"    # Embedded execution, shares parent's invocation
+    DELEGATED = "delegated"  # Delegated execution, creates new invocation
+
+
+class SubAgentInput(TypedDict):
+    """Input for SubAgent invocation.
+    
+    LLM 传的输入:
+    - agent: 调用哪个 agent
+    - task_context: 任务上下文（用户意图、背景、要求等）
+    - artifact_refs: 相关资料引用
+    
+    其他配置（mode, inherit_messages, summary_mode 等）
+    在 AgentConfig 中定义，不由 LLM 传入。
+    """
+    # Required
+    agent: str  # Agent key
+    
+    # Task context - 尽可能描述用户意图、背景信息、具体要求
+    task_context: NotRequired[str]
+    
+    # Artifact references - 相关资料 [{id, summary}, ...]
+    artifact_refs: NotRequired[list[dict[str, str]]]
+
+
+@dataclass
+class SubAgentMetadata:
+    """Metadata about sub-agent execution."""
+    child_invocation_id: str
+    agent_name: str
+    agent_type: str  # "react" | "workflow"
+    steps: int = 0
+    duration_ms: int = 0
+    token_usage: dict[str, int] = field(default_factory=dict)
+    
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "child_invocation_id": self.child_invocation_id,
+            "agent_name": self.agent_name,
+            "agent_type": self.agent_type,
+            "steps": self.steps,
+            "duration_ms": self.duration_ms,
+            "token_usage": self.token_usage,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "SubAgentMetadata":
+        return cls(
+            child_invocation_id=data["child_invocation_id"],
+            agent_name=data["agent_name"],
+            agent_type=data.get("agent_type", "react"),
+            steps=data.get("steps", 0),
+            duration_ms=data.get("duration_ms", 0),
+            token_usage=data.get("token_usage", {}),
+        )
+
+
+@dataclass
+class SubAgentResult:
+    """Result returned by sub-agent to parent agent.
+    
+    Contains both text output (for LLM context) and structured data.
+    """
+    # Text output (summary for LLM)
+    output: str
+    
+    # Execution status
+    status: Literal["completed", "aborted", "failed", "switched"]
+    
+    # Structured data (optional)
+    data: dict[str, Any] | None = None
+    
+    # State changes to merge back to parent
+    state_updates: dict[str, Any] | None = None
+    
+    # Error info (when failed)
+    error: str | None = None
+    
+    # Execution metadata
+    metadata: SubAgentMetadata | None = None
+    
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "output": self.output,
+            "status": self.status,
+            "data": self.data,
+            "state_updates": self.state_updates,
+            "error": self.error,
+            "metadata": self.metadata.to_dict() if self.metadata else None,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "SubAgentResult":
+        return cls(
+            output=data["output"],
+            status=data["status"],
+            data=data.get("data"),
+            state_updates=data.get("state_updates"),
+            error=data.get("error"),
+            metadata=SubAgentMetadata.from_dict(data["metadata"]) if data.get("metadata") else None,
+        )
+    
+    @classmethod
+    def completed(
+        cls,
+        output: str,
+        data: dict[str, Any] | None = None,
+        state_updates: dict[str, Any] | None = None,
+        metadata: SubAgentMetadata | None = None,
+    ) -> "SubAgentResult":
+        """Create a completed result."""
+        return cls(
+            output=output,
+            status="completed",
+            data=data,
+            state_updates=state_updates,
+            metadata=metadata,
+        )
+    
+    @classmethod
+    def failed(cls, error: str, output: str = "") -> "SubAgentResult":
+        """Create a failed result."""
+        return cls(
+            output=output or f"SubAgent failed: {error}",
+            status="failed",
+            error=error,
+        )
+    
+    @classmethod
+    def aborted(cls, output: str = "SubAgent was aborted") -> "SubAgentResult":
+        """Create an aborted result."""
+        return cls(output=output, status="aborted")
+
+
+__all__ = [
+    "SubAgentMode",
+    "SubAgentInput",
+    "SubAgentMetadata",
+    "SubAgentResult",
+]

aury/agents/core/types/tool.py

@@ -0,0 +1,205 @@
+"""Tool-related type definitions."""
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any, Awaitable, Callable
+
+from .session import generate_id
+
+
+@dataclass
+class ToolInfo:
+    """Tool metadata for LLM."""
+    name: str
+    description: str
+    parameters: dict[str, Any]  # JSON Schema
+    
+    def to_llm_schema(self) -> dict[str, Any]:
+        """Convert to LLM API format."""
+        return {
+            "name": self.name,
+            "description": self.description,
+            "input_schema": self.parameters,
+        }
+
+
+@dataclass
+class ToolContext:
+    """Context passed to tool execution."""
+    session_id: str
+    invocation_id: str
+    block_id: str
+    call_id: str
+    agent: str
+    abort_signal: asyncio.Event
+    update_metadata: Callable[[dict[str, Any]], Awaitable[None]]
+    
+    # Optional usage tracker
+    usage: Any | None = None  # UsageTracker
+    
+    # Branch for sub-agent isolation
+    branch: str | None = None
+    
+    # Caller's middleware chain (for sub-agent inheritance)
+    middleware: Any | None = None  # MiddlewareChain
+    
+    async def emit(self, block: Any) -> None:
+        """Emit a block event.
+        
+        Uses the global emit function via ContextVar.
+        Works automatically when called within agent.run() context.
+        
+        Args:
+            block: BlockEvent to emit
+        """
+        from ..context import emit as global_emit
+        
+        # Fill in IDs if not set
+        if hasattr(block, 'session_id') and not block.session_id:
+            block.session_id = self.session_id
+        if hasattr(block, 'invocation_id') and not block.invocation_id:
+            block.invocation_id = self.invocation_id
+        
+        await global_emit(block)
+    
+    async def emit_hitl(self, request_id: str, data: dict[str, Any]) -> None:
+        """Emit a HITL request block.
+        
+        Convenience method for tools that need user interaction.
+        The data format is flexible - can be anything the frontend understands:
+        - Choice selection (choices, radio, checkbox)
+        - Text input (text, textarea, number)
+        - Confirmation (yes/no, approve/reject)
+        - Rich content (product cards, file selection, etc.)
+        
+        Args:
+            request_id: Unique ID for this HITL request
+            data: Arbitrary data dict for frontend to render.
+                  Common fields: type, question, choices, default, context
+        """
+        from .block import BlockEvent, BlockKind
+        
+        await self.emit(BlockEvent(
+            kind=BlockKind.HITL_REQUEST,
+            data={"request_id": request_id, **data},
+        ))
+
+
+@dataclass
+class ToolResult:
+    """Tool execution result for LLM."""
+    output: str
+    is_error: bool = False
+    
+    @classmethod
+    def success(cls, output: str) -> ToolResult:
+        """Create a successful result."""
+        return cls(output=output, is_error=False)
+    
+    @classmethod
+    def error(cls, message: str) -> ToolResult:
+        """Create an error result."""
+        return cls(output=message, is_error=True)
+
+
+class ToolInvocationState(Enum):
+    """Tool invocation state machine."""
+    PARTIAL_CALL = "partial-call"  # Arguments streaming
+    CALL = "call"  # Arguments complete, ready to execute
+    RESULT = "result"  # Execution complete
+
+
+@dataclass
+class ToolInvocation:
+    """Tool invocation tracking (state machine)."""
+    tool_call_id: str
+    tool_name: str
+    state: ToolInvocationState = ToolInvocationState.PARTIAL_CALL
+    args: dict[str, Any] = field(default_factory=dict)
+    args_raw: str = ""  # Raw JSON string for streaming
+    result: str | None = None
+    is_error: bool = False
+    
+    # Timing
+    time: dict[str, datetime | None] = field(
+        default_factory=lambda: {"start": None, "end": None}
+    )
+    metadata: dict[str, Any] = field(default_factory=dict)
+    
+    def mark_call_complete(self) -> None:
+        """Mark arguments as complete."""
+        self.state = ToolInvocationState.CALL
+        self.time["start"] = datetime.now()
+    
+    def mark_result(self, result: str, is_error: bool = False) -> None:
+        """Mark execution complete."""
+        self.state = ToolInvocationState.RESULT
+        self.result = result
+        self.is_error = is_error
+        self.time["end"] = datetime.now()
+    
+    @property
+    def duration_ms(self) -> int | None:
+        """Get execution duration."""
+        if self.time["start"] and self.time["end"]:
+            return int((self.time["end"] - self.time["start"]).total_seconds() * 1000)
+        return None
+
+
+class BaseTool:
+    """Base class for tools with common functionality."""
+    
+    _name: str = "base_tool"
+    _description: str = "Base tool"
+    _parameters: dict[str, Any] = {
+        "type": "object",
+        "properties": {},
+        "required": [],
+    }
+    _config: ToolConfig | None = None
+    
+    @property
+    def name(self) -> str:
+        return self._name
+    
+    @property
+    def description(self) -> str:
+        return self._description
+    
+    @property
+    def parameters(self) -> dict[str, Any]:
+        return self._parameters
+    
+    @property
+    def config(self) -> ToolConfig:
+        """Get tool config. Returns default config if not set."""
+        return self._config or ToolConfig()
+    
+    async def execute(self, params: dict[str, Any], ctx: ToolContext) -> ToolResult:
+        """Override this method."""
+        raise NotImplementedError("Subclass must implement execute()")
+    
+    def get_info(self) -> ToolInfo:
+        """Get tool info."""
+        return ToolInfo(
+            name=self.name,
+            description=self.description,
+            parameters=self.parameters,
+        )
+
+
+@dataclass
+class ToolConfig:
+    """Tool configuration."""
+    is_resumable: bool = False  # Supports pause/resume
+    timeout: float | None = None  # Execution timeout in seconds
+    requires_permission: bool = False  # Needs HITL approval
+    permission_message: str | None = None
+    
+    # Retry configuration
+    max_retries: int = 0  # 0 = no retry
+    retry_delay: float = 1.0  # Base delay between retries (seconds)
+    retry_backoff: float = 2.0  # Exponential backoff multiplier