aury-agent 0.0.9__py3-none-any.whl → 0.0.10__py3-none-any.whl

aury/agents/backends/__init__.py

@@ -5,7 +5,7 @@ Backends provide abstracted interfaces for various capabilities:
 Data Backends (storage):
 - SessionBackend: Session management
 - InvocationBackend: Invocation management
-- MessageBackend: Message storage (truncated/raw)
+- MessageBackend: Message storage
 - MemoryBackend: Long-term memory with search
 - ArtifactBackend: File/artifact storage
 - StateBackend: Generic key-value state
@@ -26,7 +26,7 @@ from typing import TYPE_CHECKING
 # Data backends - new architecture
 from .session import SessionBackend, InMemorySessionBackend
 from .invocation import InvocationBackend, InMemoryInvocationBackend
-from .message import MessageBackend, MessageType, InMemoryMessageBackend
+from .message import MessageBackend, InMemoryMessageBackend
 from .memory import MemoryBackend, InMemoryMemoryBackend
 from .artifact import ArtifactBackend, ArtifactSource, InMemoryArtifactBackend
 
@@ -144,7 +144,6 @@ __all__ = [
     
     # Message backend
     "MessageBackend",
-    "MessageType",
     "InMemoryMessageBackend",
     
     # Memory backend

aury/agents/backends/message/__init__.py

@@ -1,9 +1,8 @@
 """Message backend."""
-from .types import MessageBackend, MessageType
+from .types import MessageBackend
 from .memory import InMemoryMessageBackend
 
 __all__ = [
     "MessageBackend",
-    "MessageType",
     "InMemoryMessageBackend",
 ]

aury/agents/backends/message/memory.py

@@ -4,45 +4,36 @@ from __future__ import annotations
 from datetime import datetime
 from typing import Any
 
-from .types import MessageType
-
 
 class InMemoryMessageBackend:
     """In-memory implementation of MessageBackend.
     
-    Stores both truncated and raw messages in separate dicts.
-    Suitable for testing and simple single-process use cases.
+    Simple in-memory storage for testing and single-process use cases.
     """
     
     def __init__(self) -> None:
         # Key format: "{session_id}" or "{session_id}:{namespace}"
         # Value: list of message dicts
-        self._truncated: dict[str, list[dict[str, Any]]] = {}
-        self._raw: dict[str, list[dict[str, Any]]] = {}
+        self._messages: dict[str, list[dict[str, Any]]] = {}
     
     def _make_key(self, session_id: str, namespace: str | None) -> str:
         if namespace:
             return f"{session_id}:{namespace}"
         return session_id
     
-    def _get_store(self, type: MessageType) -> dict[str, list[dict[str, Any]]]:
-        return self._truncated if type == "truncated" else self._raw
-    
     async def add(
         self,
         session_id: str,
         message: dict[str, Any],
-        type: MessageType = "truncated",
         agent_id: str | None = None,
         namespace: str | None = None,
         invocation_id: str | None = None,
     ) -> None:
         """Add a message."""
         key = self._make_key(session_id, namespace)
-        store = self._get_store(type)
         
-        if key not in store:
-            store[key] = []
+        if key not in self._messages:
+            self._messages[key] = []
         
         # Add metadata
         msg = {
@@ -51,20 +42,18 @@ class InMemoryMessageBackend:
             "invocation_id": invocation_id,
             "created_at": datetime.now().isoformat(),
         }
-        store[key].append(msg)
+        self._messages[key].append(msg)
     
     async def get(
         self,
         session_id: str,
-        type: MessageType = "truncated",
         agent_id: str | None = None,
         namespace: str | None = None,
         limit: int | None = None,
     ) -> list[dict[str, Any]]:
         """Get messages."""
         key = self._make_key(session_id, namespace)
-        store = self._get_store(type)
-        messages = store.get(key, [])
+        messages = self._messages.get(key, [])
         
         # Filter by agent_id if specified
         if agent_id:
@@ -80,42 +69,31 @@ class InMemoryMessageBackend:
         self,
         session_id: str,
         invocation_id: str,
-        type: MessageType | None = None,
         namespace: str | None = None,
     ) -> int:
         """Delete messages by invocation."""
         key = self._make_key(session_id, namespace)
