aury-agent 0.0.4__py3-none-any.whl

aury/agents/memory/store.py

@@ -0,0 +1,187 @@
+"""Memory store protocol and implementations."""
+from __future__ import annotations
+
+import hashlib
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any, Protocol, runtime_checkable
+
+
+@dataclass
+class MemoryEntry:
+    """A memory entry."""
+    id: str
+    content: str
+    session_id: str | None = None
+    invocation_id: str | None = None
+    created_at: datetime = field(default_factory=datetime.now)
+    metadata: dict[str, Any] = field(default_factory=dict)
+    
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "content": self.content,
+            "session_id": self.session_id,
+            "invocation_id": self.invocation_id,
+            "created_at": self.created_at.isoformat(),
+            "metadata": self.metadata,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> MemoryEntry:
+        return cls(
+            id=data["id"],
+            content=data["content"],
+            session_id=data.get("session_id"),
+            invocation_id=data.get("invocation_id"),
+            created_at=datetime.fromisoformat(data["created_at"])
+            if "created_at" in data else datetime.now(),
+            metadata=data.get("metadata", {}),
+        )
+    
+    @property
+    def content_hash(self) -> str:
+        """Get hash of content for deduplication."""
+        return hashlib.sha256(self.content.encode()).hexdigest()[:16]
+
+
+@dataclass
+class ScoredEntry:
+    """Memory entry with relevance score."""
+    entry: MemoryEntry
+    score: float
+    source: str = "default"
+
+
+@runtime_checkable
+class MemoryStore(Protocol):
+    """Memory store protocol."""
+    
+    async def add(self, entry: MemoryEntry) -> str:
+        """Add entry, return ID."""
+        ...
+    
+    async def search(
+        self,
+        query: str,
+        filter: dict[str, Any] | None = None,
+        limit: int = 10,
+    ) -> list[ScoredEntry]:
+        """Search for relevant entries."""
+        ...
+    
+    async def get(self, entry_id: str) -> MemoryEntry | None:
+        """Get entry by ID."""
+        ...
+    
+    async def remove(self, entry_id: str) -> None:
+        """Remove entry."""
+        ...
+    
+    async def revert(
+        self,
+        session_id: str,
+        after_invocation_id: str,
+    ) -> list[str]:
+        """Remove entries after specified invocation.
+        
+        Returns list of deleted IDs.
+        """
+        ...
+
+
+class InMemoryStore:
+    """Simple in-memory store for testing."""
+    
+    def __init__(self) -> None:
+        self._entries: dict[str, MemoryEntry] = {}
+        self._content_hashes: dict[str, str] = {}  # hash -> entry_id
+    
+    async def add(self, entry: MemoryEntry) -> str:
+        """Add entry with deduplication."""
+        content_hash = entry.content_hash
+        
+        # Check for duplicate
+        if content_hash in self._content_hashes:
+            return self._content_hashes[content_hash]
+        
+        self._entries[entry.id] = entry
+        self._content_hashes[content_hash] = entry.id
+        return entry.id
+    
+    async def search(
+        self,
+        query: str,
+        filter: dict[str, Any] | None = None,
+        limit: int = 10,
+    ) -> list[ScoredEntry]:
+        """Simple keyword search."""
+        filter = filter or {}
+        results = []
+        
+        query_lower = query.lower()
+        query_words = set(query_lower.split())
+        
+        for entry in self._entries.values():
+            # Apply filters
+            if filter:
+                skip = False
+                for key, value in filter.items():
+                    entry_value = getattr(entry, key, None) or entry.metadata.get(key)
+                    if entry_value != value:
+                        skip = True
+                        break
+                if skip:
+                    continue
+            
+            # Calculate simple relevance score
+            content_lower = entry.content.lower()
+            content_words = set(content_lower.split())
+            
+            # Word overlap score
+            overlap = len(query_words & content_words)
+            if overlap > 0:
+                score = overlap / len(query_words)
+                results.append(ScoredEntry(entry=entry, score=score, source="memory"))
+            elif query_lower in content_lower:
+                # Substring match
+                results.append(ScoredEntry(entry=entry, score=0.5, source="memory"))
+        
+        # Sort by score and limit
+        results.sort(key=lambda x: x.score, reverse=True)
+        return results[:limit]
+    
+    async def get(self, entry_id: str) -> MemoryEntry | None:
+        return self._entries.get(entry_id)
+    
+    async def remove(self, entry_id: str) -> None:
+        if entry_id in self._entries:
+            entry = self._entries[entry_id]
+            content_hash = entry.content_hash
+            self._content_hashes.pop(content_hash, None)
+            del self._entries[entry_id]
+    
+    async def revert(
+        self,
+        session_id: str,
+        after_invocation_id: str,
+    ) -> list[str]:
+        """Remove entries for invocations after specified one."""
+        # Simple implementation: compare invocation IDs lexicographically
+        # In production, use timestamps or sequence numbers
+        to_delete = []
+        
+        for entry_id, entry in self._entries.items():
+            if entry.session_id == session_id:
+                if entry.invocation_id and entry.invocation_id > after_invocation_id:
+                    to_delete.append(entry_id)
+        
+        for entry_id in to_delete:
+            await self.remove(entry_id)
+        
+        return to_delete
+    
+    def clear(self) -> None:
+        """Clear all entries."""
+        self._entries.clear()
+        self._content_hashes.clear()

