aury-agent 0.0.4__py3-none-any.whl

aury/agents/core/isolator.py

@@ -0,0 +1,97 @@
+"""State isolation for SubAgent execution.
+
+Provides ChainMap-based state isolation for EMBEDDED mode SubAgents.
+"""
+from __future__ import annotations
+
+from abc import abstractmethod
+from collections import ChainMap
+from typing import Any, Protocol
+
+
+class StateIsolator(Protocol):
+    """Protocol for state isolation.
+    
+    Used to isolate state changes in EMBEDDED SubAgent mode.
+    Changes in child scope don't affect parent until merged.
+    """
+    
+    @abstractmethod
+    def create_branch(self) -> "StateIsolator":
+        """Create a child branch with isolated state."""
+        ...
+    
+    @abstractmethod
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get value from state."""
+        ...
+    
+    @abstractmethod
+    def set(self, key: str, value: Any) -> None:
+        """Set value in current scope (doesn't affect parent)."""
+        ...
+    
+    @abstractmethod
+    def get_local_changes(self) -> dict[str, Any]:
+        """Get changes made in current scope only."""
+        ...
+    
+    @abstractmethod
+    def merge_to_parent(self) -> None:
+        """Merge local changes to parent scope."""
+        ...
+
+
+class ChainMapIsolator:
+    """ChainMap-based state isolator.
+    
+    Uses ChainMap to provide copy-on-write semantics.
+    Child writes go to a new dict, reads fall through to parent.
+    """
+    
+    def __init__(self, parent: "ChainMapIsolator | None" = None):
+        self._parent = parent
+        if parent is None:
+            self._chain = ChainMap({})
+        else:
+            # New child map on top of parent's chain
+            self._chain = parent._chain.new_child()
+    
+    def create_branch(self) -> "ChainMapIsolator":
+        """Create a child branch."""
+        return ChainMapIsolator(parent=self)
+    
+    def get(self, key: str, default: Any = None) -> Any:
+        """Get value, falling through to parent if not in local."""
+        return self._chain.get(key, default)
+    
+    def set(self, key: str, value: Any) -> None:
+        """Set value in local scope."""
+        self._chain.maps[0][key] = value
+    
+    def get_local_changes(self) -> dict[str, Any]:
+        """Get only the changes made in this scope."""
+        return dict(self._chain.maps[0])
+    
+    def merge_to_parent(self) -> None:
+        """Merge local changes to parent."""
+        if self._parent is not None:
+            local = self.get_local_changes()
+            for k, v in local.items():
+                self._parent.set(k, v)
+    
+    def __getitem__(self, key: str) -> Any:
+        return self._chain[key]
+    
+    def __setitem__(self, key: str, value: Any) -> None:
+        self.set(key, value)
+    
+    def __contains__(self, key: str) -> bool:
+        return key in self._chain
+    
+    def to_dict(self) -> dict[str, Any]:
+        """Flatten to dict (for serialization)."""
+        return dict(self._chain)
+
+
+__all__ = ["StateIsolator", "ChainMapIsolator"]

aury/agents/core/logging.py

@@ -0,0 +1,95 @@
+"""Centralized logging for aury-agent.
+
+Uses standard logging. In applications using foundation-kit,
+loguru will intercept these logs and add trace_id automatically.
+
+Usage:
+    from aury.agents.core.logging import get_logger
+    
+    logger = get_logger("react")  # -> "aury.agents.react"
+    logger.info("Starting agent", extra={"session_id": "..."})
+
+Pre-defined loggers:
+    - root: aury.agents (framework root)
+    - react: aury.agents.react
+    - workflow: aury.agents.workflow
+    - memory: aury.agents.memory
+    - tool: aury.agents.tool
+    - middleware: aury.agents.middleware
+    - bus: aury.agents.bus
+    - storage: aury.agents.storage
+    - session: aury.agents.session
+"""
+import logging
+from typing import Literal
+
+# Root namespace
+ROOT_NAMESPACE = "aury.agents"
+
+# Pre-defined sub-loggers
+_LOGGERS: dict[str, logging.Logger] = {}
+
+
+def get_logger(name: str | None = None) -> logging.Logger:
+    """Get a logger for the given component.
+    
+    Args:
+        name: Component name (e.g., "react", "workflow").
+              If None, returns root logger.
+              
+    Returns:
+        Logger instance for aury.agents.{name}
+    """
+    if name is None:
+        full_name = ROOT_NAMESPACE
+    else:
+        full_name = f"{ROOT_NAMESPACE}.{name}"
+    
+    if full_name not in _LOGGERS:
+        _LOGGERS[full_name] = logging.getLogger(full_name)
+    
+    return _LOGGERS[full_name]
+
+
+# Pre-instantiated loggers for common components
+logger = get_logger()  # Root logger
+react_logger = get_logger("react")
+workflow_logger = get_logger("workflow")
+memory_logger = get_logger("memory")
+tool_logger = get_logger("tool")
+middleware_logger = get_logger("middleware")
+bus_logger = get_logger("bus")
+storage_logger = get_logger("storage")
+session_logger = get_logger("session")
+context_logger = get_logger("context")
+
+
+# Convenience type for IDE completion
+LoggerName = Literal[
+    "react",
+    "workflow", 
+    "memory",
+    "tool",
+    "middleware",
+    "bus",
+    "storage",
+    "session",
+    "context",
+]
+
+
+__all__ = [
+    "get_logger",
+    "logger",
+    "react_logger",
+    "workflow_logger",
+    "memory_logger",
+    "tool_logger",
+    "middleware_logger",
+    "bus_logger",
+    "storage_logger",
+    "session_logger",
+    "context_logger",
+    "LoggerName",
+    "ROOT_NAMESPACE",
+]

aury/agents/core/parallel.py

@@ -0,0 +1,194 @@
+"""Parallel execution utilities for SubAgent.
+
+Provides utilities for parallel SubAgent execution, inspired by Google ADK.
+Supports merging async generators from multiple agents.
+"""
+from __future__ import annotations
+
+import asyncio
+from typing import Any, AsyncGenerator, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .types.block import BlockEvent
+
+
+async def merge_agent_runs(
+    agent_runs: list[AsyncGenerator["BlockEvent", None]],
+) -> AsyncGenerator["BlockEvent", None]:
+    """Merge multiple agent runs into a single stream.
+    
+    Executes agents in parallel, yielding events as they arrive.
+    Each agent's events are processed sequentially to maintain order.
+    
+    Args:
+        agent_runs: List of async generators yielding BlockEvents
+        
+    Yields:
+        BlockEvent from any of the running agents
+        
+    Example:
+        async def run_agent(agent, input):
+            async for event in agent.run(input):
+                yield event
+        
+        runs = [run_agent(a, "task") for a in agents]
+        async for event in merge_agent_runs(runs):
+            print(event)
+    """
+    if not agent_runs:
+        return
+    
+    sentinel = object()
+    queue: asyncio.Queue[tuple[Any, asyncio.Event | None]] = asyncio.Queue()
+    
+    async def process_agent(events: AsyncGenerator["BlockEvent", None]) -> None:
+        """Process single agent's events."""
+        try:
+            async for event in events:
+                # Create resume signal to wait for consumer
+                resume_signal = asyncio.Event()
+                await queue.put((event, resume_signal))
+                # Wait for upstream to consume before generating more
+                await resume_signal.wait()
+        finally:
+            # Mark this agent as finished
+            await queue.put((sentinel, None))
+    
+    # Use TaskGroup for parallel execution (Python 3.11+)
+    async with asyncio.TaskGroup() as tg:
+        for events in agent_runs:
+            tg.create_task(process_agent(events))
+        
+        sentinel_count = 0
+        # Run until all agents finished
+        while sentinel_count < len(agent_runs):
+            item, resume_signal = await queue.get()
+            
+            if item is sentinel:
+                sentinel_count += 1
+            else:
+                yield item
+                # Signal agent to continue
+                if resume_signal:
+                    resume_signal.set()
+
+
+async def run_agents_parallel(
+    agent_tasks: list[tuple[str, Any]],  # [(agent_name, input), ...]
+    agent_runner: Any,  # Callable to run agent: (name, input) -> AsyncGenerator
+    timeout: float | None = None,
+) -> dict[str, Any]:
+    """Run multiple agents in parallel and collect results.
+    
+    Unlike merge_agent_runs which streams events, this collects final results.
+    
+    Args:
+        agent_tasks: List of (agent_name, input) tuples
+        agent_runner: Async function to run agent
+        timeout: Optional timeout in seconds
+        
+    Returns:
+        Dict mapping agent_name to result or error
+        
+    Example:
+        async def run_agent(name, input):
+            agent = get_agent(name)
+            result = []
+            async for event in agent.run(input):
+                if event.kind == BlockKind.TEXT:
+                    result.append(event.data.get("content", ""))
+            return "".join(result)
+        
+        results = await run_agents_parallel(
+            [("researcher", "find data"), ("analyzer", "analyze")],
+            run_agent,
+            timeout=60.0,
+        )
+    """
+    results: dict[str, Any] = {}
+    
+    async def run_one(name: str, input: Any) -> tuple[str, Any]:
+        try:
+            result = await agent_runner(name, input)
+            return (name, result)
+        except Exception as e:
+            return (name, {"error": str(e)})
+    
+    tasks = [run_one(name, input) for name, input in agent_tasks]
+    
+    if timeout:
+        try:
+            completed = await asyncio.wait_for(
+                asyncio.gather(*tasks, return_exceptions=True),
+                timeout=timeout,
+            )
+        except asyncio.TimeoutError:
+            # Return partial results with timeout error for incomplete
+            for name, _ in agent_tasks:
+                if name not in results:
+                    results[name] = {"error": "timeout"}
+            return results
+    else:
+        completed = await asyncio.gather(*tasks, return_exceptions=True)
+    
+    for item in completed:
+        if isinstance(item, tuple):
+            name, result = item
+            results[name] = result
+        elif isinstance(item, Exception):
+            # Exception from gather
+            pass
+    
+    return results
+
+
+class ParallelSubAgentContext:
+    """Context for tracking parallel SubAgent execution.
+    
+    Manages branch isolation and result collection for parallel execution.
+    """
+    
+    def __init__(
+        self,
+        parent_invocation_id: str,
+        session_id: str,
+    ):
+        self.parent_invocation_id = parent_invocation_id
+        self.session_id = session_id
+        self.branches: dict[str, str] = {}  # agent_name -> branch_id
+        self.results: dict[str, Any] = {}
+        self.errors: dict[str, str] = {}
+        self._completed: set[str] = set()
+    
+    def create_branch(self, agent_name: str) -> str:
+        """Create isolated branch for sub-agent."""
+        branch_id = f"{self.parent_invocation_id}.{agent_name}"
+        self.branches[agent_name] = branch_id
+        return branch_id
+    
+    def mark_completed(self, agent_name: str, result: Any) -> None:
+        """Mark agent as completed with result."""
+        self._completed.add(agent_name)
+        self.results[agent_name] = result
+    
+    def mark_failed(self, agent_name: str, error: str) -> None:
+        """Mark agent as failed with error."""
+        self._completed.add(agent_name)
+        self.errors[agent_name] = error
+    
+    @property
+    def all_completed(self) -> bool:
+        """Check if all agents completed."""
+        return len(self._completed) == len(self.branches)
+    
+    @property
+    def pending_agents(self) -> list[str]:
+        """Get list of pending agent names."""
+        return [n for n in self.branches if n not in self._completed]
+
+
+__all__ = [
+    "merge_agent_runs",
+    "run_agents_parallel",
+    "ParallelSubAgentContext",
+]

aury/agents/core/runner.py

@@ -0,0 +1,139 @@
+"""Runner - Entry point coordinator for agent execution.
+
+The Runner handles:
+- Creating InvocationContext from Session/Invocation
+- Managing the execution lifecycle
+- Coordinating services (Session, Message, Storage)
+"""
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any, AsyncIterator, TYPE_CHECKING
+
+from .context import InvocationContext
+from .logging import logger
+
+if TYPE_CHECKING:
+    from .base import BaseAgent
+    from .types.block import BlockEvent
+    from .services.session import SessionService
+    from .services.message import MessageService
+    from ..backends import Backends
+    from .event_bus import Bus
+    from ..types import PromptInput
+
+
+class Runner:
+    """Entry point for agent execution.
+    
+    Coordinates the execution lifecycle:
+    1. Get or create session
+    2. Create invocation
+    3. Build InvocationContext
+    4. Execute agent
+    5. Handle completion/errors
+    
+    Example:
+        runner = Runner(
+            session_service=session_service,
+            message_service=message_service,
+            backends=backends,
+            bus=bus,
+        )
+        
+        async for response in runner.run(
+            agent_class=MyAgent,
+            session_id="sess_123",
+            input=PromptInput(text="Hello"),
+        ):
+            print(response)
+    """
+    
+    def __init__(
+        self,
+        session_service: "SessionService",
+        message_service: "MessageService",
+        backends: "Backends",
+        bus: "Bus",
+    ):
+        self.session_service = session_service
+        self.message_service = message_service
+        self.backends = backends
+        self.bus = bus
+    
+    async def run(
+        self,
+        agent_class: type["BaseAgent"],
+        session_id: str | None = None,
+        input: "PromptInput | None" = None,
+        **agent_kwargs: Any,
+    ) -> AsyncIterator["BlockEvent"]:
+        """Run an agent.
+        
+        Args:
+            agent_class: Agent class to instantiate
+            session_id: Existing session ID (creates new if None)
+            input: Input to agent
+            **agent_kwargs: Additional args for agent constructor
+            
+        Yields:
+            BlockEvent streaming events
+        """
+        from ..types import generate_id, Invocation, InvocationState
+        
+        # Get or create session
+        if session_id:
+            session = await self.session_service.get(session_id)
+            if session is None:
+                raise ValueError(f"Session not found: {session_id}")
+        else:
+            agent_name = getattr(agent_class, 'name', agent_class.__name__)
+            session = await self.session_service.create(root_agent_id=agent_name)
+        
+        logger.info(f"Starting agent execution: session={session.id}")
+        
+        # Create invocation
+        invocation = Invocation(
+            id=generate_id("inv"),
+            session_id=session.id,
+            agent_id=getattr(agent_class, 'name', agent_class.__name__),
+            state=InvocationState.RUNNING,
+            started_at=datetime.now(),
+        )
+        
+        # Build context
+        ctx = InvocationContext(
+            session_id=session.id,
+            invocation_id=invocation.id,
+            agent_id=invocation.agent_id,
+            backends=self.backends,
+            bus=self.bus,
+        )
+        
+        # Create agent instance
+        agent = agent_class(
+            ctx=ctx,
+            **agent_kwargs,
+        )
+        
+        try:
+            # Execute
+            async for response in agent.run(input):
+                yield response
+            
+            invocation.state = InvocationState.COMPLETED
+            logger.info(f"Agent completed: invocation={invocation.id}")
+            
+        except Exception as e:
+            invocation.state = InvocationState.FAILED
+            logger.error(f"Agent failed: invocation={invocation.id}, error={e}")
+            raise
+        
+        finally:
+            invocation.finished_at = datetime.now()
+            # Persist invocation
+            if self.backends.state:
+                await self.backends.state.set("invocation", invocation.id, invocation.__dict__)
+
+
+__all__ = ["Runner"]

aury/agents/core/services/__init__.py

@@ -0,0 +1,5 @@
+"""Service protocols and implementations."""
+from .session import SessionService
+from .message import MessageService
+
+__all__ = ["SessionService", "MessageService"]

aury/agents/core/services/file_session.py

@@ -0,0 +1,144 @@
+"""File-based session service implementation."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import Any
+
+from ..logging import session_logger as logger
+from ..types.session import Session, ControlFrame
+
+
+class FileSessionService:
+    """File-based session storage.
+    
+    Stores sessions as JSON files in a directory.
+    Suitable for development and single-instance deployments.
+    
+    Directory structure:
+        {base_dir}/
+            {session_id}.json
+    """
+    
+    def __init__(self, base_dir: str | Path):
+        """Initialize file session service.
+        
+        Args:
+            base_dir: Directory to store session files
+        """
+        self.base_dir = Path(base_dir)
+        self.base_dir.mkdir(parents=True, exist_ok=True)
+    
+    def _session_path(self, session_id: str) -> Path:
+        """Get file path for session."""
+        # Sanitize ID to prevent path traversal
+        safe_id = "".join(c for c in session_id if c.isalnum() or c in "-_")
+        return self.base_dir / f"{safe_id}.json"
+    
+    async def create(self, root_agent_id: str, **kwargs) -> Session:
+        """Create a new session."""
+        session = Session(root_agent_id=root_agent_id, **kwargs)
+        await self._save(session)
+        logger.debug(f"Created session: {session.id}")
+        return session
+    
+    async def get(self, session_id: str) -> Session | None:
+        """Get session by ID."""
+        path = self._session_path(session_id)
+        if not path.exists():
+            return None
+        
+        try:
+            data = json.loads(path.read_text(encoding="utf-8"))
+            return Session.from_dict(data)
+        except Exception as e:
+            logger.error(f"Failed to load session {session_id}: {e}")
+            return None
+    
+    async def update(self, session: Session) -> None:
+        """Update session."""
+        await self._save(session)
+        logger.debug(f"Updated session: {session.id}")
+    
+    async def delete(self, session_id: str) -> bool:
+        """Delete session."""
+        path = self._session_path(session_id)
+        if path.exists():
+            try:
+                path.unlink()
+                logger.debug(f"Deleted session: {session_id}")
+                return True
+            except Exception as e:
+                logger.error(f"Failed to delete session {session_id}: {e}")
+        return False
+    
+    async def list(self, limit: int = 100, offset: int = 0) -> list[Session]:
+        """List sessions, ordered by updated_at descending."""
+        sessions = []
+        
+        for path in self.base_dir.glob("*.json"):
+            try:
+                data = json.loads(path.read_text(encoding="utf-8"))
+                sessions.append(Session.from_dict(data))
+            except Exception as e:
+                logger.warning(f"Failed to load session file {path}: {e}")
+                continue
+        
+        # Sort by updated_at descending
+        sessions.sort(key=lambda s: s.updated_at, reverse=True)
+        
+        # Apply pagination
+        return sessions[offset : offset + limit]
+    
+    async def push_control(self, session_id: str, frame: ControlFrame) -> None:
+        """Push control frame to session's control stack."""
+        session = await self.get(session_id)
+        if session is None:
+            raise ValueError(f"Session not found: {session_id}")
+        
+        session.push_control(frame)
+        await self._save(session)
+        logger.debug(f"Pushed control frame to session {session_id}: {frame.agent_id}")
+    
+    async def pop_control(self, session_id: str) -> ControlFrame | None:
+        """Pop control frame from session's control stack."""
+        session = await self.get(session_id)
+        if session is None:
+            raise ValueError(f"Session not found: {session_id}")
+        
+        frame = session.pop_control()
+        if frame:
+            await self._save(session)
+            logger.debug(f"Popped control frame from session {session_id}: {frame.agent_id}")
+        return frame
+    
+    async def _save(self, session: Session) -> None:
+        """Save session to file."""
+        path = self._session_path(session.id)
+        try:
+            path.write_text(
+                json.dumps(session.to_dict(), indent=2, ensure_ascii=False),
+                encoding="utf-8"
+            )
+        except Exception as e:
+            logger.error(f"Failed to save session {session.id}: {e}")
+            raise
+    
+    async def get_by_root_agent(self, root_agent_id: str, limit: int = 100) -> list[Session]:
+        """Get sessions by root agent ID."""
+        all_sessions = await self.list(limit=10000)  # Get all
+        return [
+            s for s in all_sessions
+            if s.root_agent_id == root_agent_id
+        ][:limit]
+    
+    async def get_active(self, limit: int = 100) -> list[Session]:
+        """Get active sessions."""
+        all_sessions = await self.list(limit=10000)
+        return [
+            s for s in all_sessions
+            if s.is_active
+        ][:limit]
+
+
+__all__ = ["FileSessionService"]