agent-runtime-core 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

agent_runtime_core/persistence/manager.py

@@ -8,6 +8,16 @@ For Django integration, you can either:
 1. Pass pre-instantiated store instances
 2. Pass store classes with appropriate kwargs
 3. Use factory functions for request-scoped stores
+
+Core stores (always available):
+- MemoryStore: Key-value storage
+- ConversationStore: Conversation history
+- TaskStore: Task lists and progress
+- PreferencesStore: User/agent configuration
+
+Optional stores (must be explicitly configured):
+- KnowledgeStore: Facts, summaries, embeddings
+- AuditStore: Logs, errors, metrics
 """
 
 from dataclasses import dataclass, field
@@ -19,6 +29,8 @@ from agent_runtime_core.persistence.base import (
     ConversationStore,
     TaskStore,
     PreferencesStore,
+    KnowledgeStore,
+    AuditStore,
     Scope,
 )
 from agent_runtime_core.persistence.file import (
@@ -34,6 +46,8 @@ MemoryStoreFactory = Callable[[], MemoryStore]
 ConversationStoreFactory = Callable[[], ConversationStore]
 TaskStoreFactory = Callable[[], TaskStore]
 PreferencesStoreFactory = Callable[[], PreferencesStore]
+KnowledgeStoreFactory = Callable[[], KnowledgeStore]
+AuditStoreFactory = Callable[[], AuditStore]
 
 
 @dataclass
@@ -46,6 +60,9 @@ class PersistenceConfig:
     - A pre-instantiated store instance
     - A factory function that returns a store instance
 
+    Core stores (memory, conversations, tasks, preferences) have file-based
+    defaults. Optional stores (knowledge, audit) must be explicitly configured.
+
     Example for Django:
         from myapp.stores import DjangoMemoryStore, DjangoConversationStore
 
@@ -65,31 +82,47 @@ class PersistenceConfig:
         config = PersistenceConfig(
             memory_store_factory=lambda: DjangoMemoryStore(user=get_current_user()),
         )
+
+        # Option 4: Enable optional stores
+        config = PersistenceConfig(
+            knowledge_store=DjangoKnowledgeStore(user=request.user),
+            audit_store=DjangoAuditStore(user=request.user),
+        )
     """
 
-    # Backend classes (can be swapped for custom implementations)
+    # Backend classes for core stores (can be swapped for custom implementations)
     memory_store_class: Type[MemoryStore] = FileMemoryStore
     conversation_store_class: Type[ConversationStore] = FileConversationStore
     task_store_class: Type[TaskStore] = FileTaskStore
     preferences_store_class: Type[PreferencesStore] = FilePreferencesStore
 
+    # Backend classes for optional stores (no defaults - must be explicitly set)
+    knowledge_store_class: Optional[Type[KnowledgeStore]] = None
+    audit_store_class: Optional[Type[AuditStore]] = None
+
     # Pre-instantiated store instances (takes precedence over classes)
     memory_store: Optional[MemoryStore] = None
     conversation_store: Optional[ConversationStore] = None
     task_store: Optional[TaskStore] = None
     preferences_store: Optional[PreferencesStore] = None
+    knowledge_store: Optional[KnowledgeStore] = None
+    audit_store: Optional[AuditStore] = None
 
     # Factory functions (takes precedence over classes, but not instances)
     memory_store_factory: Optional[MemoryStoreFactory] = None
     conversation_store_factory: Optional[ConversationStoreFactory] = None
     task_store_factory: Optional[TaskStoreFactory] = None
     preferences_store_factory: Optional[PreferencesStoreFactory] = None
+    knowledge_store_factory: Optional[KnowledgeStoreFactory] = None
+    audit_store_factory: Optional[AuditStoreFactory] = None
 
     # Kwargs passed to store class constructors (only used with classes)
     memory_store_kwargs: dict = field(default_factory=dict)
     conversation_store_kwargs: dict = field(default_factory=dict)
     task_store_kwargs: dict = field(default_factory=dict)
     preferences_store_kwargs: dict = field(default_factory=dict)
+    knowledge_store_kwargs: dict = field(default_factory=dict)
+    audit_store_kwargs: dict = field(default_factory=dict)
 
     # Project directory (convenience for file-based stores)
     # Only used if store_kwargs doesn't already have project_dir