-        deleted = 0
-        
-        types_to_delete = [type] if type else ["truncated", "raw"]
         
-        for t in types_to_delete:
-            store = self._get_store(t)
-            if key in store:
-                original = store[key]
-                store[key] = [m for m in original if m.get("invocation_id") != invocation_id]
-                deleted += len(original) - len(store[key])
+        if key not in self._messages:
+            return 0
         
-        return deleted
+        original = self._messages[key]
+        self._messages[key] = [m for m in original if m.get("invocation_id") != invocation_id]
+        return len(original) - len(self._messages[key])
     
     async def clear(
         self,
         session_id: str,
-        type: MessageType | None = None,
         namespace: str | None = None,
     ) -> int:
         """Clear all messages for a session."""
         key = self._make_key(session_id, namespace)
-        deleted = 0
-        
-        types_to_clear = [type] if type else ["truncated", "raw"]
         
-        for t in types_to_clear:
-            store = self._get_store(t)
-            if key in store:
-                deleted += len(store[key])
-                del store[key]
+        if key not in self._messages:
+            return 0
         
+        deleted = len(self._messages[key])
+        del self._messages[key]
         return deleted
 
 

aury/agents/backends/message/types.py

@@ -1,49 +1,29 @@
 """Message backend types and protocols."""
 from __future__ import annotations
 
-from typing import Any, Literal, Protocol, runtime_checkable
-
-
-MessageType = Literal["truncated", "raw"]
+from typing import Any, Protocol, runtime_checkable
 
 
 @runtime_checkable
 class MessageBackend(Protocol):
     """Protocol for message storage.
     
-    Handles both truncated (context window) and raw (full history) messages
-    through a unified interface with type parameter.
-    
-    - truncated: Messages kept in context window, may be summarized/trimmed
-    - raw: Full original messages for audit/replay
+    Simple interface for message persistence.
+    Storage details (raw/truncated handling) are left to the application layer.
     
     Example usage:
-        # Add truncated message (for LLM context)
         await backend.add(
             session_id="sess_123",
             message={"role": "user", "content": "Hello"},
-            type="truncated",
-        )
-        
-        # Add raw message (for audit)
-        await backend.add(
-            session_id="sess_123",
-            message={"role": "user", "content": "Hello", "attachments": [...]},
-            type="raw",
         )
         
-        # Get messages for LLM
-        messages = await backend.get("sess_123", type="truncated", limit=50)
-        
-        # Get raw history
-        raw_messages = await backend.get("sess_123", type="raw")
+        messages = await backend.get("sess_123", limit=50)
     """
     
     async def add(
         self,
         session_id: str,
         message: dict[str, Any],
-        type: MessageType = "truncated",
         agent_id: str | None = None,
         namespace: str | None = None,
         invocation_id: str | None = None,
@@ -53,7 +33,6 @@ class MessageBackend(Protocol):
         Args:
             session_id: Session ID
             message: Message dict (role, content, tool_call_id, etc.)
-            type: Message type - "truncated" or "raw"
             agent_id: Optional agent ID
             namespace: Optional namespace for sub-agent isolation
             invocation_id: Optional invocation ID for grouping
@@ -63,7 +42,6 @@ class MessageBackend(Protocol):
     async def get(
         self,
         session_id: str,
-        type: MessageType = "truncated",
         agent_id: str | None = None,
         namespace: str | None = None,
         limit: int | None = None,
@@ -72,7 +50,6 @@ class MessageBackend(Protocol):
         
         Args:
             session_id: Session ID
-            type: Message type - "truncated" or "raw"
             agent_id: Optional filter by agent
             namespace: Optional namespace filter
             limit: Max messages to return (None = all)
@@ -86,7 +63,6 @@ class MessageBackend(Protocol):
         self,
         session_id: str,
         invocation_id: str,
-        type: MessageType | None = None,
         namespace: str | None = None,
     ) -> int:
         """Delete messages by invocation (for revert).
@@ -94,7 +70,6 @@ class MessageBackend(Protocol):
         Args:
             session_id: Session ID
             invocation_id: Invocation ID to delete
-            type: Message type to delete, None = both types
             namespace: Optional namespace filter
             
         Returns:
@@ -105,14 +80,12 @@ class MessageBackend(Protocol):
     async def clear(
         self,
         session_id: str,
-        type: MessageType | None = None,
         namespace: str | None = None,
     ) -> int:
         """Clear all messages for a session.
         
         Args:
             session_id: Session ID
-            type: Message type to clear, None = both types
             namespace: Optional namespace filter
             
         Returns:
@@ -121,4 +94,4 @@ class MessageBackend(Protocol):
         ...
 
 
-__all__ = ["MessageBackend", "MessageType"]
+__all__ = ["MessageBackend"]

aury/agents/context_providers/message.py

@@ -79,7 +79,6 @@ class MessageContextProvider(BaseContextProvider):
         if ctx.backends is not None and ctx.backends.message is not None:
             messages = await ctx.backends.message.get(
                 session_id=ctx.session.id,
-                type="truncated",
                 limit=self.max_messages,
             )
             # Convert to LLM format (include tool_call_id for tool messages)

aury/agents/core/context.py

@@ -728,11 +728,11 @@ class InvocationContext:
                 **{k: v for k, v in request.items() if k not in ("messages", "stream")}
             ):
                 if self.middleware:
