aury-agent 0.0.4__py3-none-any.whl → 0.0.6__py3-none-any.whl

aury/agents/messages/store.py

@@ -72,10 +72,22 @@ class InMemoryMessageStore:
         message: Message,
         namespace: str | None = None,
     ) -> None:
+        from ..core.logging import storage_logger as logger
+        
         key = self._make_key(session_id, namespace)
         if key not in self._messages:
             self._messages[key] = []
         self._messages[key].append(message)
+        
+        logger.debug(
+            "Message stored",
+            extra={
+                "session_id": session_id,
+                "invocation_id": getattr(message, "invocation_id", None),
+                "role": getattr(message, "role", None),
+                "namespace": namespace,
+            },
+        )
     
     async def get_all(
         self,
@@ -101,6 +113,8 @@ class InMemoryMessageStore:
         invocation_id: str,
         namespace: str | None = None,
     ) -> int:
+        from ..core.logging import storage_logger as logger
+        
         key = self._make_key(session_id, namespace)
         if key not in self._messages:
             return 0
@@ -109,7 +123,19 @@ class InMemoryMessageStore:
         self._messages[key] = [
             m for m in original if m.invocation_id != invocation_id
         ]
-        return len(original) - len(self._messages[key])
+        deleted_count = len(original) - len(self._messages[key])
+        
+        if deleted_count > 0:
+            logger.debug(
+                "Messages deleted by invocation",
+                extra={
+                    "session_id": session_id,
+                    "invocation_id": invocation_id,
+                    "count": deleted_count,
+                },
+            )
+        
+        return deleted_count
 
 
 __all__ = [

aury/agents/middleware/base.py

@@ -87,6 +87,22 @@ class Middleware(Protocol):
         """
         ...
     
+    async def on_thinking_stream(
+        self,
+        chunk: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Process thinking stream chunk.
+        
+        Args:
+            chunk: The thinking chunk with {"delta": str}
+            context: Execution context
+            
+        Returns:
+            Modified chunk, or None to skip
+        """
+        ...
+    
     # ========== Agent Lifecycle Hooks ==========
     
     async def on_agent_start(
@@ -143,6 +159,29 @@ class Middleware(Protocol):
         """
         ...
     
+    async def on_tool_call_delta(
+        self,
+        call_id: str,
+        tool_name: str,
+        delta: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Called during streaming tool argument generation.
+        
+        Only triggered for tools with stream_arguments=True.
+        Receives incremental updates as LLM generates tool parameters.
+        
+        Args:
+            call_id: Tool call identifier
+            tool_name: Name of the tool being called
+            delta: Incremental parameter update (e.g. {"content": "more text"})
+            context: Execution context with 'accumulated_args' containing current state
+            
+        Returns:
+            Modified delta, or None to skip emitting this delta
+        """
+        ...
+    
     async def on_tool_end(
         self,
         tool: "BaseTool",
@@ -268,6 +307,14 @@ class BaseMiddleware:
         """Default: pass through."""
         return chunk
     
+    async def on_thinking_stream(
+        self,
+        chunk: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Default: pass through."""
+        return chunk
+    
     # ========== Agent Lifecycle Hooks ==========
     
     async def on_agent_start(
@@ -297,6 +344,16 @@ class BaseMiddleware:
         """Default: continue."""
         return HookResult.proceed()
     
+    async def on_tool_call_delta(
+        self,
+        call_id: str,
+        tool_name: str,
+        delta: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Default: pass through."""
+        return delta
+    
     async def on_tool_end(
         self,
         tool: "BaseTool",

aury/agents/middleware/chain.py

@@ -29,6 +29,7 @@ class MiddlewareChain:
         
         # Add initial middlewares if provided
         if middlewares:
+            logger.debug(f"MiddlewareChain init with {len(middlewares)} middlewares")
             for mw in middlewares:
                 self.use(mw)
     
@@ -55,6 +56,7 @@ class MiddlewareChain:
         entry = MiddlewareEntry(middleware=middleware, inherit=effective_inherit)
         self._entries.append(entry)
         self._entries.sort(key=lambda e: e.middleware.config.priority)
+        logger.debug(f"Added middleware to chain, priority={middleware.config.priority}, inherit={effective_inherit}, total={len(self._entries)}")
         return self
     
     def remove(self, middleware: Middleware) -> "MiddlewareChain":
@@ -114,13 +116,16 @@ class MiddlewareChain:
     ) -> dict[str, Any] | None:
         """Process request through all middlewares."""
         current = request
+        logger.debug(f"Processing request through {len(self._middlewares)} middlewares")
         
-        for mw in self._middlewares:
+        for i, mw in enumerate(self._middlewares):
             result = await mw.on_request(current, context)
             if result is None:
+                logger.info(f"Middleware #{i} blocked request")
                 return None
             current = result
         
+        logger.debug("Request processing completed")
         return current
     
     async def process_response(
@@ -130,13 +135,16 @@ class MiddlewareChain:
     ) -> dict[str, Any] | None:
         """Process response through all middlewares (reverse order)."""
         current = response
+        logger.debug(f"Processing response through {len(self._middlewares)} middlewares (reverse order)")
         
-        for mw in reversed(self._middlewares):
+        for i, mw in enumerate(reversed(self._middlewares)):
             result = await mw.on_response(current, context)
             if result is None:
+                logger.info(f"Middleware #{i} blocked response")
                 return None
             current = result
         
+        logger.debug("Response processing completed")
         return current
     
     async def process_error(
@@ -146,13 +154,16 @@ class MiddlewareChain:
     ) -> Exception | None:
         """Process error through all middlewares."""
         current = error
+        logger.debug(f"Processing error {type(error).__name__} through {len(self._middlewares)} middlewares")
         
-        for mw in self._middlewares:
+        for i, mw in enumerate(self._middlewares):
             result = await mw.on_error(current, context)
             if result is None:
+                logger.info(f"Middleware #{i} suppressed error")
                 return None
             current = result
         
+        logger.debug("Error processing completed")
         return current
     
     async def process_stream_chunk(
@@ -164,18 +175,55 @@ class MiddlewareChain:
         text = chunk.get("text", chunk.get("delta", ""))
         self._token_buffer += text
         self._token_count += 1
+        logger.debug(f"Processing stream chunk, token_count={self._token_count}, triggered_middlewares=?")
         
         current = chunk
+        triggered_count = 0
         
-        for mw in self._middlewares:
+        for i, mw in enumerate(self._middlewares):
             should_trigger = self._should_trigger(mw, text)
             
             if should_trigger:
+                triggered_count += 1
+                logger.debug(f"Middleware #{i} triggered, mode={mw.config.trigger_mode}")
                 result = await mw.on_model_stream(current, context)
                 if result is None:
+                    logger.info(f"Middleware #{i} blocked stream chunk")
                     return None
                 current = result
         
+        logger.debug(f"Stream chunk processing completed, {triggered_count} middlewares triggered")
+        return current
+    
+    async def process_tool_call_delta(
+        self,
+        call_id: str,
+        tool_name: str,
+        delta: dict[str, Any],
+        context: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Process tool call delta through all middlewares.
+        
+        Args:
+            call_id: Tool call identifier
+            tool_name: Name of the tool being called
+            delta: Incremental parameter update
+            context: Execution context with 'accumulated_args'
+            
+        Returns:
+            Modified delta, or None to skip emitting
+        """
+        current = delta
+        logger.debug(f"Processing tool_call_delta for {tool_name} (call_id={call_id}) through {len(self._middlewares)} middlewares")
+        
+        for i, mw in enumerate(self._middlewares):
+            result = await mw.on_tool_call_delta(call_id, tool_name, current, context)
+            if result is None:
+                logger.info(f"Middleware #{i} blocked tool_call_delta")
+                return None
+            current = result
+        
+        logger.debug("Tool call delta processing completed")
         return current
     
     def _should_trigger(self, middleware: Middleware, text: str) -> bool:
@@ -198,6 +246,7 @@ class MiddlewareChain:
     
     def reset_stream_state(self) -> None:
         """Reset streaming state (call at start of new stream)."""
+        logger.debug("Resetting stream state")
         self._token_buffer = ""
         self._token_count = 0
     
@@ -219,12 +268,14 @@ class MiddlewareChain:
         Returns:
             First non-CONTINUE result, or CONTINUE if all pass
         """
-        for mw in self._middlewares:
+        logger.debug(f"Processing agent_start for agent_id={agent_id}, {len(self._middlewares)} middlewares")
+        for i, mw in enumerate(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")
+                    logger.info(f"Middleware #{i} returned {result.action} on agent_start")
                     return result
+        logger.debug("Agent start processing completed, all middlewares passed")
         return HookResult.proceed()
     
     async def process_agent_end(
@@ -234,12 +285,14 @@ class MiddlewareChain:
         context: dict[str, Any],
     ) -> HookResult:
         """Process agent end through all middlewares (reverse order)."""
-        for mw in reversed(self._middlewares):
+        logger.debug(f"Processing agent_end for agent_id={agent_id}, {len(self._middlewares)} middlewares (reverse order)")
+        for i, mw in enumerate(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")
+                    logger.info(f"Middleware #{i} returned {hook_result.action} on agent_end")
                     return hook_result
+        logger.debug("Agent end processing completed, all middlewares passed")
         return HookResult.proceed()
     
     async def process_tool_call(
@@ -253,12 +306,14 @@ class MiddlewareChain:
         Returns:
             SKIP to skip tool, RETRY with modified_data to change params
         """
-        for mw in self._middlewares:
+        logger.debug(f"Processing tool_call for tool={tool.name}, {len(self._middlewares)} middlewares")
+        for i, mw in enumerate(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")
+                    logger.info(f"Middleware #{i} returned {result.action} on tool_call for tool={tool.name}")
                     return result
+        logger.debug("Tool call processing completed, all middlewares passed")
         return HookResult.proceed()
     
     async def process_tool_end(
@@ -268,12 +323,14 @@ class MiddlewareChain:
         context: dict[str, Any],
     ) -> HookResult:
         """Process tool end through all middlewares (reverse order)."""
-        for mw in reversed(self._middlewares):
+        logger.debug(f"Processing tool_end for tool={tool.name}, {len(self._middlewares)} middlewares (reverse order)")
+        for i, mw in enumerate(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")
+                    logger.info(f"Middleware #{i} returned {hook_result.action} on tool_end for tool={tool.name}")
                     return hook_result
+        logger.debug("Tool end processing completed, all middlewares passed")
         return HookResult.proceed()
     
     async def process_subagent_start(
@@ -284,14 +341,16 @@ class MiddlewareChain:
         context: dict[str, Any],
     ) -> HookResult:
         """Process sub-agent start through all middlewares."""
-        for mw in self._middlewares:
+        logger.debug(f"Processing subagent_start, parent={parent_agent_id}, child={child_agent_id}, mode={mode}, {len(self._middlewares)} middlewares")
+        for i, mw in enumerate(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")
+                    logger.info(f"Middleware #{i} returned {result.action} on subagent_start")
                     return result
+        logger.debug("Subagent start processing completed, all middlewares passed")
         return HookResult.proceed()
     
     async def process_subagent_end(
@@ -302,14 +361,16 @@ class MiddlewareChain:
         context: dict[str, Any],
     ) -> HookResult:
         """Process sub-agent end through all middlewares (reverse order)."""
-        for mw in reversed(self._middlewares):
+        logger.debug(f"Processing subagent_end, parent={parent_agent_id}, child={child_agent_id}, {len(self._middlewares)} middlewares (reverse order)")
+        for i, mw in enumerate(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")
+                    logger.info(f"Middleware #{i} returned {hook_result.action} on subagent_end")
                     return hook_result
+        logger.debug("Subagent end processing completed, all middlewares passed")
         return HookResult.proceed()
     
     async def process_message_save(
@@ -327,15 +388,17 @@ class MiddlewareChain:
             Modified message, or None to skip saving
         """
         current = message
+        logger.debug(f"Processing message_save, role={message.get('role')}, {len(self._middlewares)} middlewares")
         
-        for mw in self._middlewares:
+        for i, mw in enumerate(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")
+                    logger.info(f"Middleware #{i} blocked message save for role={message.get('role')}")
                     return None
                 current = result
         
+        logger.debug("Message save processing completed, all middlewares passed")
         return current