@@ -100,8 +133,11 @@ class PersistenceManager:
     """
     Unified manager for all persistence stores.
 
-    Provides access to memory, conversations, tasks, and preferences
-    with pluggable backends.
+    Provides access to core stores (memory, conversations, tasks, preferences)
+    and optional stores (knowledge, audit) with pluggable backends.
+
+    Core stores have file-based defaults. Optional stores return None
+    unless explicitly configured.
 
     Example:
         # Use default file-based storage
@@ -120,8 +156,15 @@ class PersistenceManager:
         config = PersistenceConfig(
             memory_store=DjangoMemoryStore(user=request.user),
             conversation_store=DjangoConversationStore(user=request.user),
+            # Enable optional stores
+            knowledge_store=DjangoKnowledgeStore(user=request.user),
+            audit_store=DjangoAuditStore(user=request.user),
         )
         manager = PersistenceManager(config)
+
+        # Check if optional stores are available
+        if manager.knowledge:
+            await manager.knowledge.save_fact(fact)
     """
 
     def __init__(self, config: Optional[PersistenceConfig] = None):
@@ -130,6 +173,11 @@ class PersistenceManager:
         self._conversations: Optional[ConversationStore] = None
         self._tasks: Optional[TaskStore] = None
         self._preferences: Optional[PreferencesStore] = None
+        self._knowledge: Optional[KnowledgeStore] = None
+        self._audit: Optional[AuditStore] = None
+        # Track if optional stores have been initialized
+        self._knowledge_initialized = False
+        self._audit_initialized = False
 
     def _build_kwargs(self, store_kwargs: dict) -> dict:
         """Build kwargs for store instantiation."""
@@ -193,6 +241,54 @@ class PersistenceManager:
                 self._preferences = self._config.preferences_store_class(**kwargs)
         return self._preferences
 
+    @property
+    def knowledge(self) -> Optional[KnowledgeStore]:
+        """
+        Get the knowledge store (optional).
+
+        Returns None if not configured. Check before using:
+            if manager.knowledge:
+                await manager.knowledge.save_fact(fact)
+        """
+        if not self._knowledge_initialized:
+            self._knowledge_initialized = True
+            if self._config.knowledge_store is not None:
+                self._knowledge = self._config.knowledge_store
+            elif self._config.knowledge_store_factory is not None:
+                self._knowledge = self._config.knowledge_store_factory()
+            elif self._config.knowledge_store_class is not None:
+                kwargs = self._build_kwargs(self._config.knowledge_store_kwargs)
+                self._knowledge = self._config.knowledge_store_class(**kwargs)
+        return self._knowledge
+
+    @property
+    def audit(self) -> Optional[AuditStore]:
+        """
+        Get the audit store (optional).
+
+        Returns None if not configured. Check before using:
+            if manager.audit:
+                await manager.audit.log_event(entry)
+        """
+        if not self._audit_initialized:
+            self._audit_initialized = True
+            if self._config.audit_store is not None:
+                self._audit = self._config.audit_store
+            elif self._config.audit_store_factory is not None:
+                self._audit = self._config.audit_store_factory()
+            elif self._config.audit_store_class is not None:
+                kwargs = self._build_kwargs(self._config.audit_store_kwargs)
+                self._audit = self._config.audit_store_class(**kwargs)
+        return self._audit
+
+    def has_knowledge(self) -> bool:
+        """Check if knowledge store is configured."""
+        return self.knowledge is not None
+
+    def has_audit(self) -> bool:
+        """Check if audit store is configured."""
+        return self.audit is not None
+
     async def close(self) -> None:
         """Close all stores."""
         if self._memory:
@@ -203,6 +299,10 @@ class PersistenceManager:
             await self._tasks.close()
         if self._preferences:
             await self._preferences.close()
+        if self._knowledge:
+            await self._knowledge.close()
+        if self._audit:
+            await self._audit.close()
 
 
 # Global manager instance