-                    chunk_dict = {"chunk": chunk}
-                    processed = await self.middleware.process_stream_chunk(chunk_dict)
+                    chunk_dict = {"delta": chunk}
+                    processed = await self.middleware.process_text_stream(chunk_dict)
                     if processed is None:
                         continue
-                    chunk = processed.get("chunk", chunk)
+                    chunk = processed.get("delta", chunk)
                 yield chunk
                 
         except Exception as e:

aury/agents/core/types/tool.py

@@ -90,14 +90,41 @@ class ToolContext:
 
 @dataclass
 class ToolResult:
-    """Tool execution result for LLM."""
-    output: str
+    """Tool execution result for LLM.
+    
+    Supports dual output for context management:
+    - output: Complete output (raw), for storage and recall
+    - truncated_output: Shortened output for context window
+    
+    If truncated_output is not provided, it defaults to output.
+    """
+    output: str  # Complete output (raw)
     is_error: bool = False
+    truncated_output: str | None = None  # Shortened output (defaults to output)
+    
+    def __post_init__(self):
+        # Default truncated to output if not provided
+        if self.truncated_output is None:
+            self.truncated_output = self.output
     
     @classmethod
-    def success(cls, output: str) -> ToolResult:
-        """Create a successful result."""
-        return cls(output=output, is_error=False)
+    def success(
+        cls,
+        output: str,
+        *,
+        truncated_output: str | None = None,
+    ) -> ToolResult:
+        """Create a successful result.
+        
+        Args:
+            output: Complete output (raw)
+            truncated_output: Shortened output for context (defaults to output)
+        """
+        return cls(
+            output=output,
+            is_error=False,
+            truncated_output=truncated_output,
+        )
     
     @classmethod
     def error(cls, message: str) -> ToolResult:

aury/agents/llm/provider.py

@@ -133,18 +133,31 @@ class LLMMessage:
         - system: System prompt
         - user: User message (can include images)
         - assistant: Assistant response (can include tool_calls)