aury/agents/memory/types.py

@@ -0,0 +1,137 @@
+"""Memory type definitions.
+
+Core types for the memory system:
+- MemorySummary: Compressed overview of conversation history
+- MemoryRecall: Key points extracted from invocations
+- MemoryContext: Combined context for LLM (summary + recalls)
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Any
+
+
+@dataclass
+class MemorySummary:
+    """Compressed overview of conversation history.
+    
+    Represents the "big picture" of a session's conversation.
+    Updated incrementally as invocations complete.
+    """
+    session_id: str
+    content: str  # Summary text
+    last_invocation_id: str  # Last invocation included in summary
+    updated_at: datetime = field(default_factory=datetime.now)
+    token_count: int = 0  # Estimated token count
+    metadata: dict[str, Any] = field(default_factory=dict)
+    
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "session_id": self.session_id,
+            "content": self.content,
+            "last_invocation_id": self.last_invocation_id,
+            "updated_at": self.updated_at.isoformat(),
+            "token_count": self.token_count,
+            "metadata": self.metadata,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "MemorySummary":
+        return cls(
+            session_id=data["session_id"],
+            content=data["content"],
+            last_invocation_id=data["last_invocation_id"],
+            updated_at=datetime.fromisoformat(data["updated_at"]),
+            token_count=data.get("token_count", 0),
+            metadata=data.get("metadata", {}),
+        )
+
+
+@dataclass
+class MemoryRecall:
+    """Key point extracted from an invocation.
+    
+    Represents an important piece of information that should be
+    recalled when building LLM context. Linked to specific invocation
+    for isolation and revert support.
+    """
+    id: str
+    session_id: str
+    invocation_id: str  # Which invocation this came from
+    content: str  # Recall content
+    importance: float = 0.5  # 0.0 - 1.0, higher = more important
+    tags: list[str] = field(default_factory=list)
+    created_at: datetime = field(default_factory=datetime.now)
+    metadata: dict[str, Any] = field(default_factory=dict)
+    
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "session_id": self.session_id,
+            "invocation_id": self.invocation_id,
+            "content": self.content,
+            "importance": self.importance,
+            "tags": self.tags,
+            "created_at": self.created_at.isoformat(),
+            "metadata": self.metadata,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "MemoryRecall":
+        return cls(
+            id=data["id"],
+            session_id=data["session_id"],
+            invocation_id=data["invocation_id"],
+            content=data["content"],
+            importance=data.get("importance", 0.5),
+            tags=data.get("tags", []),
+            created_at=datetime.fromisoformat(data["created_at"]),
+            metadata=data.get("metadata", {}),
+        )
+
+
+@dataclass
+class MemoryContext:
+    """Combined memory context for LLM.
+    
+    This is what gets injected into the LLM prompt to provide
+    historical context. Contains:
+    - summary: High-level overview of conversation
+    - recalls: Specific important points
+    """
+    summary: MemorySummary | None = None
+    recalls: list[MemoryRecall] = field(default_factory=list)
+    
+    def to_system_message(self) -> str:
+        """Format as system message content for LLM."""
+        parts = []
+        
+        if self.summary and self.summary.content:
+            parts.append(f"## Conversation History Overview\n{self.summary.content}")
+        
+        if self.recalls:
+            recalls_text = "\n".join([
+                f"- [{', '.join(r.tags) if r.tags else 'note'}] {r.content}"
+                for r in sorted(self.recalls, key=lambda x: x.importance, reverse=True)
+            ])
+            parts.append(f"## Key Points\n{recalls_text}")
+        
+        return "\n\n".join(parts) if parts else ""
+    
+    @property
+    def is_empty(self) -> bool:
+        """Check if context has any content."""
+        return (not self.summary or not self.summary.content) and not self.recalls
+    
+    @property
+    def total_recalls(self) -> int:
+        """Get total number of recalls."""
+        return len(self.recalls)
+
+
+__all__ = [
+    "MemorySummary",
+    "MemoryRecall", 
+    "MemoryContext",
+]

aury/agents/messages/__init__.py

@@ -0,0 +1,40 @@
+"""Message types for conversation history.
+
+Note: Message persistence is now handled by:
+- MessageBackend (backends/message/): Storage layer
+- MessageBackendMiddleware (middleware/message.py): Save via on_message_save hook
+- MessageContextProvider (context_providers/message.py): Fetch for context
+
+This module provides message types used across the system.
+"""
+from .types import (
+    MessageRole,
+    Message,
+)
+from .store import (
+    MessageStore,
+    InMemoryMessageStore,
+)
+from .raw_store import (
+    RawMessageStore,
+    StateBackendRawMessageStore,
+    InMemoryRawMessageStore,
+)
+from .config import (
+    MessageConfig,
+)
+
+__all__ = [
+    # Types
+    "MessageRole",
+    "Message",
+    # Store (protocol + in-memory for testing)
+    "MessageStore",
+    "InMemoryMessageStore",
+    # Raw Store
+    "RawMessageStore",
+    "StateBackendRawMessageStore",
+    "InMemoryRawMessageStore",
+    # Config
+    "MessageConfig",
+]

aury/agents/messages/config.py

@@ -0,0 +1,47 @@
+"""Message storage configuration.
+
+Configures how messages are stored and retrieved.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+
+@dataclass
+class MessageConfig:
+    """Configuration for message storage.
+    
+    Controls raw message storage and recall behavior.
+    
+    Attributes:
+        enable_raw_store: Whether to store complete messages in RawMessageStore.
+                         Required for HITL recovery and full-context recall.
+        persist_raw: Whether to keep raw messages after invocation completes.
+                    False = clean up after invocation (default, saves space)
+                    True = keep forever (for audit/recall)
+        recall_mode: How to build context for LLM recall.
+                    "mixed" = previous invocations truncated + current invocation raw
+                    "raw" = all raw messages (requires persist_raw=True or current inv only)
+    
+    Example:
+        # Default: raw for recovery, clean up after
+        config = MessageConfig()
+        
+        # Keep all raw messages for full recall
+        config = MessageConfig(persist_raw=True, recall_mode="raw")
+        
+        # Disable raw storage (no HITL recovery)
+        config = MessageConfig(enable_raw_store=False)
+    """
+    enable_raw_store: bool = True
+    persist_raw: bool = False
+    recall_mode: Literal["mixed", "raw"] = "mixed"
+    
+    def __post_init__(self):
+        # Validate: raw recall mode requires raw storage
+        if self.recall_mode == "raw" and not self.enable_raw_store:
+            raise ValueError("recall_mode='raw' requires enable_raw_store=True")
+
+
+__all__ = ["MessageConfig"]

