aury-agent 0.0.4__py3-none-any.whl

aury/agents/core/context_builder.py

@@ -0,0 +1,303 @@
+"""Context builder for LLM communication.
+
+ContextBuilder constructs the full context to send to LLM, combining:
+- System prompt
+- Knowledge (user-defined, via Middleware)
+- Summary (compressed history)
+- Recalls (session key points)
+- Recent messages
+- Current input
+
+Context is runtime-built, not stored.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Protocol, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from .types.message import Message, PromptInput
+    from .types.recall import Recall, Summary
+    from ..backends.state import StateBackend
+
+
+@dataclass
+class ContextConfig:
+    """Configuration for context building."""
+
+    # Token limits
+    max_tokens: int = 8000
+
+    # Message control
+    max_recent_messages: int = 50
+
+    # Memory control
+    include_summary: bool = True
+    include_recalls: bool = True
+    recall_limit: int = 20
+
+    # Knowledge control (user-defined via Middleware)
+    include_knowledge: bool = True
+    knowledge_limit: int = 10
+
+    # Compression
+    enable_compression: bool = True
+    compression_threshold: int = 6000
+
+
+@dataclass
+class LLMContext:
+    """Context prepared for LLM.
+
+    This is the final object sent to LLM, not stored.
+    """
+
+    # Final messages (LLM API format)
+    messages: list[dict[str, Any]] = field(default_factory=list)
+
+    # Sources (for debugging/tracing)
+    source_summary: "Summary | None" = None
+    source_recalls: list["Recall"] = field(default_factory=list)
+    source_knowledge: list[dict[str, Any]] = field(default_factory=list)  # User-defined structure
+    source_messages: list["Message"] = field(default_factory=list)
+
+    # Token stats
+    estimated_tokens: int = 0
+
+    def to_llm_messages(self) -> list[dict[str, Any]]:
+        """Get messages in LLM API format."""
+        return self.messages
+
+
+class ContextBuilder(Protocol):
+    """Protocol for context building.
+
+    Users can implement custom builders for specialized context management.
+    """
+
+    async def build(
+        self,
+        session_id: str,
+        invocation_id: str,
+        current_input: "PromptInput",
+        branch: str | None = None,
+        config: ContextConfig | None = None,
+    ) -> LLMContext:
+        """Build context for LLM.
+
+        Args:
+            session_id: Current session ID
+            invocation_id: Current invocation ID
+            current_input: User input for this turn
+            branch: SubAgent branch (for isolation)
+            config: Context configuration
+
+        Returns:
+            LLMContext ready for LLM API call
+        """
+        ...
+
+
+class DefaultContextBuilder:
+    """Default implementation of context builder.
+
+    Build order:
+    1. System prompt
+    2. Knowledge (injected via Middleware)
+    3. Summary (compressed history)
+    4. Recalls (key points)
+    5. Recent messages
+    6. Current input
+    """
+
+    def __init__(
+        self,
+        storage: "StateBackend",
+        system_prompt: str | None = None,
+    ):
+        """Initialize context builder.
+
+        Args:
+            storage: State backend for loading messages/recalls
+            system_prompt: System prompt template (can include {var} for state.vars)
+        """
+        self._storage = storage
+        self._system_prompt = system_prompt
+
+    async def build(
+        self,
+        session_id: str,
+        invocation_id: str,
+        current_input: "PromptInput",
+        branch: str | None = None,
+        config: ContextConfig | None = None,
+        *,
+        state_vars: dict[str, Any] | None = None,
+    ) -> LLMContext:
+        """Build context for LLM.
+
+        Args:
+            session_id: Current session ID
+            invocation_id: Current invocation ID
+            current_input: User input for this turn
+            branch: SubAgent branch (for isolation)
+            config: Context configuration
+            state_vars: Variables for prompt formatting (from state.vars)
+        """
+        from .types.message import Message, MessageRole
+        from .types.recall import Recall, Summary
+
+        cfg = config or ContextConfig()
+        context = LLMContext()
+        messages: list[dict[str, Any]] = []
+
+        # 1. System prompt
+        if self._system_prompt:
+            system_text = self._system_prompt
+            if state_vars:
+                try:
+                    system_text = system_text.format(**state_vars)
+                except KeyError:
+                    pass  # Ignore missing vars
+            messages.append(
+                {
+                    "role": "system",
+                    "content": system_text,
+                }
+            )
+
+        # 2. Summary (compressed history)
+        if cfg.include_summary:
+            summary = await self._load_summary(session_id)
+            if summary:
+                context.source_summary = summary
+                messages.append(
+                    {
+                        "role": "system",
+                        "content": f"[Previous Conversation Summary]\n{summary.content}",
+                    }
+                )
+
+        # 3. Recalls (key points)
+        if cfg.include_recalls:
+            recalls = await self._load_recalls(session_id, branch, cfg.recall_limit)
+            if recalls:
+                context.source_recalls = recalls
+                recalls_text = "\n".join([f"- {r.content}" for r in recalls])
+                messages.append(
+                    {
+                        "role": "system",
+                        "content": f"[Key Information]\n{recalls_text}",
+                    }
+                )
+
+        # 4. Recent messages
+        recent_messages = await self._load_messages(
+            session_id,
+            branch,
+            cfg.max_recent_messages,
+        )
+        context.source_messages = recent_messages
+
+        for msg in recent_messages:
+            llm_msg = msg.to_llm_format()
+            messages.append(llm_msg)
+
+        # 5. Current input
+        input_msg = current_input.to_message(session_id, invocation_id)
+        messages.append(input_msg.to_llm_format())
+
+        context.messages = messages
+        context.estimated_tokens = self._estimate_tokens(messages)
+
+        return context
+
+    async def _load_summary(self, session_id: str) -> "Summary | None":
+        """Load session summary."""
+        from .types.recall import Summary
+
+        data = await self._storage.get("summaries", session_id)
+        if data:
+            return Summary.from_dict(data)
+        return None
+
+    async def _load_recalls(
+        self,
+        session_id: str,
+        branch: str | None,
+        limit: int,
+    ) -> list["Recall"]:
+        """Load session recalls."""
+        from .types.recall import Recall
+
+        # Load all recalls for session
+        keys = await self._storage.list("recalls", prefix=session_id)
+        recalls: list[Recall] = []
+
+        for key in keys[: limit * 2]:  # Load extra to filter by branch
+            data = await self._storage.get("recalls", key)
+            if data:
+                recall = Recall.from_dict(data)
+                # Filter by branch
+                if branch is None or recall.branch is None or recall.branch == branch:
+                    recalls.append(recall)
+                    if len(recalls) >= limit:
+                        break
+
+        # Sort by importance
+        recalls.sort(key=lambda r: r.importance, reverse=True)
+        return recalls[:limit]
+
+    async def _load_messages(
+        self,
+        session_id: str,
+        branch: str | None,
+        limit: int,
+    ) -> list["Message"]:
+        """Load recent messages."""
+        from .types.message import Message
+
+        # Load all message keys for session
+        keys = await self._storage.list("messages", prefix=session_id)
+        messages: list[Message] = []
+
+        # Load in reverse (most recent first)
+        for key in reversed(keys):
+            if len(messages) >= limit:
+                break
+
+            data = await self._storage.get("messages", key)
+            if data:
+                msg = Message.from_dict(data)
+                # Filter by branch
+                if branch is None or msg.branch is None or msg.branch == branch:
+                    messages.append(msg)
+
+        # Reverse to chronological order
+        messages.reverse()
+        return messages
+
+    def _estimate_tokens(self, messages: list[dict[str, Any]]) -> int:
+        """Estimate token count (rough approximation)."""
+        total_chars = 0
+        for msg in messages:
+            content = msg.get("content", "")
+            if isinstance(content, str):
+                total_chars += len(content)
+            elif isinstance(content, list):
+                for part in content:
+                    if isinstance(part, dict):
+                        text = part.get("text", "") or part.get("content", "")
+                        total_chars += len(str(text))
+
+        # Rough estimate: 4 chars per token
+        return total_chars // 4
+
+
+__all__ = [
+    "ContextConfig",
+    "LLMContext",
+    "ContextBuilder",
+    "DefaultContextBuilder",
+]

aury/agents/core/event_bus/__init__.py

@@ -0,0 +1,15 @@
+"""Event bus for pub/sub messaging."""
+from .bus import (
+    Events,
+    EventHandler,
+    EventBus,
+    EventCollector,
+)
+
+
+__all__ = [
+    "Events",
+    "EventHandler",
+    "EventBus",
+    "EventCollector",
+]

aury/agents/core/event_bus/bus.py

@@ -0,0 +1,203 @@
+"""Event bus for pub/sub messaging."""
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Callable, Awaitable
+from enum import Enum
+
+from ..logging import bus_logger as logger
+
+
+class Events:
+    """Predefined event types."""
+    
+    # Session lifecycle
+    SESSION_CREATED = "session.created"
+    SESSION_UPDATED = "session.updated"
+    SESSION_ENDED = "session.ended"
+    
+    # Invocation lifecycle
+    INVOCATION_START = "invocation.start"
+    INVOCATION_END = "invocation.end"
+    INVOCATION_ERROR = "invocation.error"
+    INVOCATION_CANCELLED = "invocation.cancelled"
+    
+    # Unified block event (used by ctx.emit())
+    BLOCK = "block"  # Main streaming event for all block outputs
+    
+    # Block lifecycle events (for persistence/logging)
+    BLOCK_CREATED = "block.created"
+    BLOCK_UPDATED = "block.updated"
+    BLOCK_DELTA = "block.delta"
+    BLOCK_CLOSED = "block.closed"
+    
+    # Tool events
+    TOOL_START = "tool.start"
+    TOOL_END = "tool.end"
+    TOOL_ERROR = "tool.error"
+    
+    # Workflow node events
+    NODE_START = "node.start"
+    NODE_END = "node.end"
+    NODE_ERROR = "node.error"
+    NODE_SKIPPED = "node.skipped"
+    
+    # LLM events
+    LLM_START = "llm.start"
+    LLM_END = "llm.end"
+    LLM_STREAM = "llm.stream"
+    
+    # Memory events
+    MEMORY_ADD = "memory.add"
+    MEMORY_SEARCH = "memory.search"
+    
+    # Usage events
+    USAGE_RECORDED = "usage.recorded"
+    
+    # Permission events
+    PERMISSION_REQUESTED = "permission.requested"
+    PERMISSION_RESOLVED = "permission.resolved"
+    
+    # State events
+    STATE_CHANGED = "state.changed"
+    STATE_REVERTED = "state.reverted"
+
+
+EventHandler = Callable[[str, Any], Awaitable[None]] | Callable[[str, Any], None]
+
+
+class EventBus:
+    """Event bus with pub/sub support.
+    
+    Features:
+    - Async and sync handlers
+    - Wildcard subscription with "*"
+    - Handler errors don't block other handlers
+    """
+    
+    def __init__(self) -> None:
+        self._subscriptions: dict[str, list[EventHandler]] = {}
+        self._lock = asyncio.Lock()
+    
+    async def publish(self, event_type: str, payload: Any) -> None:
+        """Publish event to all subscribers.
+        
+        Args:
+            event_type: Event type string
+            payload: Event payload (any data)
+        """
+        handlers: list[EventHandler] = []
+        
+        async with self._lock:
+            # Exact match
+            handlers.extend(self._subscriptions.get(event_type, []))
+            # Wildcard match
+            handlers.extend(self._subscriptions.get("*", []))
+        
+        logger.debug("Publishing event", extra={"event_type": event_type})
+        
+        for handler in handlers:
+            try:
+                if asyncio.iscoroutinefunction(handler):
+                    await handler(event_type, payload)
+                else:
+                    handler(event_type, payload)
+            except Exception as e:
+                logger.error(
+                    "Event handler error",
+                    extra={"event_type": event_type, "error": str(e)},
+                    exc_info=True,
+                )
+    
+    def subscribe(self, event_type: str, handler: EventHandler) -> Callable[[], None]:
+        """Subscribe to event type.
+        
+        Args:
+            event_type: Event type to subscribe to, or "*" for all
+            handler: Callback function (async or sync)
+        
+        Returns:
+            Unsubscribe function
+        """
+        if event_type not in self._subscriptions:
+            self._subscriptions[event_type] = []
+        
+        self._subscriptions[event_type].append(handler)
+        
+        def unsubscribe() -> None:
+            if event_type in self._subscriptions:
+                try:
+                    self._subscriptions[event_type].remove(handler)
+                except ValueError:
+                    pass
+        
+        return unsubscribe
+    
+    def on(self, event_type: str) -> Callable[[EventHandler], EventHandler]:
+        """Decorator for subscribing to events.
+        
+        Usage:
+            @bus.on("session.created")
+            async def handler(event_type, payload):
+                ...
+        """
+        def decorator(handler: EventHandler) -> EventHandler:
+            self.subscribe(event_type, handler)
+            return handler
+        return decorator
+    
+    def clear(self) -> None:
+        """Clear all subscriptions."""
+        self._subscriptions.clear()
+    
+    def subscriber_count(self, event_type: str | None = None) -> int:
+        """Get number of subscribers.
+        
+        Args:
+            event_type: Specific event type, or None for total count
+        """
+        if event_type:
+            return len(self._subscriptions.get(event_type, []))
+        return sum(len(handlers) for handlers in self._subscriptions.values())
+
+
+class EventCollector:
+    """Collects events for testing/debugging.
+    
+    Usage:
+        collector = EventCollector(bus)
+        collector.start()
+        # ... do stuff ...
+        events = collector.stop()
+    """
+    
+    def __init__(self, bus: EventBus, event_types: list[str] | None = None):
+        self.bus = bus
+        self.event_types = event_types  # None means all
+        self.events: list[tuple[str, Any]] = []
+        self._unsubscribers: list[Callable[[], None]] = []
+    
+    def _handler(self, event_type: str, payload: Any) -> None:
+        if self.event_types is None or event_type in self.event_types:
+            self.events.append((event_type, payload))
+    
+    def start(self) -> None:
+        """Start collecting events."""
+        if self.event_types:
+            for event_type in self.event_types:
+                unsub = self.bus.subscribe(event_type, self._handler)
+                self._unsubscribers.append(unsub)
+        else:
+            unsub = self.bus.subscribe("*", self._handler)
+            self._unsubscribers.append(unsub)
+    
+    def stop(self) -> list[tuple[str, Any]]:
+        """Stop collecting and return events."""
+        for unsub in self._unsubscribers:
+            unsub()
+        self._unsubscribers.clear()
+        return self.events
+    
+    def clear(self) -> None:
+        """Clear collected events."""
+        self.events.clear()

aury/agents/core/factory.py

@@ -0,0 +1,169 @@
+"""Agent factory for creating agent instances.
+
+Provides unified creation of both ReactAgent and WorkflowAgent:
+    factory = AgentFactory()
+    factory.register("researcher", ResearcherAgent)  # ReactAgent subclass
+    factory.register("pipeline", PipelineWorkflow)   # WorkflowAgent subclass
+    
+    agent = factory.create("researcher", ctx)  # Works for both types
+"""
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+from .base import AgentConfig
+
+if TYPE_CHECKING:
+    from .base import BaseAgent
+    from .context import InvocationContext
+
+
+class AgentFactory:
+    """Factory for creating agent instances.
+    
+    Unified factory for both ReactAgent and WorkflowAgent.
+    All agents use the same constructor:
+        __init__(self, ctx: InvocationContext, config: AgentConfig | None = None)
+    
+    Usage:
+        factory = AgentFactory()
+        
+        # Register agent classes
+        factory.register("researcher", ResearcherAgent)
+        factory.register("coder", CoderAgent)
+        factory.register("pipeline", PipelineWorkflow)
+        
+        # Create agents
+        agent = factory.create("researcher", ctx)
+        agent = factory.create("researcher", ctx, config=custom_config)
+        
+        # Auto-register from class
+        factory.register_class(ResearcherAgent)  # Uses class.name
+    """
+    
+    def __init__(self):
+        self._registry: dict[str, type["BaseAgent"]] = {}
+    
+    def register(self, name: str, agent_class: type["BaseAgent"]) -> None:
+        """Register an agent class with a name.
+        
+        Args:
+            name: Name to register under
+            agent_class: Agent class (must have unified constructor)
+        """
+        self._registry[name] = agent_class
+    
+    def register_class(self, agent_class: type["BaseAgent"]) -> None:
+        """Register an agent class using its class-level name.
+        
+        Args:
+            agent_class: Agent class with 'name' class attribute
+        """
+        name = getattr(agent_class, 'name', agent_class.__name__)
+        self._registry[name] = agent_class
+    
+    def register_all(self, *agent_classes: type["BaseAgent"]) -> None:
+        """Register multiple agent classes.
+        
+        Args:
+            agent_classes: Agent classes with 'name' class attribute
+        """
+        for agent_class in agent_classes:
+            self.register_class(agent_class)
+    
+    def create(
+        self,
+        agent_type: str,
+        ctx: "InvocationContext",
+        config: AgentConfig | None = None,
+    ) -> "BaseAgent":
+        """Create an agent instance.
+        
+        All agents are created with the same signature:
+            agent = AgentClass(ctx, config)
+        
+        Args:
+            agent_type: Registered agent type name
+            ctx: InvocationContext with all services
+            config: Agent configuration (optional)
+            
+        Returns:
+            Agent instance (ReactAgent or WorkflowAgent)
+            
+        Raises:
+            KeyError: If agent type not registered
+        """
+        if agent_type not in self._registry:
+            available = ", ".join(self._registry.keys()) or "none"
+            raise KeyError(
+                f"Unknown agent type: {agent_type}. Available: {available}"
+            )
+        
+        agent_class = self._registry[agent_type]
+        return agent_class(ctx, config)
+    
+    def create_subagent(
+        self,
+        agent_type: str,
+        parent_ctx: "InvocationContext",
+        mode: str = "delegated",
+        config: AgentConfig | None = None,
+    ) -> "BaseAgent":
+        """Create a sub-agent with child context.
+        
+        Convenience method that creates child context and agent.
+        
+        Args:
+            agent_type: Registered agent type name
+            parent_ctx: Parent's InvocationContext
+            mode: Execution mode ('delegated' or 'embedded')
+            config: Agent configuration (optional)
+            
+        Returns:
+            Agent instance with child context
+        """
+        child_ctx = parent_ctx.create_child(agent_id=agent_type, mode=mode)
+        return self.create(agent_type, child_ctx, config)
+    
+    def list_types(self) -> list[str]:
+        """List registered agent types."""
+        return list(self._registry.keys())
+    
+    def get_class(self, agent_type: str) -> type["BaseAgent"] | None:
+        """Get agent class by type name."""
+        return self._registry.get(agent_type)
+    
+    def is_registered(self, agent_type: str) -> bool:
+        """Check if agent type is registered."""
+        return agent_type in self._registry
+    
+    def get_info(self, agent_type: str) -> dict[str, Any] | None:
+        """Get agent info (name, description, type).
+        
+        Returns:
+            Dict with name, description, agent_type, or None if not found
+        """
+        agent_class = self._registry.get(agent_type)
+        if agent_class is None:
+            return None
+        
+        return {
+            "name": getattr(agent_class, 'name', agent_type),
+            "description": getattr(agent_class, 'description', ''),
+            "agent_type": getattr(agent_class, 'agent_type', 'react'),
+            "sub_agents": [
+                getattr(sa, 'name', sa.__name__)
+                for sa in getattr(agent_class, 'sub_agents', [])
+            ],
+        }
+    
+    def list_info(self) -> list[dict[str, Any]]:
+        """List info for all registered agents."""
+        return [
+            self.get_info(name)
+            for name in self._registry.keys()
+            if self.get_info(name) is not None
+        ]
+
+
+__all__ = ["AgentFactory"]