aury-agent 0.0.4__py3-none-any.whl

aury/agents/middleware/chain.py

@@ -0,0 +1,342 @@
+"""Middleware chain for sequential processing."""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, TYPE_CHECKING
+
+from ..core.logging import middleware_logger as logger
+from .types import TriggerMode, HookAction, HookResult
+from .base import Middleware
+
+if TYPE_CHECKING:
+    from ..core.types.tool import BaseTool, ToolResult
+
+
+@dataclass
+class MiddlewareEntry:
+    """Entry in middleware chain with inherit override."""
+    middleware: Middleware
+    inherit: bool  # Effective inherit value (config default or overridden)
+
+
+class MiddlewareChain:
+    """Chain of middlewares for sequential processing."""
+    
+    def __init__(self, middlewares: list[Middleware] | None = None) -> None:
+        self._entries: list[MiddlewareEntry] = []
+        self._token_buffer: str = ""
+        self._token_count: int = 0
+        
+        # Add initial middlewares if provided
+        if middlewares:
+            for mw in middlewares:
+                self.use(mw)
+    
+    @property
+    def _middlewares(self) -> list[Middleware]:
+        """Get middleware list."""
+        return [e.middleware for e in self._entries]
+    
+    def use(
+        self,
+        middleware: Middleware,
+        *,
+        inherit: bool | None = None,
+    ) -> "MiddlewareChain":
+        """Add middleware to chain.
+        
+        Args:
+            middleware: The middleware to add
+            inherit: Override inherit setting (None = use middleware's config default)
+        
+        Maintains sorted order by priority.
+        """
+        effective_inherit = inherit if inherit is not None else middleware.config.inherit
+        entry = MiddlewareEntry(middleware=middleware, inherit=effective_inherit)
+        self._entries.append(entry)
+        self._entries.sort(key=lambda e: e.middleware.config.priority)
+        return self
+    
+    def remove(self, middleware: Middleware) -> "MiddlewareChain":
+        """Remove middleware from chain."""
+        self._entries = [e for e in self._entries if e.middleware != middleware]
+        return self
+    
+    def clear(self) -> "MiddlewareChain":
+        """Clear all middlewares."""
+        self._entries.clear()
+        return self
+    
+    def get_inheritable(self) -> list[MiddlewareEntry]:
+        """Get entries that should be inherited by sub-agents."""
+        return [e for e in self._entries if e.inherit]
+    
+    def merge(self, other: "MiddlewareChain | None") -> "MiddlewareChain":
+        """Merge this chain's inheritable middlewares with another chain.
+        
+        Creates a new chain with:
+        - This chain's inheritable middlewares
+        - All of other chain's middlewares
+        
+        Args:
+            other: Chain to merge with (sub-agent's own middlewares)
+            
+        Returns:
+            New merged MiddlewareChain
+        """
+        merged = MiddlewareChain()
+        
+        # Add inheritable from this chain
+        for entry in self.get_inheritable():
+            merged._entries.append(MiddlewareEntry(
+                middleware=entry.middleware,
+                inherit=entry.inherit,
+            ))
+        
+        # Add all from other chain
+        if other:
+            for entry in other._entries:
+                # Avoid duplicates (same middleware instance)
+                if entry.middleware not in [e.middleware for e in merged._entries]:
+                    merged._entries.append(MiddlewareEntry(
+                        middleware=entry.middleware,
+                        inherit=entry.inherit,
+                    ))
+        
+        # Re-sort by priority
+        merged._entries.sort(key=lambda e: e.middleware.config.priority)
+        return merged
+    
+    async def process_request(
+        self,
+        request: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Process request through all middlewares."""
+        current = request
+        
+        for mw in self._middlewares:
+            result = await mw.on_request(current, context)
+            if result is None:
+                return None
+            current = result
+        
+        return current
+    
+    async def process_response(
+        self,
+        response: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Process response through all middlewares (reverse order)."""
+        current = response
+        
+        for mw in reversed(self._middlewares):
+            result = await mw.on_response(current, context)
+            if result is None:
+                return None
+            current = result
+        
+        return current
+    
+    async def process_error(
+        self,
+        error: Exception,
+        context: dict[str, Any],
+    ) -> Exception | None:
+        """Process error through all middlewares."""
+        current = error
+        
+        for mw in self._middlewares:
+            result = await mw.on_error(current, context)
+            if result is None:
+                return None
+            current = result
+        
+        return current
+    
+    async def process_stream_chunk(
+        self,
+        chunk: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Process streaming chunk through middlewares based on trigger mode."""
+        text = chunk.get("text", chunk.get("delta", ""))
+        self._token_buffer += text
+        self._token_count += 1
+        
+        current = chunk
+        
+        for mw in self._middlewares:
+            should_trigger = self._should_trigger(mw, text)
+            
+            if should_trigger:
+                result = await mw.on_model_stream(current, context)
+                if result is None:
+                    return None
+                current = result
+        
+        return current
+    
+    def _should_trigger(self, middleware: Middleware, text: str) -> bool:
+        """Check if middleware should be triggered."""
+        mode = middleware.config.trigger_mode
+        
+        if mode == TriggerMode.EVERY_TOKEN:
+            return True
+        elif mode == TriggerMode.EVERY_N_TOKENS:
+            return self._token_count % middleware.config.trigger_n == 0
+        elif mode == TriggerMode.ON_BOUNDARY:
+            return self._is_boundary(text)
+        
+        return True
+    
+    def _is_boundary(self, text: str) -> bool:
+        """Check if text ends with a sentence/paragraph boundary."""
+        boundaries = (".", "。", "\n", "!", "?", "！", "？", ";", "；")
+        return text.rstrip().endswith(boundaries)
+    
+    def reset_stream_state(self) -> None:
+        """Reset streaming state (call at start of new stream)."""
+        self._token_buffer = ""
+        self._token_count = 0
+    
+    @property
+    def middlewares(self) -> list[Middleware]:
+        """Get list of middlewares (read-only)."""
+        return list(self._middlewares)
+    
+    # ========== Lifecycle Hook Processing ==========
+    
+    async def process_agent_start(
+        self,
+        agent_id: str,
+        input_data: Any,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Process agent start through all middlewares.
+        
+        Returns:
+            First non-CONTINUE result, or CONTINUE if all pass
+        """
+        for mw in self._middlewares:
+            if hasattr(mw, 'on_agent_start'):
+                result = await mw.on_agent_start(agent_id, input_data, context)
+                if result.action != HookAction.CONTINUE:
+                    logger.debug(f"Middleware returned {result.action} on agent_start")
+                    return result
+        return HookResult.proceed()
+    
+    async def process_agent_end(
+        self,
+        agent_id: str,
+        result: Any,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Process agent end through all middlewares (reverse order)."""
+        for mw in reversed(self._middlewares):
+            if hasattr(mw, 'on_agent_end'):
+                hook_result = await mw.on_agent_end(agent_id, result, context)
+                if hook_result.action != HookAction.CONTINUE:
+                    logger.debug(f"Middleware returned {hook_result.action} on agent_end")
+                    return hook_result
+        return HookResult.proceed()
+    
+    async def process_tool_call(
+        self,
+        tool: "BaseTool",
+        params: dict[str, Any],
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Process tool call through all middlewares.
+        
+        Returns:
+            SKIP to skip tool, RETRY with modified_data to change params
+        """
+        for mw in self._middlewares:
+            if hasattr(mw, 'on_tool_call'):
+                result = await mw.on_tool_call(tool, params, context)
+                if result.action != HookAction.CONTINUE:
+                    logger.debug(f"Middleware returned {result.action} on tool_call")
+                    return result
+        return HookResult.proceed()
+    
+    async def process_tool_end(
+        self,
+        tool: "BaseTool",
+        result: "ToolResult",
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Process tool end through all middlewares (reverse order)."""
+        for mw in reversed(self._middlewares):
+            if hasattr(mw, 'on_tool_end'):
+                hook_result = await mw.on_tool_end(tool, result, context)
+                if hook_result.action != HookAction.CONTINUE:
+                    logger.debug(f"Middleware returned {hook_result.action} on tool_end")
+                    return hook_result
+        return HookResult.proceed()
+    
+    async def process_subagent_start(
+        self,
+        parent_agent_id: str,
+        child_agent_id: str,
+        mode: str,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Process sub-agent start through all middlewares."""
+        for mw in self._middlewares:
+            if hasattr(mw, 'on_subagent_start'):
+                result = await mw.on_subagent_start(
+                    parent_agent_id, child_agent_id, mode, context
+                )
+                if result.action != HookAction.CONTINUE:
+                    logger.debug(f"Middleware returned {result.action} on subagent_start")
+                    return result
+        return HookResult.proceed()
+    
+    async def process_subagent_end(
+        self,
+        parent_agent_id: str,
+        child_agent_id: str,
+        result: Any,
+        context: dict[str, Any],
+    ) -> HookResult:
+        """Process sub-agent end through all middlewares (reverse order)."""
+        for mw in reversed(self._middlewares):
+            if hasattr(mw, 'on_subagent_end'):
+                hook_result = await mw.on_subagent_end(
+                    parent_agent_id, child_agent_id, result, context
+                )
+                if hook_result.action != HookAction.CONTINUE:
+                    logger.debug(f"Middleware returned {hook_result.action} on subagent_end")
+                    return hook_result
+        return HookResult.proceed()
+    
+    async def process_message_save(
+        self,
+        message: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Process message save through all middlewares.
+        
+        Args:
+            message: Message to be saved
+            context: Execution context
+            
+        Returns:
+            Modified message, or None to skip saving
+        """
+        current = message
+        
+        for mw in self._middlewares:
+            if hasattr(mw, 'on_message_save'):
+                result = await mw.on_message_save(current, context)
+                if result is None:
+                    logger.debug("Middleware blocked message save")
+                    return None
+                current = result
+        
+        return current
+
+
+__all__ = ["MiddlewareChain"]

aury/agents/middleware/message.py

@@ -0,0 +1,129 @@
+"""Message persistence middleware.
+
+Uses backends.message directly to persist messages.
+
+Usage:
+    middleware = MessageBackendMiddleware()
+    
+    agent = ReactAgent.create(
+        llm=llm,
+        middlewares=[middleware],
+    )
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+from .base import BaseMiddleware
+
+if TYPE_CHECKING:
+    from ..backends import MessageBackend
+
+
+class MessageBackendMiddleware(BaseMiddleware):
+    """Middleware that persists messages directly via MessageBackend.
+    
+    This is the recommended middleware for message persistence.
+    Uses backends.message from the context for direct storage.
+    
+    Features:
+    - Saves both truncated (for LLM context) and raw (for audit) messages
+    - Supports namespace isolation for sub-agents
+    - Works with any MessageBackend implementation
+    
+    Example:
+        agent = ReactAgent.create(
+            llm=llm,
+            middlewares=[MessageBackendMiddleware()],
+        )
+    """
+    
+    def __init__(
+        self,
+        *,
+        save_raw: bool = False,
+        max_history: int = 100,
+    ):
+        """Initialize MessageBackendMiddleware.
+        
+        Args:
+            save_raw: Also save raw (untruncated) messages for audit
+            max_history: Max messages to keep in history (for truncation)
+        """
+        self.save_raw = save_raw
+        self.max_history = max_history
+    
+    async def on_message_save(
+        self,
+        message: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Save message via backends.message.
+        
+        Args:
+            message: Message dict with 'role', 'content', etc.
+            context: Execution context with 'session_id', 'agent_id', 'backends'
+            
+        Returns:
+            The message (pass through to other middlewares)
+        """
+        from ..core.context import get_current_ctx_or_none
+        
+        session_id = context.get("session_id", "")
+        if not session_id:
+            return message
+        
+        # Get MessageBackend from context
+        ctx = get_current_ctx_or_none()
+        if ctx is None or ctx.backends is None or ctx.backends.message is None:
+            # No backend available, pass through
+            return message
+        
+        backend = ctx.backends.message
+        
+        # Extract message fields
+        role = message.get("role", "")
+        content = message.get("content", "")
+        invocation_id = context.get("invocation_id", "")
+        agent_id = context.get("agent_id")
+        namespace = context.get("namespace")
+        tool_call_id = message.get("tool_call_id")
+        
+        # Build message dict for backend
+        msg_dict = {
+            "role": role,
+            "content": content,
+        }
+        if tool_call_id:
+            msg_dict["tool_call_id"] = tool_call_id
+        
+        # Include tool_calls if present (for assistant messages)
+        if "tool_calls" in message:
+            msg_dict["tool_calls"] = message["tool_calls"]
+        
+        # Save truncated message (for LLM context)
+        await backend.add(
+            session_id=session_id,
+            message=msg_dict,
+            type="truncated",
+            agent_id=agent_id,
+            namespace=namespace,
+            invocation_id=invocation_id,
+        )
+        
+        # Optionally save raw message (for audit)
+        if self.save_raw:
+            await backend.add(
+                session_id=session_id,
+                message=message,  # Full original message
+                type="raw",
+                agent_id=agent_id,
+                namespace=namespace,
+                invocation_id=invocation_id,
+            )
+        
+        # Pass through to other middlewares
+        return message
+
+
+__all__ = ["MessageBackendMiddleware"]

aury/agents/middleware/message_container.py

@@ -0,0 +1,126 @@
+"""MessageContainerMiddleware - Groups thinking and text blocks under a message container.
+
+This middleware creates a "message" container block when LLM request starts,
+and all thinking/text blocks emitted during the LLM call will have parent_id
+pointing to this container.
+
+Usage:
+    from aury.agents.middleware import MessageContainerMiddleware, MiddlewareChain
+    
+    chain = MiddlewareChain()
+    chain.use(MessageContainerMiddleware())
+    
+    agent = ReactAgent.create(llm=llm, middleware=chain)
+
+Result structure:
+    message (block_id: blk_abc)
+    ├── thinking (parent_id: blk_abc)
+    └── text (parent_id: blk_abc)
+    
+    tool_use (parent_id: None) - not grouped
+    tool_result (parent_id: None) - not grouped
+"""
+from __future__ import annotations
+
+from typing import Any, TYPE_CHECKING
+
+from .base import BaseMiddleware
+from .types import HookResult
+from ..core.context import set_parent_id, reset_parent_id, emit, get_parent_id
+from ..core.types.session import generate_id
+from ..core.types.block import BlockEvent, BlockKind, BlockOp
+from ..core.logging import middleware_logger as logger
+
+if TYPE_CHECKING:
+    from ..core.types.tool import BaseTool, ToolResult
+
+
+class MessageContainerMiddleware(BaseMiddleware):
+    """Groups thinking and text blocks under a message container.
+    
+    When an LLM request starts, creates a "message" container block and sets
+    the parent_id ContextVar. Only thinking and text blocks will inherit this
+    parent_id (tool_use, tool_result, etc. are not grouped).
+    
+    This allows frontend to group thinking + text as a single unit for display.
+    
+    Args:
+        apply_to_kinds: Set of block kinds that should inherit the container's
+                       parent_id. Defaults to {"thinking", "text"}.
+    """
+    
+    # Key to store the token in middleware context
+    _TOKEN_KEY = "_message_container_token"
+    _BLOCK_ID_KEY = "_message_container_block_id"
+    
+    # Default kinds that should be grouped under message container
+    DEFAULT_KINDS = {"thinking", "text"}
+    
+    def __init__(self, apply_to_kinds: set[str] | None = None):
+        """Initialize with optional custom kinds filter.
+        
+        Args:
+            apply_to_kinds: Block kinds to group. Defaults to {"thinking", "text"}.
+        """
+        super().__init__()
+        self.apply_to_kinds = apply_to_kinds or self.DEFAULT_KINDS
+    
+    async def on_request(
+        self,
+        request: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Create message container block and set parent_id."""
+        # Generate container block ID
+        message_block_id = generate_id("blk")
+        
+        logger.debug(
+            "[MessageContainerMiddleware] Creating container",
+            extra={"block_id": message_block_id, "apply_to_kinds": list(self.apply_to_kinds)}
+        )
+        
+        # Emit the container block
+        await emit(BlockEvent(
+            block_id=message_block_id,
+            kind="message",  # Container type
+            op=BlockOp.APPLY,
+            data={
+                "type": "llm_response",
+                "step": context.get("step"),
+            },
+            session_id=context.get("session_id"),
+            invocation_id=context.get("invocation_id"),
+        ))
+        
+        # Set parent_id in ContextVar with apply_to_kinds filter
+        # Only blocks matching apply_to_kinds will inherit this parent_id
+        token = set_parent_id(message_block_id, apply_to_kinds=self.apply_to_kinds)
+        context[self._TOKEN_KEY] = token
+        context[self._BLOCK_ID_KEY] = message_block_id
+        
+        return request
+    
+    async def on_response(
+        self,
+        response: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Reset parent_id to previous value."""
+        token = context.get(self._TOKEN_KEY)
+        if token is not None:
+            reset_parent_id(token)
+        return response
+    
+    async def on_error(
+        self,
+        error: Exception,
+        context: dict[str, Any],
+    ) -> Exception | None:
+        """Reset parent_id on error too."""
+        token = context.get(self._TOKEN_KEY)
+        if token is not None:
+            reset_parent_id(token)
+        return error
+
+
+__all__ = ["MessageContainerMiddleware"]