agno 2.2.13__py3-none-any.whl → 2.4.3__py3-none-any.whl

agno/run/cancellation_management/in_memory_cancellation_manager.py

@@ -0,0 +1,100 @@
+"""Run cancellation management."""
+
+import asyncio
+import threading
+from typing import Dict
+
+from agno.exceptions import RunCancelledException
+from agno.run.cancellation_management.base import BaseRunCancellationManager
+from agno.utils.log import logger
+
+
+class InMemoryRunCancellationManager(BaseRunCancellationManager):
+    def __init__(self):
+        self._cancelled_runs: Dict[str, bool] = {}
+        self._lock = threading.Lock()
+        self._async_lock = asyncio.Lock()
+
+    def register_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled."""
+        with self._lock:
+            self._cancelled_runs[run_id] = False
+
+    async def aregister_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled (async version)."""
+        async with self._async_lock:
+            self._cancelled_runs[run_id] = False
+
+    def cancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled.
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        with self._lock:
+            if run_id in self._cancelled_runs:
+                self._cancelled_runs[run_id] = True
+                logger.info(f"Run {run_id} marked for cancellation")
+                return True
+            else:
+                logger.warning(f"Attempted to cancel unknown run {run_id}")
+                return False
+
+    async def acancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled (async version).
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        async with self._async_lock:
+            if run_id in self._cancelled_runs:
+                self._cancelled_runs[run_id] = True
+                logger.info(f"Run {run_id} marked for cancellation")
+                return True
+            else:
+                logger.warning(f"Attempted to cancel unknown run {run_id}")
+                return False
+
+    def is_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled."""
+        with self._lock:
+            return self._cancelled_runs.get(run_id, False)
+
+    async def ais_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled (async version)."""
+        async with self._async_lock:
+            return self._cancelled_runs.get(run_id, False)
+
+    def cleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes)."""
+        with self._lock:
+            if run_id in self._cancelled_runs:
+                del self._cancelled_runs[run_id]
+
+    async def acleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes) (async version)."""
+        async with self._async_lock:
+            if run_id in self._cancelled_runs:
+                del self._cancelled_runs[run_id]
+
+    def raise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so."""
+        if self.is_cancelled(run_id):
+            logger.info(f"Cancelling run {run_id}")
+            raise RunCancelledException(f"Run {run_id} was cancelled")
+
+    async def araise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so (async version)."""
+        if await self.ais_cancelled(run_id):
+            logger.info(f"Cancelling run {run_id}")
+            raise RunCancelledException(f"Run {run_id} was cancelled")
+
+    def get_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status."""
+        with self._lock:
+            return self._cancelled_runs.copy()
+
+    async def aget_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status (async version)."""
+        async with self._async_lock:
+            return self._cancelled_runs.copy()

agno/run/cancellation_management/redis_cancellation_manager.py