-        - tool: Tool result (requires tool_call_id)
+        - tool: Tool result (requires tool_call_id and name)
     """
     role: Literal["system", "user", "assistant", "tool"]
     content: str | list[dict[str, Any]]
     tool_call_id: str | None = None  # Required for tool role
+    name: str | None = None  # Tool name, required for Gemini compatibility
     
     def to_dict(self) -> dict[str, Any]:
         d = {"role": self.role, "content": self.content}
         if self.tool_call_id:
             d["tool_call_id"] = self.tool_call_id
+        if self.name:
+            d["name"] = self.name
         return d
     
+    def get(self, key: str, default: Any = None) -> Any:
+        """Dict-like access for middleware compatibility."""
+        return getattr(self, key, default)
+    
+    def __getitem__(self, key: str) -> Any:
+        """Dict-like access via []."""
+        if hasattr(self, key):
+            return getattr(self, key)
+        raise KeyError(key)
+    
     @classmethod
     def system(cls, content: str) -> "LLMMessage":
         """Create system message."""
@@ -161,9 +174,15 @@ class LLMMessage:
         return cls(role="assistant", content=content)
     
     @classmethod
-    def tool(cls, content: str, tool_call_id: str) -> "LLMMessage":
-        """Create tool result message."""
-        return cls(role="tool", content=content, tool_call_id=tool_call_id)
+    def tool(cls, content: str, tool_call_id: str, name: str | None = None) -> "LLMMessage":
+        """Create tool result message.
+        
+        Args:
+            content: Tool result content
+            tool_call_id: ID of the tool call this result is for
+            name: Tool name (required for Gemini compatibility)
+        """
+        return cls(role="tool", content=content, tool_call_id=tool_call_id, name=name)
 
 
 @runtime_checkable

aury/agents/messages/__init__.py

@@ -15,11 +15,6 @@ from .store import (
     MessageStore,
     InMemoryMessageStore,
 )
-from .raw_store import (
-    RawMessageStore,
-    StateBackendRawMessageStore,
-    InMemoryRawMessageStore,
-)
 from .config import (
     MessageConfig,
 )
@@ -31,10 +26,6 @@ __all__ = [
     # Store (protocol + in-memory for testing)
     "MessageStore",
     "InMemoryMessageStore",
-    # Raw Store
-    "RawMessageStore",
-    "StateBackendRawMessageStore",
-    "InMemoryRawMessageStore",
     # Config
     "MessageConfig",
 ]

aury/agents/middleware/__init__.py

@@ -13,7 +13,6 @@ from .chain import MiddlewareChain
 from .message_container import MessageContainerMiddleware
 from .message import MessageBackendMiddleware
 from .truncation import MessageTruncationMiddleware
-from .raw_message import RawMessageMiddleware
 
 __all__ = [
     "TriggerMode",
@@ -27,5 +26,4 @@ __all__ = [
     "MessageContainerMiddleware",
     "MessageBackendMiddleware",
     "MessageTruncationMiddleware",
-    "RawMessageMiddleware",
 ]

aury/agents/middleware/base.py

@@ -73,17 +73,28 @@ class Middleware(Protocol):
         """
         ...
     