@@ -215,27 +315,31 @@ def configure_persistence(
     conversation_store_class: Optional[Type[ConversationStore]] = None,
     task_store_class: Optional[Type[TaskStore]] = None,
     preferences_store_class: Optional[Type[PreferencesStore]] = None,
+    knowledge_store_class: Optional[Type[KnowledgeStore]] = None,
+    audit_store_class: Optional[Type[AuditStore]] = None,
     project_dir: Optional[Path] = None,
     **kwargs,
 ) -> PersistenceConfig:
     """
     Configure the global persistence manager.
-    
+
     Args:
         memory_store_class: Custom memory store implementation
         conversation_store_class: Custom conversation store implementation
         task_store_class: Custom task store implementation
         preferences_store_class: Custom preferences store implementation
+        knowledge_store_class: Custom knowledge store implementation (optional)
+        audit_store_class: Custom audit store implementation (optional)
         project_dir: Project directory for PROJECT scope
         **kwargs: Additional store-specific configuration
-    
+
     Returns:
         The configured PersistenceConfig
     """
     global _config, _manager
-    
+
     config = PersistenceConfig(project_dir=project_dir)
-    
+
     if memory_store_class:
         config.memory_store_class = memory_store_class
     if conversation_store_class:
@@ -244,23 +348,27 @@ def configure_persistence(
         config.task_store_class = task_store_class
     if preferences_store_class:
         config.preferences_store_class = preferences_store_class
-    
+    if knowledge_store_class:
+        config.knowledge_store_class = knowledge_store_class
+    if audit_store_class:
+        config.audit_store_class = audit_store_class
+
     _config = config
     _manager = None  # Reset manager to use new config
-    
+
     return config
 
 
 def get_persistence_manager() -> PersistenceManager:
     """
     Get the global persistence manager.
-    
+
     Creates a new manager with default config if not configured.
     """
     global _manager
-    
+
     if _manager is None:
         _manager = PersistenceManager(_config)
-    
+
     return _manager
 

agent_runtime_core/steps.py

@@ -0,0 +1,373 @@
+"""
+Step executor for long-running multi-step agent operations.
+
+This module provides a structured way to execute multi-step operations
+with automatic checkpointing, resume capability, retries, and progress
+reporting.
+
+Example usage:
+    from agent_runtime_core.steps import StepExecutor, Step
+
+    class MyAgent(AgentRuntime):
+        async def run(self, ctx: RunContext) -> RunResult:
+            executor = StepExecutor(ctx)
+
+            result = await executor.run([
+                Step("fetch", self.fetch_data),
+                Step("process", self.process_data, retries=3),
+                Step("validate", self.validate),
+            ])
+
+            return RunResult(final_output=result)
+"""
+
+import asyncio
+import traceback
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any, Awaitable, Callable, Optional, TypeVar, Union
+from uuid import UUID, uuid4
+
+from agent_runtime_core.interfaces import EventType, RunContext
+
+
+class StepStatus(str, Enum):
+    """Status of a step execution."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    SKIPPED = "skipped"
+    CANCELLED = "cancelled"
+
+
+# Type for step functions: async def step_fn(ctx, state) -> result
+StepFunction = Callable[[RunContext, dict], Awaitable[Any]]
+
+
+@dataclass
+class Step:
+    """
+    Definition of a single step in a multi-step operation.
+
+    Args:
+        name: Unique identifier for this step
+        fn: Async function to execute. Receives (ctx, state) and returns result.
+        retries: Number of retry attempts on failure (default: 0)
+        retry_delay: Seconds to wait between retries (default: 1.0)
+        timeout: Optional timeout in seconds for this step
+        description: Human-readable description for progress reporting
+        checkpoint: Whether to checkpoint after this step (default: True)
+    """
+
+    name: str
+    fn: StepFunction
+    retries: int = 0
+    retry_delay: float = 1.0
+    timeout: Optional[float] = None
+    description: Optional[str] = None
+    checkpoint: bool = True
+
+
+@dataclass
+class StepResult:
+    """Result of executing a single step."""
+
+    name: str
+    status: StepStatus
+    result: Any = None
+    error: Optional[str] = None
+    attempts: int = 1
+    started_at: Optional[datetime] = None
+    completed_at: Optional[datetime] = None
+    duration_ms: Optional[float] = None
+
+
+@dataclass
+class ExecutionState:
+    """
+    State of a multi-step execution.
+
+    This is what gets checkpointed and can be used to resume.
+    """
+
+    execution_id: UUID = field(default_factory=uuid4)
+    current_step_index: int = 0
+    completed_steps: list[str] = field(default_factory=list)
+    step_results: dict[str, Any] = field(default_factory=dict)
+    started_at: datetime = field(default_factory=datetime.utcnow)
+    custom_state: dict = field(default_factory=dict)
+
+    def to_dict(self) -> dict:
+        """Convert to dictionary for checkpointing."""
+        return {
+            "execution_id": str(self.execution_id),
+            "current_step_index": self.current_step_index,
+            "completed_steps": self.completed_steps,
+            "step_results": self.step_results,
+            "started_at": self.started_at.isoformat(),
+            "custom_state": self.custom_state,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict) -> "ExecutionState":
+        """Restore from checkpointed dictionary."""
+        return cls(
+            execution_id=UUID(data["execution_id"]),
+            current_step_index=data["current_step_index"],
+            completed_steps=data["completed_steps"],
+            step_results=data["step_results"],
+            started_at=datetime.fromisoformat(data["started_at"]),
+            custom_state=data.get("custom_state", {}),
+        )
+
+
+class StepExecutionError(Exception):
+    """Raised when step execution fails after all retries."""
+
+    def __init__(self, step_name: str, message: str, attempts: int):
+        self.step_name = step_name
+        self.attempts = attempts
+        super().__init__(f"Step '{step_name}' failed after {attempts} attempts: {message}")
+
+
+class StepCancelledError(Exception):
+    """Raised when execution is cancelled."""
+
+    def __init__(self, step_name: str):
+        self.step_name = step_name
+        super().__init__(f"Execution cancelled during step '{step_name}'")
+
+
+class StepExecutor:
+    """
+    Executes a sequence of steps with checkpointing and resume capability.
+
+    Features:
+    - Automatic checkpointing after each step
+    - Resume from last checkpoint on restart
+    - Per-step retries with configurable delay
+    - Progress reporting via events
+    - Cancellation support
+    - Step-level timeouts
+
+    Example:
+        executor = StepExecutor(ctx)
+
+        result = await executor.run([
+            Step("fetch", fetch_data),
+            Step("process", process_data, retries=3),
+            Step("save", save_results),
+        ])
+    """
+
+    def __init__(
+        self,
+        ctx: RunContext,
+        *,
+        checkpoint_key: str = "_step_executor_state",
+        cancel_check_interval: float = 0.5,
+    ):
+        """
+        Initialize the step executor.
+
+        Args:
+            ctx: The run context from the agent runtime
+            checkpoint_key: Key used for storing execution state
+            cancel_check_interval: How often to check for cancellation (seconds)
+        """
+        self.ctx = ctx
+        self.checkpoint_key = checkpoint_key
+        self.cancel_check_interval = cancel_check_interval
+        self._state: Optional[ExecutionState] = None
+
+    async def run(
+        self,
+        steps: list[Step],
+        *,
+        initial_state: Optional[dict] = None,
+        resume: bool = True,
+    ) -> dict[str, Any]:
+        """
+        Execute a sequence of steps.
+
+        Args:
+            steps: List of steps to execute
+            initial_state: Optional initial custom state
+            resume: Whether to resume from checkpoint if available
+
+        Returns:
+            Dictionary mapping step names to their results
+
+        Raises:
+            StepExecutionError: If a step fails after all retries
+            StepCancelledError: If execution is cancelled
+        """
+        # Try to resume from checkpoint
+        if resume:
+            self._state = await self._load_state()
+
+        if self._state is None:
+            self._state = ExecutionState(
+                custom_state=initial_state or {}
+            )
+
+        total_steps = len(steps)
+
+        for i, step in enumerate(steps):
+            # Skip already completed steps
+            if step.name in self._state.completed_steps:
+                await self.ctx.emit(EventType.STEP_SKIPPED, {
+                    "step_name": step.name,
+                    "step_index": i,
+                    "total_steps": total_steps,
+                    "reason": "already_completed",
+                })
+                continue
+
+            # Check for cancellation
+            if self.ctx.cancelled():
+                raise StepCancelledError(step.name)
+
+            # Update state
+            self._state.current_step_index = i
+
+            # Execute the step
+            result = await self._execute_step(step, i, total_steps)
+
+            # Record completion
+            self._state.completed_steps.append(step.name)
+            self._state.step_results[step.name] = result.result
+
+            # Checkpoint if enabled
+            if step.checkpoint:
+                await self._save_state()
+
+        return self._state.step_results
+
+    async def _execute_step(
+        self,
+        step: Step,
+        index: int,
+        total: int,
+    ) -> StepResult:
+        """Execute a single step with retries."""
+        attempts = 0
+        last_error: Optional[str] = None
+
+        while attempts <= step.retries:
+            attempts += 1
+
+            # Emit started event
+            await self.ctx.emit(EventType.STEP_STARTED, {
+                "step_name": step.name,
+                "step_index": index,
+                "total_steps": total,
+                "attempt": attempts,
+                "max_attempts": step.retries + 1,
+                "description": step.description,
+            })
+
+            # Emit progress
+            await self.ctx.emit(EventType.PROGRESS_UPDATE, {
+                "step_name": step.name,
+                "step_index": index,
+                "total_steps": total,
+                "progress_percent": (index / total) * 100,
+                "description": step.description or f"Executing {step.name}",
+            })
+
+            started_at = datetime.utcnow()
+
+            try:
+                # Execute with optional timeout
+                if step.timeout:
+                    result = await asyncio.wait_for(
+                        step.fn(self.ctx, self._state.custom_state),
+                        timeout=step.timeout,
+                    )
+                else:
+                    result = await step.fn(self.ctx, self._state.custom_state)
+
+                completed_at = datetime.utcnow()
+                duration_ms = (completed_at - started_at).total_seconds() * 1000
+
+                # Emit completed event
+                await self.ctx.emit(EventType.STEP_COMPLETED, {
+                    "step_name": step.name,
+                    "step_index": index,
+                    "total_steps": total,
+                    "attempt": attempts,
+                    "duration_ms": duration_ms,
+                })
+
+                return StepResult(
+                    name=step.name,
+                    status=StepStatus.COMPLETED,
+                    result=result,
+                    attempts=attempts,
+                    started_at=started_at,
+                    completed_at=completed_at,
+                    duration_ms=duration_ms,
+                )
+
+            except asyncio.CancelledError:
+                raise StepCancelledError(step.name)
+
+            except asyncio.TimeoutError:
+                last_error = f"Step timed out after {step.timeout}s"
+
+            except Exception as e:
+                last_error = f"{type(e).__name__}: {str(e)}"
+
+            # Check if we should retry
+            if attempts <= step.retries:
+                await self.ctx.emit(EventType.STEP_RETRYING, {
+                    "step_name": step.name,
+                    "step_index": index,
+                    "attempt": attempts,
+                    "max_attempts": step.retries + 1,
+                    "error": last_error,
+                    "retry_delay": step.retry_delay,
+                })
+                await asyncio.sleep(step.retry_delay)
+
+        # All retries exhausted
+        await self.ctx.emit(EventType.STEP_FAILED, {
+            "step_name": step.name,
+            "step_index": index,
+            "total_steps": total,
+            "attempts": attempts,
+            "error": last_error,
+        })
+
+        raise StepExecutionError(step.name, last_error or "Unknown error", attempts)
+
+    async def _load_state(self) -> Optional[ExecutionState]:
+        """Load execution state from checkpoint."""
+        checkpoint = await self.ctx.get_state()
+        if checkpoint and self.checkpoint_key in checkpoint:
+            try:
+                return ExecutionState.from_dict(checkpoint[self.checkpoint_key])
+            except (KeyError, ValueError):
+                return None
+        return None
+
+    async def _save_state(self) -> None:
+        """Save execution state to checkpoint."""
+        checkpoint = await self.ctx.get_state() or {}
+        checkpoint[self.checkpoint_key] = self._state.to_dict()
+        await self.ctx.checkpoint(checkpoint)
+
+    @property
+    def state(self) -> Optional[ExecutionState]:
+        """Get the current execution state."""
+        return self._state
+
+    def update_custom_state(self, updates: dict) -> None:
+        """Update custom state (will be checkpointed with next step)."""
+        if self._state:
+            self._state.custom_state.update(updates)
+