aury/agents/messages/raw_store.py

@@ -0,0 +1,224 @@
+"""Raw message store for complete message storage.
+
+RawMessageStore stores complete, untruncated messages for:
+- HITL recovery (restore exact state)
+- Full-context recall (when truncated history is insufficient)
+- Audit/debugging
+
+Messages are stored per invocation and can be cleaned up after invocation completes.
+"""
+from __future__ import annotations
+
+from typing import Any, Protocol, runtime_checkable
+from datetime import datetime
+
+from ..core.types.session import generate_id
+
+
+@runtime_checkable
+class RawMessageStore(Protocol):
+    """Protocol for raw message storage.
+    
+    Stores complete messages per invocation.
+    """
+    
+    async def add(
+        self,
+        invocation_id: str,
+        message: dict[str, Any],
+    ) -> str:
+        """Add a message and return its ID.
+        
+        Args:
+            invocation_id: Invocation this message belongs to
+            message: Complete message dict
+            
+        Returns:
+            Generated message ID
+        """
+        ...
+    
+    async def get(self, msg_id: str) -> dict[str, Any] | None:
+        """Get a single message by ID.
+        
+        Args:
+            msg_id: Message ID
+            
+        Returns:
+            Message dict or None if not found
+        """
+        ...
+    
+    async def get_many(self, msg_ids: list[str]) -> list[dict[str, Any]]:
+        """Get multiple messages by IDs.
+        
+        Args:
+            msg_ids: List of message IDs
+            
+        Returns:
+            List of message dicts (in same order, None for missing)
+        """
+        ...
+    
+    async def get_by_invocation(self, invocation_id: str) -> list[dict[str, Any]]:
+        """Get all messages for an invocation.
+        
+        Args:
+            invocation_id: Invocation ID
+            
+        Returns:
+            List of messages in chronological order
+        """
+        ...
+    
+    async def delete_by_invocation(self, invocation_id: str) -> int:
+        """Delete all messages for an invocation.
+        
+        Args:
+            invocation_id: Invocation ID
+            
+        Returns:
+            Number of messages deleted
+        """
+        ...
+
+
+class StateBackendRawMessageStore:
+    """RawMessageStore implementation backed by StateBackend."""
+    
+    def __init__(self, backend: Any):  # StateBackend
+        """Initialize with StateBackend.
+        
+        Args:
+            backend: StateBackend instance
+        """
+        self._backend = backend
+    
+    async def add(
+        self,
+        invocation_id: str,
+        message: dict[str, Any],
+    ) -> str:
+        """Add a message using StateBackend."""
+        msg_id = generate_id("rmsg")
+        
+        # Store message with metadata
+        await self._backend.set("raw_messages", msg_id, {
+            "id": msg_id,
+            "invocation_id": invocation_id,
+            "message": message,
+            "created_at": datetime.now().isoformat(),
+        })
+        
+        # Add to invocation's message list
+        inv_key = f"raw_msg_ids:{invocation_id}"
+        msg_ids = await self._backend.get("raw_messages", inv_key) or []
+        msg_ids.append(msg_id)
+        await self._backend.set("raw_messages", inv_key, msg_ids)
+        
+        return msg_id
+    
+    async def get(self, msg_id: str) -> dict[str, Any] | None:
+        """Get a single message."""
+        data = await self._backend.get("raw_messages", msg_id)
+        if data:
+            return data.get("message")
+        return None
+    
+    async def get_many(self, msg_ids: list[str]) -> list[dict[str, Any]]:
+        """Get multiple messages."""
+        results = []
+        for msg_id in msg_ids:
+            msg = await self.get(msg_id)
+            if msg:
+                results.append(msg)
+        return results
+    
+    async def get_by_invocation(self, invocation_id: str) -> list[dict[str, Any]]:
+        """Get all messages for an invocation."""
+        inv_key = f"raw_msg_ids:{invocation_id}"
+        msg_ids = await self._backend.get("raw_messages", inv_key) or []
+        return await self.get_many(msg_ids)
+    
+    async def delete_by_invocation(self, invocation_id: str) -> int:
+        """Delete all messages for an invocation."""
+        inv_key = f"raw_msg_ids:{invocation_id}"
+        msg_ids = await self._backend.get("raw_messages", inv_key) or []
+        
+        # Delete each message
+        for msg_id in msg_ids:
+            await self._backend.remove("raw_messages", msg_id)
+        
+        # Delete the index
+        await self._backend.remove("raw_messages", inv_key)
+        
+        return len(msg_ids)
+
+
+class InMemoryRawMessageStore:
+    """In-memory raw message store for testing."""
+    
+    def __init__(self) -> None:
+        self._messages: dict[str, dict[str, Any]] = {}  # msg_id -> message data
+        self._invocation_index: dict[str, list[str]] = {}  # inv_id -> [msg_ids]
+    
+    async def add(
+        self,
+        invocation_id: str,
+        message: dict[str, Any],
+    ) -> str:
+        """Add a message."""
+        msg_id = generate_id("rmsg")
+        
+        self._messages[msg_id] = {
+            "id": msg_id,
+            "invocation_id": invocation_id,
+            "message": message,
+            "created_at": datetime.now().isoformat(),
+        }
+        
+        if invocation_id not in self._invocation_index:
+            self._invocation_index[invocation_id] = []
+        self._invocation_index[invocation_id].append(msg_id)
+        
+        return msg_id
+    
+    async def get(self, msg_id: str) -> dict[str, Any] | None:
+        """Get a single message."""
+        data = self._messages.get(msg_id)
+        if data:
+            return data.get("message")
+        return None
+    
+    async def get_many(self, msg_ids: list[str]) -> list[dict[str, Any]]:
+        """Get multiple messages."""
+        results = []
+        for msg_id in msg_ids:
+            msg = await self.get(msg_id)
+            if msg:
+                results.append(msg)
+        return results
+    
+    async def get_by_invocation(self, invocation_id: str) -> list[dict[str, Any]]:
+        """Get all messages for an invocation."""
+        msg_ids = self._invocation_index.get(invocation_id, [])
+        return await self.get_many(msg_ids)
+    
+    async def delete_by_invocation(self, invocation_id: str) -> int:
+        """Delete all messages for an invocation."""
+        msg_ids = self._invocation_index.pop(invocation_id, [])
+        for msg_id in msg_ids:
+            self._messages.pop(msg_id, None)
+        return len(msg_ids)
+    
+    def clear(self) -> None:
+        """Clear all messages (for testing)."""
+        self._messages.clear()
+        self._invocation_index.clear()
+
+
+__all__ = [
+    "RawMessageStore",
+    "StateBackendRawMessageStore",
+    "InMemoryRawMessageStore",
+]