-    async def on_model_stream(
+    async def on_text_stream(
         self,
         chunk: dict[str, Any],
     ) -> dict[str, Any] | None:
-        """Process streaming chunk (triggered by trigger_mode).
+        """Process text streaming chunk.
         
         Args:
-            chunk: The streaming chunk
+            chunk: The text chunk with {"delta": str}
             
         Returns:
-            Modified chunk, or None to skip further processing
+            Modified chunk, or None to skip
+        """
+        ...
+    
+    async def on_text_stream_end(self) -> dict[str, Any] | None:
+        """Called when text stream ends.
+        
+        Use this to flush any buffered text content.
+        
+        Returns:
+            Optional dict with {"delta": str} to emit final content,
+            or None if no additional content.
         """
         ...
     
@@ -101,6 +112,17 @@ class Middleware(Protocol):
         """
         ...
     
+    async def on_thinking_stream_end(self) -> dict[str, Any] | None:
+        """Called when thinking stream ends.
+        
+        Use this to flush any buffered thinking content.
+        
+        Returns:
+            Optional dict with {"delta": str} to emit final thinking content,
+            or None if no additional content.
+        """
+        ...
+    
     # ========== Agent Lifecycle Hooks ==========
     
     async def on_agent_start(
@@ -283,13 +305,17 @@ class BaseMiddleware:
         """Default: re-raise error."""
         return error
     
-    async def on_model_stream(
+    async def on_text_stream(
         self,
         chunk: dict[str, Any],
     ) -> dict[str, Any] | None:
         """Default: pass through."""
         return chunk
     
+    async def on_text_stream_end(self) -> dict[str, Any] | None:
+        """Default: no additional content."""
+        return None
+    
     async def on_thinking_stream(
         self,
         chunk: dict[str, Any],
@@ -297,6 +323,10 @@ class BaseMiddleware:
         """Default: pass through."""
         return chunk
     
+    async def on_thinking_stream_end(self) -> dict[str, Any] | None:
+        """Default: no additional content."""
+        return None
+    
     # ========== Agent Lifecycle Hooks ==========
     
     async def on_agent_start(

aury/agents/middleware/chain.py

@@ -163,32 +163,44 @@ class MiddlewareChain:
         logger.debug("Error processing completed")
         return current
     
-    async def process_stream_chunk(
+    async def process_text_stream(
         self,
         chunk: dict[str, Any],
     ) -> dict[str, Any] | None:
-        """Process streaming chunk through middlewares based on trigger mode."""
-        text = chunk.get("text", chunk.get("delta", ""))
+        """Process text streaming chunk through middlewares based on trigger mode."""
+        text = chunk.get("delta", "")
         self._token_buffer += text
         self._token_count += 1
         
         current = chunk
-        triggered_count = 0
         
         for i, mw in enumerate(self._middlewares):
             should_trigger = self._should_trigger(mw, text)
             
             if should_trigger:
-                triggered_count += 1
-                result = await mw.on_model_stream(current)
+                result = await mw.on_text_stream(current)
                 if result is None:
-                    logger.info(f"Middleware #{i} blocked stream chunk")
                     return None
                 current = result
         
         # Log only every 50 tokens to reduce noise
         if self._token_count % 50 == 0:
-            logger.debug(f"Stream progress: token_count={self._token_count}, middlewares={len(self._middlewares)}")
+            logger.debug(f"Text stream progress: token_count={self._token_count}")
+        
+        return current
+    
+    async def process_thinking_stream(
+        self,
+        chunk: dict[str, Any],
+    ) -> dict[str, Any] | None:
+        """Process thinking streaming chunk through all middlewares."""
+        current = chunk
+        
+        for i, mw in enumerate(self._middlewares):
+            result = await mw.on_thinking_stream(current)
+            if result is None:
+                return None
+            current = result
         
         return current
     
@@ -241,6 +253,50 @@ class MiddlewareChain:
         boundaries = (".", "。", "\n", "!", "?", "！", "？", ";", "；")
         return text.rstrip().endswith(boundaries)
     
+    async def process_text_stream_end(self) -> list[dict[str, Any]]:
+        """Process text stream end through all middlewares.
+        
+        Called when text stream ends, before on_response.
+        Allows middlewares to flush any buffered content.
+        
+        Returns:
+            List of final chunks to emit (may be empty)
+        """
+        final_chunks: list[dict[str, Any]] = []
+        logger.debug(f"Processing text_stream_end through {len(self._middlewares)} middlewares")
+        
+        for i, mw in enumerate(self._middlewares):
+            if hasattr(mw, 'on_text_stream_end'):
+                result = await mw.on_text_stream_end()
+                if result is not None:
+                    logger.debug(f"Middleware #{i} returned final chunk on text_stream_end")
+                    final_chunks.append(result)
+        
+        logger.debug(f"Text stream end processing completed, {len(final_chunks)} final chunks")
+        return final_chunks
+    
+    async def process_thinking_stream_end(self) -> list[dict[str, Any]]:
+        """Process thinking stream end through all middlewares.
+        
+        Called when thinking stream ends.
+        Allows middlewares to flush any buffered thinking content.
+        
+        Returns:
+            List of final thinking chunks to emit (may be empty)
+        """
+        final_chunks: list[dict[str, Any]] = []
+        logger.debug(f"Processing thinking_stream_end through {len(self._middlewares)} middlewares")
+        
+        for i, mw in enumerate(self._middlewares):
+            if hasattr(mw, 'on_thinking_stream_end'):
+                result = await mw.on_thinking_stream_end()
+                if result is not None:
+                    logger.debug(f"Middleware #{i} returned final chunk on thinking_stream_end")
+                    final_chunks.append(result)
+        
+        logger.debug(f"Thinking stream end processing completed, {len(final_chunks)} final chunks")
+        return final_chunks
+    
     def reset_stream_state(self) -> None:
         """Reset streaming state (call at start of new stream)."""
         logger.debug("Resetting stream state")