aury-agent 0.0.4__py3-none-any.whl

aury/agents/context_providers/artifact.py

@@ -0,0 +1,299 @@
+"""ArtifactContextProvider - provides artifacts context and tools.
+
+Features:
+- Artifact index in system_content
+- ReadArtifactTool for LLM to read artifact content
+- Pluggable Loader system for different artifact types
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+
+from .base import BaseContextProvider, AgentContext
+from ..core.types.tool import BaseTool, ToolResult, ToolContext
+
+if TYPE_CHECKING:
+    from ..core.context import InvocationContext
+    from ..backends.artifact import ArtifactBackend, StoredArtifact
+
+
+# ============================================================
+# Loader Protocol and Registry
+# ============================================================
+
+@runtime_checkable
+class ArtifactLoader(Protocol):
+    """Protocol for loading artifact content.
+    
+    Implement this to support custom artifact types.
+    """
+    
+    async def load(self, artifact: "StoredArtifact") -> str:
+        """Load artifact content.
+        
+        Args:
+            artifact: Stored artifact with metadata and data
+            
+        Returns:
+            String content for LLM context
+        """
+        ...
+
+
+class DefaultLoader:
+    """Default loader - returns summary or data as string."""
+    
+    async def load(self, artifact: "StoredArtifact") -> str:
+        # Try summary first
+        if artifact.summary:
+            return artifact.summary
+        
+        # Try data
+        if artifact.data:
+            import json
+            return json.dumps(artifact.data, ensure_ascii=False, indent=2)
+        
+        return "No content"
+
+
+class SearchResultsLoader:
+    """Loader for search_results artifacts."""
+    
+    async def load(self, artifact: "StoredArtifact") -> str:
+        data = artifact.data or {}
+        query = data.get("query", "")
+        results = data.get("results", [])
+        
+        if not results:
+            return f"Search `{query}` returned no results"
+        
+        lines = [f"Search: {query}", f"Total {len(results)} results:", ""]
+        for i, r in enumerate(results, 1):
+            lines.append(f"{i}. {r.get('title', 'Untitled')}")
+            snippet = r.get('snippet', '')[:150]
+            if snippet:
+                lines.append(f"   {snippet}")
+            url = r.get('url', '')
+            if url:
+                lines.append(f"   URL: {url}")
+            lines.append("")
+        
+        return "\n".join(lines)
+
+
+class FileLoader:
+    """Loader for file-based artifacts (file:// URLs)."""
+    
+    async def load(self, artifact: "StoredArtifact") -> str:
+        data = artifact.data or {}
+        url = data.get("url")
+        
+        if not url:
+            return artifact.summary or "No content"
+        
+        # Load from file:// URL
+        if url.startswith("file://"):
+            file_path = url[7:]  # Remove "file://"
+            try:
+                import aiofiles
+                async with aiofiles.open(file_path, "r", encoding="utf-8") as f:
+                    return await f.read()
+            except ImportError:
+                # Fallback to sync read
+                with open(file_path, "r", encoding="utf-8") as f:
+                    return f.read()
+            except Exception as e:
+                return artifact.summary or f"Failed to read file: {file_path} ({e})"
+        
+        return artifact.summary or f"Unsupported URL scheme: {url}"
+
+
+# Global loader registry
+_LOADERS: dict[str, ArtifactLoader] = {
+    "search_results": SearchResultsLoader(),
+    "report": FileLoader(),
+    "file": FileLoader(),
+}
+
+_DEFAULT_LOADER = DefaultLoader()
+
+
+def register_loader(kind: str, loader: ArtifactLoader) -> None:
+    """Register a custom loader for artifact kind.
+    
+    Args:
+        kind: Artifact kind (e.g., "pdf", "image")
+        loader: Loader instance implementing ArtifactLoader protocol
+        
+    Example:
+        class PDFLoader:
+            async def load(self, artifact):
+                # Custom PDF loading logic
+                return extracted_text
+        
+        register_loader("pdf", PDFLoader())
+    """
+    _LOADERS[kind] = loader
+
+
+def get_loader(kind: str) -> ArtifactLoader:
+    """Get loader for artifact kind.
+    
+    Returns registered loader or default loader if not found.
+    """
+    return _LOADERS.get(kind, _DEFAULT_LOADER)
+
+
+# ============================================================
+# ReadArtifactTool
+# ============================================================
+
+class ReadArtifactTool(BaseTool):
+    """Tool for reading artifact content.
+    
+    Uses registered loaders to load content based on artifact kind.
+    """
+    
+    _name = "read_artifact"
+    _description = """Read the full content of an artifact.
+
+Use this to get detailed content of artifacts like search results, reports, etc.
+"""
+    
+    _parameters: dict[str, Any] = {
+        "type": "object",
+        "properties": {
+            "artifact_id": {
+                "type": "string",
+                "description": "Artifact ID to read",
+            },
+        },
+        "required": ["artifact_id"],
+    }
+    
+    def __init__(self, backend: "ArtifactBackend") -> None:
+        """Initialize with artifact backend.
+        
+        Args:
+            backend: Backend for artifact storage/retrieval
+        """
+        self._backend = backend
+    
+    async def execute(self, params: dict[str, Any], ctx: ToolContext) -> ToolResult:
+        artifact_id = params.get("artifact_id")
+        if not artifact_id:
+            return ToolResult.error("artifact_id is required")
+        
+        # Fetch artifact
+        stored = await self._backend.get(artifact_id)
+        if not stored:
+            return ToolResult.error(f"Artifact not found: {artifact_id}")
+        
+        # Load content using appropriate loader
+        loader = get_loader(stored.kind)
+        try:
+            content = await loader.load(stored)
+        except Exception as e:
+            return ToolResult.error(f"Failed to load artifact: {e}")
+        
+        return ToolResult.success(output=content)
+
+
+# ============================================================
+# ArtifactContextProvider
+# ============================================================
+
+class ArtifactContextProvider(BaseContextProvider):
+    """Artifact context provider.
+    
+    Provides:
+    - system_content: Artifact index for LLM awareness
+    - tools: ReadArtifactTool for reading artifact content
+    
+    Example:
+        provider = ArtifactContextProvider(
+            backend=my_artifact_backend,
+            max_summary_items=10,
+        )
+        
+        # Register custom loader
+        register_loader("pdf", MyPDFLoader())
+    """
+    
+    _name = "artifacts"
+    
+    def __init__(
+        self,
+        backend: "ArtifactBackend",
+        max_summary_items: int = 10,
+        enable_tools: bool = True,
+    ):
+        """Initialize ArtifactContextProvider.
+        
+        Args:
+            backend: Artifact storage backend
+            max_summary_items: Max artifacts to show in summary
+            enable_tools: Whether to provide ReadArtifactTool
+        """
+        self.backend = backend
+        self.max_summary_items = max_summary_items
+        self.enable_tools = enable_tools
+    
+    async def fetch(self, ctx: "InvocationContext") -> AgentContext:
+        """Fetch artifact context.
+        
+        Returns:
+            AgentContext with system_content (index) and ReadArtifactTool
+        """
+        # Get artifacts for session
+        artifacts = await self.backend.list(session_id=ctx.session.id)
+        
+        if not artifacts:
+            # No artifacts, but still provide tool if enabled
+            tools: list[BaseTool] = []
+            if self.enable_tools:
+                tools = [ReadArtifactTool(self.backend)]
+            return AgentContext(tools=tools) if tools else AgentContext.empty()
+        
+        # Build summary (artifact index)
+        summary = self._build_summary(artifacts[:self.max_summary_items])
+        
+        # Build tools
+        tools = []
+        if self.enable_tools:
+            tools = [ReadArtifactTool(self.backend)]
+        
+        return AgentContext(
+            system_content=summary,
+            tools=tools,
+        )
+    
+    def _build_summary(self, artifacts: list[Any]) -> str:
+        """Build artifact summary for LLM context."""
+        lines = ["## Available Artifacts", ""]
+        for a in artifacts:
+            title = getattr(a, 'title', 'Untitled')
+            summary = getattr(a, 'summary', None)
+            artifact_id = getattr(a, 'id', 'unknown')
+            kind = getattr(a, 'kind', 'unknown')
+            
+            line = f"- [{artifact_id}] ({kind}) {title}"
+            if summary:
+                line += f": {summary[:100]}"
+            lines.append(line)
+        
+        lines.append("")
+        lines.append("Use `read_artifact` tool to get full content.")
+        return "\n".join(lines)
+
+
+__all__ = [
+    "ArtifactContextProvider",
+    "ArtifactLoader",
+    "ReadArtifactTool",
+    "register_loader",
+    "get_loader",
+    "DefaultLoader",
+    "SearchResultsLoader",
+    "FileLoader",
+]

aury/agents/context_providers/base.py

@@ -0,0 +1,177 @@
+"""Base ContextProvider protocol and AgentContext.
+
+ContextProvider is the unified abstraction for providing LLM context.
+All context sources (memory, artifacts, skills, subagents)
+are ContextProviders that implement fetch(ctx) -> AgentContext.
+
+Design principle:
+- Providers only provide DATA (content, subagents list, skills list, etc.)
+- Agent consumes this data and decides how to use it (e.g., create DelegateTool)
+"""
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Protocol, runtime_checkable
+
+if TYPE_CHECKING:
+    from ..core.context import InvocationContext
+    from ..core.types.tool import BaseTool
+    from ..backends.subagent import AgentConfig
+    from ..skill import Skill
+
+
+@dataclass
+class AgentContext:
+    """Context provided by ContextProvider.fetch().
+    
+    Contains all context that a ContextProvider provides:
+    - system_content: Injected into system prompt
+    - user_content: Injected into user message (before User: input)
+    - tools: Tools to register (from MemoryContextProvider, ArtifactContextProvider, etc.)
+    - messages: History messages (from MessageContextProvider)
+    - subagents: Sub-agent configs (from SubAgentContextProvider) - Agent creates DelegateTool
+    - skills: Skill definitions (from SkillContextProvider)
+    
+    Multiple AgentContexts are merged by the Agent before LLM call.
+    """
+    # Content injection
+    system_content: str | None = None
+    user_content: str | None = None
+    
+    # Tools to register (paired tools from providers like MemorySearchTool)
+    tools: list["BaseTool"] = field(default_factory=list)
+    
+    # Messages (from MessageContextProvider)
+    messages: list[dict[str, Any]] = field(default_factory=list)
+    
+    # SubAgents (from SubAgentContextProvider) - Agent creates DelegateTool from this
+    subagents: list["AgentConfig"] = field(default_factory=list)
+    
+    # Skills (from SkillContextProvider)
+    skills: list["Skill"] = field(default_factory=list)
+    
+    @staticmethod
+    def empty() -> "AgentContext":
+        """Create an empty AgentContext."""
+        return AgentContext()
+    
+    @staticmethod
+    def merge(outputs: list["AgentContext"]) -> "AgentContext":
+        """Merge multiple AgentContexts into one.
+        
+        - system_content: Concatenated with newlines
+        - user_content: Concatenated with newlines
+        - tools: Combined list (deduplicated by tool.name)
+        - messages: Combined list
+        - subagents: Combined list (deduplicated by key)
+        - skills: Combined list (deduplicated by name)
+        """
+        system_parts: list[str] = []
+        user_parts: list[str] = []
+        all_tools: list["BaseTool"] = []
+        all_messages: list[dict[str, Any]] = []
+        all_subagents: list["AgentConfig"] = []
+        all_skills: list["Skill"] = []
+        seen_tool_names: set[str] = set()
+        seen_agent_keys: set[str] = set()
+        seen_skill_names: set[str] = set()
+        
+        for output in outputs:
+            if output.system_content:
+                system_parts.append(output.system_content)
+            if output.user_content:
+                user_parts.append(output.user_content)
+            
+            # Deduplicate tools by name
+            for tool in output.tools:
+                if tool.name not in seen_tool_names:
+                    seen_tool_names.add(tool.name)
+                    all_tools.append(tool)
+            
+            all_messages.extend(output.messages)
+            
+            # Deduplicate subagents by key
+            for agent in output.subagents:
+                if agent.key not in seen_agent_keys:
+                    seen_agent_keys.add(agent.key)
+                    all_subagents.append(agent)
+            
+            # Deduplicate skills by name
+            for skill in output.skills:
+                if skill.name not in seen_skill_names:
+                    seen_skill_names.add(skill.name)
+                    all_skills.append(skill)
+        
+        return AgentContext(
+            system_content="\n\n".join(system_parts) if system_parts else None,
+            user_content="\n\n".join(user_parts) if user_parts else None,
+            tools=all_tools,
+            messages=all_messages,
+            subagents=all_subagents,
+            skills=all_skills,
+        )
+
+
+@runtime_checkable
+class ContextProvider(Protocol):
+    """Protocol for all ContextProviders.
+    
+    ContextProviders provide context for LLM calls. Each ContextProvider:
+    1. Has a unique name
+    2. Implements fetch(ctx) to return AgentContext
+    3. Provides DATA only - Agent decides how to consume it
+    
+    Examples:
+        - MessageContextProvider: Returns messages list
+        - MemoryContextProvider: Returns system_content (summary) + MemorySearchTool
+        - ArtifactContextProvider: Returns system_content (index) + ReadArtifactTool
+        - SubAgentContextProvider: Returns subagents list (Agent creates DelegateTool)
+        - SkillContextProvider: Returns skills list
+    """
+    
+    @property
+    def name(self) -> str:
+        """Unique name for this provider."""
+        ...
+    
+    async def fetch(self, ctx: "InvocationContext") -> AgentContext:
+        """Fetch context for LLM call.
+        
+        Called before each LLM call. Can return dynamic content
+        based on current context (state, session, ctx.input, etc.).
+        
+        Args:
+            ctx: Current invocation context (includes ctx.input with runtime vars)
+            
+        Returns:
+            AgentContext with content and tools
+        """
+        ...
+
+
+class BaseContextProvider(ABC):
+    """Abstract base class for ContextProviders.
+    
+    Provides default implementation structure. Subclass and implement
+    fetch() to create custom providers.
+    """
+    
+    _name: str = "base"
+    
+    @property
+    def name(self) -> str:
+        """Provider name."""
+        return self._name
+    
+    @abstractmethod
+    async def fetch(self, ctx: "InvocationContext") -> AgentContext:
+        """Fetch context. Subclasses must implement."""
+        ...
+
+
+__all__ = [
+    "ContextProvider",
+    "AgentContext",
+    "BaseContextProvider",
+]

aury/agents/context_providers/memory.py

@@ -0,0 +1,70 @@
+"""MemoryContextProvider - provides memory context.
+
+This provider ONLY fetches memory context (summary/recalls).
+Memory writing is handled by Middleware (e.g., in on_message_save hook).
+"""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from .base import BaseContextProvider, AgentContext
+
+if TYPE_CHECKING:
+    from ..core.context import InvocationContext
+    from ..memory import MemoryManager
+
+
+class MemoryContextProvider(BaseContextProvider):
+    """Memory context provider.
+    
+    Provides:
+    - system_content: Memory summary and recalls for LLM awareness
+    
+    Note: Memory WRITING is not done here.
+    Use Middleware (on_message_save hook) to write to MemoryManager.
+    
+    Example:
+        provider = MemoryContextProvider(
+            memory_manager=my_memory_manager,
+            recall_limit=10,
+        )
+    """
+    
+    _name = "memory"
+    
+    def __init__(
+        self,
+        memory_manager: "MemoryManager",
+        recall_limit: int = 10,
+    ):
+        """Initialize MemoryContextProvider.
+        
+        Args:
+            memory_manager: The underlying MemoryManager
+            recall_limit: Max recalls to include in context
+        """
+        self.memory = memory_manager
+        self.recall_limit = recall_limit
+    
+    async def fetch(self, ctx: "InvocationContext") -> AgentContext:
+        """Fetch memory context.
+        
+        Returns:
+            AgentContext with system_content (memory summary/recalls)
+        """
+        # Get memory context
+        memory_ctx = await self.memory.get_context(
+            session_id=ctx.session.id,
+            recall_limit=self.recall_limit,
+        )
+        
+        # Build system content
+        if not memory_ctx or memory_ctx.is_empty:
+            return AgentContext.empty()
+        
+        return AgentContext(
+            system_content=memory_ctx.to_system_message(),
+        )
+
+
+__all__ = ["MemoryContextProvider"]

aury/agents/context_providers/message.py

@@ -0,0 +1,130 @@
+"""MessageContextProvider - provides conversation history messages.
+
+This provider ONLY fetches message history for context.
+Message saving is handled by Middleware (on_message_save hook).
+
+Recovery Strategy:
+- Check State for complete messages (from pending/crashed invocation)
+- If found, use State messages (complete, not truncated)
+- Otherwise, load from MessageBackend (truncated historical messages)
+"""
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, Any
+
+from .base import BaseContextProvider, AgentContext
+
+if TYPE_CHECKING:
+    from ..core.context import InvocationContext
+    from ..backends import MessageBackend
+
+logger = logging.getLogger(__name__)
+
+
+class MessageContextProvider(BaseContextProvider):
+    """Message history context provider.
+    
+    Provides conversation history messages for LLM context.
+    Uses ctx.backends.message directly.
+    
+    Features:
+    - Load messages via MessageBackend
+    - Priority: State (complete) > MessageBackend (truncated)
+    - Turn limits handled by backend
+    
+    Note: Message SAVING is not done here.
+    Use MessageBackendMiddleware for saving messages.
+    """
+    
+    _name = "messages"
+    
+    def __init__(self, *, max_messages: int = 100):
+        """Initialize MessageContextProvider.
+        
+        Args:
+            max_messages: Max messages to fetch
+        """
+        self.max_messages = max_messages
+    
+    async def fetch(self, ctx: "InvocationContext") -> AgentContext:
+        """Fetch conversation history.
+        
+        Priority:
+        1. Check State for complete messages (from pending/crashed inv)
+        2. Fall back to MessageBackend (truncated historical messages)
+        
+        Returns:
+            AgentContext with messages list
+        """
+        # Try to get complete messages from State (pending/crashed recovery)
+        state_messages = await self._get_messages_from_state(ctx)
+        if state_messages:
+            logger.debug(
+                f"Loaded {len(state_messages)} messages from State (complete)",
+                extra={"session_id": ctx.session.id},
+            )
+            return AgentContext(messages=state_messages)
+        
+        # Fall back to MessageBackend
+        messages = await self._fetch_from_backend(ctx)
+        logger.debug(
+            f"Loaded {len(messages)} messages from backend (may be truncated)",
+            extra={"session_id": ctx.session.id},
+        )
+        return AgentContext(messages=messages)
+    
+    async def _fetch_from_backend(self, ctx: "InvocationContext") -> list[dict[str, Any]]:
+        """Fetch messages from MessageBackend."""
+        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
+            return [
+                {"role": m["role"], "content": m["content"]}
+                for m in messages
+            ]
+        
+        # No backend available
+        logger.warning(
+            "No message backend available for MessageContextProvider",
+            extra={"session_id": ctx.session.id},
+        )
+        return []
+    
+    async def _get_messages_from_state(self, ctx: "InvocationContext") -> list[dict[str, Any]] | None:
+        """Try to get complete messages from State.
+        
+        State stores complete (not truncated) messages during execution.
+        This allows recovery from:
+        - HITL suspend
+        - Process crash
+        - Abnormal termination
+        
+        Returns:
+            List of message dicts if found, None otherwise
+        """
+        if not ctx.state:
+            return None
+        
+        # Check if State has message_history
+        messages_data = ctx.state.get("agent.message_history")
+        if not messages_data:
+            return None
+        
+        if not isinstance(messages_data, list):
+            return None
+        
+        # Validate and convert to expected format
+        messages: list[dict[str, Any]] = []
+        for msg in messages_data:
+            if isinstance(msg, dict) and "role" in msg:
+                messages.append(msg)
+        
+        return messages if messages else None
+
+
+__all__ = ["MessageContextProvider"]