@@ -0,0 +1,236 @@
+"""Redis-based run cancellation management."""
+
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+
+from agno.exceptions import RunCancelledException
+from agno.run.cancellation_management.base import BaseRunCancellationManager
+from agno.utils.log import logger
+
+# Defer import error until class instantiation
+_redis_available = True
+_redis_import_error: Optional[str] = None
+
+try:
+    from redis import Redis, RedisCluster
+    from redis.asyncio import Redis as AsyncRedis
+    from redis.asyncio import RedisCluster as AsyncRedisCluster
+except ImportError:
+    _redis_available = False
+    _redis_import_error = "`redis` not installed. Please install it using `pip install redis`"
+    # Type hints for when redis is not installed
+    if TYPE_CHECKING:
+        from redis import Redis, RedisCluster
+        from redis.asyncio import Redis as AsyncRedis
+        from redis.asyncio import RedisCluster as AsyncRedisCluster
+    else:
+        Redis = Any
+        RedisCluster = Any
+        AsyncRedis = Any
+        AsyncRedisCluster = Any
+
+
+class RedisRunCancellationManager(BaseRunCancellationManager):
+    """Redis-based cancellation manager for distributed run cancellation.
+    This manager stores run cancellation state in Redis, enabling cancellation
+    across multiple processes or services.
+
+    To use: call the set_cancellation_manager function to set the cancellation manager.
+    Args:
+        redis_client: Sync Redis client for sync methods. Can be Redis or RedisCluster.
+        async_redis_client: Async Redis client for async methods. Can be AsyncRedis or AsyncRedisCluster.
+        key_prefix: Prefix for Redis keys. Defaults to "agno:run:cancellation:".
+        ttl_seconds: TTL for keys in seconds. Defaults to 86400 (1 day).
+            Keys auto-expire to prevent orphaned keys if runs aren't cleaned up.
+            Set to None to disable expiration.
+    """
+
+    DEFAULT_TTL_SECONDS = 60 * 60 * 24  # 1 day
+
+    def __init__(
+        self,
+        redis_client: Optional[Union[Redis, RedisCluster]] = None,
+        async_redis_client: Optional[Union[AsyncRedis, AsyncRedisCluster]] = None,
+        key_prefix: str = "agno:run:cancellation:",
+        ttl_seconds: Optional[int] = DEFAULT_TTL_SECONDS,
+    ):
+        if not _redis_available:
+            raise ImportError(_redis_import_error)
+
+        super().__init__()
+        self.redis_client = redis_client
+        self.async_redis_client = async_redis_client
+        self.key_prefix = key_prefix
+        self.ttl_seconds = ttl_seconds
+
+        if redis_client is None and async_redis_client is None:
+            raise ValueError("At least one of redis_client or async_redis_client must be provided")
+
+    def _get_key(self, run_id: str) -> str:
+        """Get the Redis key for a run ID."""
+        return f"{self.key_prefix}{run_id}"
+
+    def _ensure_sync_client(self) -> Union[Redis, RedisCluster]:
+        """Ensure sync client is available."""
+        if self.redis_client is None:
+            raise RuntimeError("Sync Redis client not provided. Use async methods or provide a sync client.")
+        return self.redis_client
+
+    def _ensure_async_client(self) -> Union[AsyncRedis, AsyncRedisCluster]:
+        """Ensure async client is available."""
+        if self.async_redis_client is None:
+            raise RuntimeError("Async Redis client not provided. Use sync methods or provide an async client.")
+        return self.async_redis_client
+
+    def register_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled."""
+        client = self._ensure_sync_client()
+        key = self._get_key(run_id)
+        client.set(key, "0", ex=self.ttl_seconds)
+
+    async def aregister_run(self, run_id: str) -> None:
+        """Register a new run as not cancelled (async version)."""
+        client = self._ensure_async_client()
+        key = self._get_key(run_id)
+        await client.set(key, "0", ex=self.ttl_seconds)
+
+    def cancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled.
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        client = self._ensure_sync_client()
+        key = self._get_key(run_id)
+
+        # Atomically set to "1" only if key exists (XX flag)
+        result = client.set(key, "1", ex=self.ttl_seconds, xx=True)
+
+        if result:
+            logger.info(f"Run {run_id} marked for cancellation")
+            return True
+        else:
+            logger.warning(f"Attempted to cancel unknown run {run_id}")
+            return False
+
+    async def acancel_run(self, run_id: str) -> bool:
+        """Cancel a run by marking it as cancelled (async version).
+
+        Returns:
+            bool: True if run was found and cancelled, False if run not found.
+        """
+        client = self._ensure_async_client()
+        key = self._get_key(run_id)
+
+        # Atomically set to "1" only if key exists (XX flag)
+        result = await client.set(key, "1", ex=self.ttl_seconds, xx=True)
+
+        if result:
+            logger.info(f"Run {run_id} marked for cancellation")
+            return True
+        else:
+            logger.warning(f"Attempted to cancel unknown run {run_id}")
+            return False
+
+    def is_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled."""
+        client = self._ensure_sync_client()
+        key = self._get_key(run_id)
+        value = client.get(key)
+        if value is None:
+            return False
+        # Redis returns bytes, handle both bytes and str
+        if isinstance(value, bytes):
+            return value == b"1"
+        return value == "1"
+
+    async def ais_cancelled(self, run_id: str) -> bool:
+        """Check if a run is cancelled (async version)."""
+        client = self._ensure_async_client()
+        key = self._get_key(run_id)
+        value = await client.get(key)
+        if value is None:
+            return False
+        # Redis returns bytes, handle both bytes and str
+        if isinstance(value, bytes):
+            return value == b"1"
+        return value == "1"
+
+    def cleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes)."""
+        client = self._ensure_sync_client()
+        key = self._get_key(run_id)
+        client.delete(key)
+
+    async def acleanup_run(self, run_id: str) -> None:
+        """Remove a run from tracking (called when run completes) (async version)."""
+        client = self._ensure_async_client()
+        key = self._get_key(run_id)
+        await client.delete(key)
+
+    def raise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so."""
+        if self.is_cancelled(run_id):
+            logger.info(f"Cancelling run {run_id}")
+            raise RunCancelledException(f"Run {run_id} was cancelled")
+
+    async def araise_if_cancelled(self, run_id: str) -> None:
+        """Check if a run should be cancelled and raise exception if so (async version)."""
+        if await self.ais_cancelled(run_id):
+            logger.info(f"Cancelling run {run_id}")
+            raise RunCancelledException(f"Run {run_id} was cancelled")
+
+    def get_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status.
+
+        Note: Uses scan_iter which works correctly with both standalone Redis
+        and Redis Cluster (scans all nodes in cluster mode).
+        """
+        client = self._ensure_sync_client()
+        result: Dict[str, bool] = {}
+
+        # scan_iter handles cluster mode correctly (scans all nodes)
+        pattern = f"{self.key_prefix}*"
+        for key in client.scan_iter(match=pattern, count=100):
+            # Extract run_id from key
+            if isinstance(key, bytes):
+                key = key.decode("utf-8")
+            run_id = key[len(self.key_prefix) :]
+
+            # Get value
+            value = client.get(key)
+            if value is not None:
+                if isinstance(value, bytes):
+                    is_cancelled = value == b"1"
+                else:
+                    is_cancelled = value == "1"
+                result[run_id] = is_cancelled
+
+        return result
+
+    async def aget_active_runs(self) -> Dict[str, bool]:
+        """Get all currently tracked runs and their cancellation status (async version).
+
+        Note: Uses scan_iter which works correctly with both standalone Redis
+        and Redis Cluster (scans all nodes in cluster mode).
+        """
+        client = self._ensure_async_client()
+        result: Dict[str, bool] = {}
+
+        # scan_iter handles cluster mode correctly (scans all nodes)
+        pattern = f"{self.key_prefix}*"
+        async for key in client.scan_iter(match=pattern, count=100):
+            # Extract run_id from key
+            if isinstance(key, bytes):
+                key = key.decode("utf-8")
+            run_id = key[len(self.key_prefix) :]
+
+            # Get value
+            value = await client.get(key)
+            if value is not None:
+                if isinstance(value, bytes):
+                    is_cancelled = value == b"1"
+                else:
+                    is_cancelled = value == "1"
+                result[run_id] = is_cancelled
+
+        return result

agno/run/requirement.py

@@ -0,0 +1,181 @@
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from uuid import uuid4
+
+from agno.models.response import ToolExecution, UserInputField
+
+if TYPE_CHECKING:
+    pass
+
+
+@dataclass
+class RunRequirement:
+    """Requirement to complete a paused run (used in HITL flows)"""
+
+    tool_execution: Optional[ToolExecution] = None
+    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+    # User confirmation
+    confirmation: Optional[bool] = None
+    confirmation_note: Optional[str] = None
+
+    # User input
+    user_input_schema: Optional[List[UserInputField]] = None
+
+    # External execution
+    external_execution_result: Optional[str] = None
+
+    def __init__(
+        self,
+        tool_execution: ToolExecution,
+        id: Optional[str] = None,
+        created_at: Optional[datetime] = None,
+    ):
+        self.id = id or str(uuid4())
+        self.tool_execution = tool_execution
+        self.user_input_schema = tool_execution.user_input_schema if tool_execution else None
+        self.created_at = created_at or datetime.now(timezone.utc)
+        self.confirmation = None
+        self.confirmation_note = None
+        self.external_execution_result = None
+
+    @property
+    def needs_confirmation(self) -> bool:
+        if self.confirmation is not None:
+            return False
+        if not self.tool_execution:
+            return False
+        if self.tool_execution.confirmed is True:
+            return True
+
+        return self.tool_execution.requires_confirmation or False
+
+    @property
+    def needs_user_input(self) -> bool:
+        if not self.tool_execution:
+            return False
+        if self.tool_execution.answered is True:
+            return False
+        if self.user_input_schema and not all(field.value is not None for field in self.user_input_schema):
+            return True
+
+        return self.tool_execution.requires_user_input or False
+
+    @property
+    def needs_external_execution(self) -> bool:
+        if not self.tool_execution:
+            return False
+        if self.external_execution_result is not None:
+            return True
+
+        return self.tool_execution.external_execution_required or False
+
+    def confirm(self):
+        if not self.needs_confirmation:
+            raise ValueError("This requirement does not require confirmation")
+        self.confirmation = True
+        if self.tool_execution:
+            self.tool_execution.confirmed = True
+
+    def reject(self):
+        if not self.needs_confirmation:
+            raise ValueError("This requirement does not require confirmation")
+        self.confirmation = False
+        if self.tool_execution:
+            self.tool_execution.confirmed = False
+
+    def set_external_execution_result(self, result: str):
+        if not self.needs_external_execution:
+            raise ValueError("This requirement does not require external execution")
+        self.external_execution_result = result
+        if self.tool_execution:
+            self.tool_execution.result = result
+
+    def update_tool(self):
+        if not self.tool_execution:
+            return
+        if self.confirmation is True:
+            self.tool_execution.confirmed = True
+        elif self.confirmation is False:
+            self.tool_execution.confirmed = False
+        else:
+            raise ValueError("This requirement does not require confirmation or user input")
+
+    def is_resolved(self) -> bool:
+        """Return True if the requirement has been resolved"""
+        return not self.needs_confirmation and not self.needs_user_input and not self.needs_external_execution
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to JSON-serializable dictionary for storage."""
+        _dict: Dict[str, Any] = {
+            "id": self.id,
+            "created_at": self.created_at.isoformat() if isinstance(self.created_at, datetime) else self.created_at,
+            "confirmation": self.confirmation,
+            "confirmation_note": self.confirmation_note,
+            "external_execution_result": self.external_execution_result,
+        }
+
+        if self.tool_execution is not None:
+            _dict["tool_execution"] = (
+                self.tool_execution.to_dict() if isinstance(self.tool_execution, ToolExecution) else self.tool_execution
+            )
+
+        if self.user_input_schema is not None:
+            _dict["user_input_schema"] = [f.to_dict() if hasattr(f, "to_dict") else f for f in self.user_input_schema]
+
+        return {k: v for k, v in _dict.items() if v is not None}
+
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "RunRequirement":
+        """Reconstruct from stored dictionary."""
+        if data is None:
+            raise ValueError("RunRequirement.from_dict() requires a non-None dict")
+
+        # Handle tool_execution
+        tool_data = data.get("tool_execution")
+        tool_execution: Optional[ToolExecution] = None
+        if isinstance(tool_data, ToolExecution):
+            tool_execution = tool_data
+        elif isinstance(tool_data, dict):
+            tool_execution = ToolExecution.from_dict(tool_data)
+
+        # Handle created_at (ISO string or datetime)
+        created_at_raw = data.get("created_at")
+        created_at: Optional[datetime] = None
+        if isinstance(created_at_raw, datetime):
+            created_at = created_at_raw
+        elif isinstance(created_at_raw, str):
+            try:
+                created_at = datetime.fromisoformat(created_at_raw)
+            except ValueError:
+                created_at = None
+
+        # Build requirement - tool_execution is required by __init__
+        # For legacy data without tool_execution, create a minimal placeholder
+        if tool_execution is None:
+            tool_execution = ToolExecution(tool_name="unknown", tool_args={})
+
+        requirement = cls(
+            tool_execution=tool_execution,
+            id=data.get("id"),
+            created_at=created_at,
+        )
+
+        # Set optional fields
+        requirement.confirmation = data.get("confirmation")
+        requirement.confirmation_note = data.get("confirmation_note")
+        requirement.external_execution_result = data.get("external_execution_result")
+
+        # Handle user_input_schema
+        schema_raw = data.get("user_input_schema")
+        if schema_raw is not None:
+            rebuilt_schema: List[UserInputField] = []
+            for item in schema_raw:
+                if isinstance(item, UserInputField):
+                    rebuilt_schema.append(item)
+                elif isinstance(item, dict):
+                    rebuilt_schema.append(UserInputField.from_dict(item))
+            requirement.user_input_schema = rebuilt_schema if rebuilt_schema else None
+
+        return requirement