lionagi 0.13.7__py3-none-any.whl → 0.14.1__py3-none-any.whl

lionagi/libs/concurrency/__init__.py

@@ -0,0 +1,25 @@
+"""Structured concurrency primitives.
+
+This module provides structured concurrency primitives using AnyIO,
+which allows for consistent behavior across asyncio and trio backends.
+"""
+
+from .cancel import CancelScope, fail_after, move_on_after
+from .errors import get_cancelled_exc_class, shield
+from .primitives import CapacityLimiter, Condition, Event, Lock, Semaphore
+from .task import TaskGroup, create_task_group
+
+__all__ = [
+    "TaskGroup",
+    "create_task_group",
+    "CancelScope",
+    "move_on_after",
+    "fail_after",
+    "Lock",
+    "Semaphore",
+    "CapacityLimiter",
+    "Event",
+    "Condition",
+    "get_cancelled_exc_class",
+    "shield",
+]

lionagi/libs/concurrency/cancel.py

@@ -0,0 +1,134 @@
+"""Cancellation scope implementation for structured concurrency."""
+
+import time
+from collections.abc import Iterator
+from contextlib import contextmanager
+from types import TracebackType
+from typing import Optional, TypeVar
+
+import anyio
+
+T = TypeVar("T")
+
+
+class CancelScope:
+    """A context manager for controlling cancellation of tasks."""
+
+    def __init__(self, deadline: float | None = None, shield: bool = False):
+        """Initialize a new cancel scope.
+
+        Args:
+            deadline: The time (in seconds since the epoch) when this scope should be cancelled
+            shield: If True, this scope shields its contents from external cancellation
+        """
+        self._scope = None
+        self._deadline = deadline
+        self._shield = shield
+        self.cancel_called = False
+        self.cancelled_caught = False
+
+    def cancel(self) -> None:
+        """Cancel this scope.
+
+        This will cause all tasks within this scope to be cancelled.
+        """
+        self.cancel_called = True
+        if self._scope is not None:
+            self._scope.cancel()
+
+    def __enter__(self) -> "CancelScope":
+        """Enter the cancel scope context.
+
+        Returns:
+            The cancel scope instance.
+        """
+        # Use math.inf as the default deadline (no timeout)
+        import math
+
+        deadline = self._deadline if self._deadline is not None else math.inf
+        self._scope = anyio.CancelScope(deadline=deadline, shield=self._shield)
+        if self.cancel_called:
+            self._scope.cancel()
+        self._scope.__enter__()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> bool:
+        """Exit the cancel scope context.
+
+        Returns:
+            True if the exception was handled, False otherwise.
+        """
+        if self._scope is None:
+            return False
+
+        try:
+            result = self._scope.__exit__(exc_type, exc_val, exc_tb)
+            self.cancelled_caught = self._scope.cancelled_caught
+            return result
+        finally:
+            self._scope = None
+
+
+@contextmanager
+def move_on_after(seconds: float | None) -> Iterator[CancelScope]:
+    """Return a context manager that cancels its contents after the given number of seconds.
+
+    Args:
+        seconds: The number of seconds to wait before cancelling, or None to disable the timeout
+
+    Returns:
+        A cancel scope that will be cancelled after the specified time
+
+    Example:
+        with move_on_after(5) as scope:
+            await long_running_operation()
+            if scope.cancelled_caught:
+                print("Operation timed out")
+    """
+    deadline = None if seconds is None else time.time() + seconds
+    scope = CancelScope(deadline=deadline)
+    with scope:
+        yield scope
+
+
+@contextmanager
+def fail_after(seconds: float | None) -> Iterator[CancelScope]:
+    """Return a context manager that raises TimeoutError if its contents take longer than the given time.
+
+    Args:
+        seconds: The number of seconds to wait before raising TimeoutError, or None to disable the timeout
+
+    Returns:
+        A cancel scope that will raise TimeoutError after the specified time
+
+    Raises:
+        TimeoutError: If the operation takes longer than the specified time
+
+    Example:
+        try:
+            with fail_after(5):
+                await long_running_operation()
+        except TimeoutError:
+            print("Operation timed out")
+    """
+    if seconds is None:
+        # No timeout
+        scope = CancelScope(shield=True)
+        with scope:
+            yield scope
+    else:
+        deadline = time.time() + seconds
+        scope = CancelScope(deadline=deadline)
+        try:
+            with scope:
+                yield scope
+        finally:
+            if scope.cancelled_caught:
+                raise TimeoutError(
+                    f"Operation took longer than {seconds} seconds"
+                )

lionagi/libs/concurrency/errors.py

@@ -0,0 +1,35 @@
+"""Error handling utilities for structured concurrency."""
+
+from collections.abc import Awaitable, Callable
+from typing import Any, TypeVar
+
+import anyio
+
+T = TypeVar("T")
+
+
+def get_cancelled_exc_class() -> type[BaseException]:
+    """Get the exception class used for cancellation.
+
+    Returns:
+        The exception class used for cancellation (CancelledError for asyncio,
+        Cancelled for trio).
+    """
+    return anyio.get_cancelled_exc_class()
+
+
+async def shield(
+    func: Callable[..., Awaitable[T]], *args: Any, **kwargs: Any
+) -> T:
+    """Run a coroutine function with protection from cancellation.
+
+    Args:
+        func: The coroutine function to call
+        *args: Positional arguments to pass to the function
+        **kwargs: Keyword arguments to pass to the function
+
+    Returns:
+        The return value of the function
+    """
+    with anyio.CancelScope(shield=True):
+        return await func(*args, **kwargs)

lionagi/libs/concurrency/patterns.py

@@ -0,0 +1,252 @@
+"""Common concurrency patterns for structured concurrency."""
+
+import math
+from collections.abc import Awaitable, Callable
+from types import TracebackType
+from typing import Any, Optional, TypeVar
+
+import anyio
+
+from .cancel import move_on_after
+from .primitives import CapacityLimiter, Lock
+from .task import create_task_group
+
+T = TypeVar("T")
+R = TypeVar("R")
+Response = TypeVar("Response")
+
+
+class ConnectionPool:
+    """A pool of reusable connections."""
+
+    def __init__(
+        self,
+        max_connections: int,
+        connection_factory: Callable[[], Awaitable[T]],
+    ):
+        """Initialize a new connection pool.
+
+        Args:
+            max_connections: The maximum number of connections in the pool
+            connection_factory: A factory function that creates new connections
+        """
+        self._connection_factory = connection_factory
+        self._limiter = CapacityLimiter(max_connections)
+        self._connections: list[T] = []
+        self._lock = Lock()
+
+    async def acquire(self) -> T:
+        """Acquire a connection from the pool.
+
+        Returns:
+            A connection from the pool, or a new connection if the pool is empty.
+        """
+        async with self._limiter:
+            async with self._lock:
+                if self._connections:
+                    return self._connections.pop()
+
+            # No connections available, create a new one
+            return await self._connection_factory()
+
+    async def release(self, connection: T) -> None:
+        """Release a connection back to the pool.
+
+        Args:
+            connection: The connection to release
+        """
+        async with self._lock:
+            self._connections.append(connection)
+
+    async def __aenter__(self) -> "ConnectionPool":
+        """Enter the connection pool context.
+
+        Returns:
+            The connection pool instance.
+        """
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit the connection pool context, closing all connections."""
+        async with self._lock:
+            for connection in self._connections:
+                if hasattr(connection, "close"):
+                    await connection.close()
+                elif hasattr(connection, "disconnect"):
+                    await connection.disconnect()
+            self._connections.clear()
+
+
+async def parallel_requests(
+    urls: list[str],
+    fetch_func: Callable[[str], Awaitable[Response]],
+    max_concurrency: int = 10,
+) -> list[Response]:
+    """Fetch multiple URLs in parallel with limited concurrency.
+
+    Args:
+        urls: The URLs to fetch
+        fetch_func: The function to use for fetching
+        max_concurrency: The maximum number of concurrent requests
+
+    Returns:
+        A list of responses in the same order as the URLs
+    """
+    limiter = CapacityLimiter(max_concurrency)
+    results: list[Response | None] = [None] * len(urls)
+    exceptions: list[Exception | None] = [None] * len(urls)
+
+    async def fetch_with_limit(index: int, url: str) -> None:
+        async with limiter:
+            try:
+                results[index] = await fetch_func(url)
+            except Exception as exc:
+                exceptions[index] = exc
+
+    async with create_task_group() as tg:
+        for i, url in enumerate(urls):
+            await tg.start_soon(fetch_with_limit, i, url)
+
+    # Check for exceptions
+    for i, exc in enumerate(exceptions):
+        if exc is not None:
+            raise exc
+
+    return results  # type: ignore
+
+
+async def retry_with_timeout(
+    func: Callable[..., Awaitable[T]],
+    *args: Any,
+    max_retries: int = 3,
+    timeout: float = 5.0,
+    retry_exceptions: list[type[Exception]] | None = None,
+    **kwargs: Any,
+) -> T:
+    """Execute a function with retry logic and timeout.
+
+    Args:
+        func: The function to call
+        *args: Positional arguments to pass to the function
+        max_retries: The maximum number of retry attempts
+        timeout: The timeout for each attempt in seconds
+        retry_exceptions: List of exception types to retry on, or None to retry on any exception
+        **kwargs: Keyword arguments to pass to the function
+
+    Returns:
+        The return value of the function
+
+    Raises:
+        TimeoutError: If all retry attempts time out
+        Exception: If the function raises an exception after all retry attempts
+    """
+    retry_exceptions = retry_exceptions or [Exception]
+    last_exception = None
+
+    for attempt in range(max_retries):
+        try:
+            timed_out = False
+            with move_on_after(timeout) as scope:
+                result = await func(*args, **kwargs)
+                if not scope.cancelled_caught:
+                    return result
+                timed_out = True
+
+            # If we get here, the operation timed out
+            if timed_out:
+                if attempt == max_retries - 1:
+                    raise TimeoutError(
+                        f"Operation timed out after {max_retries} attempts"
+                    )
+
+                # Wait before retrying (exponential backoff)
+                await anyio.sleep(2**attempt)
+
+        except tuple(retry_exceptions) as exc:
+            last_exception = exc
+            if attempt == max_retries - 1:
+                raise
+
+            # Wait before retrying (exponential backoff)
+            await anyio.sleep(2**attempt)
+
+    # This should never be reached, but makes the type checker happy
+    if last_exception:
+        raise last_exception
+    raise RuntimeError("Unreachable code")
+
+
+class WorkerPool:
+    """A pool of worker tasks that process items from a queue."""
+
+    def __init__(
+        self, num_workers: int, worker_func: Callable[[Any], Awaitable[None]]
+    ):
+        """Initialize a new worker pool.
+
+        Args:
+            num_workers: The number of worker tasks to create
+            worker_func: The function that each worker will run
+        """
+        self._num_workers = num_workers
+        self._worker_func = worker_func
+        self._queue = anyio.create_memory_object_stream(math.inf)
+        self._task_group = None
+
+    async def start(self) -> None:
+        """Start the worker pool."""
+        if self._task_group is not None:
+            raise RuntimeError("Worker pool already started")
+
+        self._task_group = create_task_group()
+
+        async with self._task_group as tg:
+            for _ in range(self._num_workers):
+                tg.start_soon(self._worker_loop)
+
+    async def stop(self) -> None:
+        """Stop the worker pool."""
+        if self._task_group is None:
+            return
+
+        # Signal workers to stop
+        for _ in range(self._num_workers):
+            await self._queue[0].send(None)
+
+        # Wait for workers to finish
+        await self._task_group.__aexit__(None, None, None)
+        self._task_group = None
+
+    async def submit(self, item: Any) -> None:
+        """Submit an item to be processed by a worker.
+
+        Args:
+            item: The item to process
+        """
+        if self._task_group is None:
+            raise RuntimeError("Worker pool not started")
+
+        await self._queue[0].send(item)
+
+    async def _worker_loop(self) -> None:
+        """The main loop for each worker task."""
+        while True:
+            try:
+                item = await self._queue[1].receive()
+
+                # None is a signal to stop
+                if item is None:
+                    break
+
+                try:
+                    await self._worker_func(item)
+                except Exception as exc:
+                    # Log the exception but keep the worker running
+                    print(f"Worker error: {exc}")
+            except anyio.EndOfStream:
+                break

lionagi/libs/concurrency/primitives.py

@@ -0,0 +1,242 @@
+"""Resource management primitives for structured concurrency."""
+
+from types import TracebackType
+from typing import Optional
+
+import anyio
+
+
+class Lock:
+    """A mutex lock for controlling access to a shared resource.
+
+    This lock is reentrant, meaning the same task can acquire it multiple times
+    without deadlocking.
+    """
+
+    def __init__(self):
+        """Initialize a new lock."""
+        self._lock = anyio.Lock()
+
+    async def __aenter__(self) -> None:
+        """Acquire the lock.
+
+        If the lock is already held by another task, this will wait until it's released.
+        """
+        await self.acquire()
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Release the lock."""
+        self.release()
+
+    async def acquire(self) -> bool:
+        """Acquire the lock.
+
+        Returns:
+            True if the lock was acquired, False otherwise.
+        """
+        await self._lock.acquire()
+        return True
+
+    def release(self) -> None:
+        """Release the lock.
+
+        Raises:
+            RuntimeError: If the lock is not currently held by this task.
+        """
+        self._lock.release()
+
+
+class Semaphore:
+    """A semaphore for limiting concurrent access to a resource."""
+
+    def __init__(self, initial_value: int):
+        """Initialize a new semaphore.
+
+        Args:
+            initial_value: The initial value of the semaphore (must be >= 0)
+        """
+        if initial_value < 0:
+            raise ValueError("The initial value must be >= 0")
+        self._semaphore = anyio.Semaphore(initial_value)
+
+    async def __aenter__(self) -> None:
+        """Acquire the semaphore.
+
+        If the semaphore value is zero, this will wait until it's released.
+        """
+        await self.acquire()
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Release the semaphore."""
+        self.release()
+
+    async def acquire(self) -> None:
+        """Acquire the semaphore.
+
+        If the semaphore value is zero, this will wait until it's released.
+        """
+        await self._semaphore.acquire()
+
+    def release(self) -> None:
+        """Release the semaphore, incrementing its value."""
+        self._semaphore.release()
+
+
+class CapacityLimiter:
+    """A context manager for limiting the number of concurrent operations."""
+
+    def __init__(self, total_tokens: float):
+        """Initialize a new capacity limiter.
+
+        Args:
+            total_tokens: The maximum number of tokens (>= 1)
+        """
+        if total_tokens < 1:
+            raise ValueError("The total number of tokens must be >= 1")
+        self._limiter = anyio.CapacityLimiter(total_tokens)
+
+    async def __aenter__(self) -> None:
+        """Acquire a token.
+
+        If no tokens are available, this will wait until one is released.
+        """
+        await self.acquire()
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Release the token."""
+        self.release()
+
+    async def acquire(self) -> None:
+        """Acquire a token.
+
+        If no tokens are available, this will wait until one is released.
+        """
+        await self._limiter.acquire()
+
+    def release(self) -> None:
+        """Release a token.
+
+        Raises:
+            RuntimeError: If the current task doesn't hold any tokens.
+        """
+        self._limiter.release()
+
+    @property
+    def total_tokens(self) -> float:
+        """The total number of tokens."""
+        return self._limiter.total_tokens
+
+    @total_tokens.setter
+    def total_tokens(self, value: float) -> None:
+        """Set the total number of tokens.
+
+        Args:
+            value: The new total number of tokens (>= 1)
+        """
+        if value < 1:
+            raise ValueError("The total number of tokens must be >= 1")
+        self._limiter.total_tokens = value
+
+    @property
+    def borrowed_tokens(self) -> int:
+        """The number of tokens currently borrowed."""
+        return self._limiter.borrowed_tokens
+
+    @property
+    def available_tokens(self) -> float:
+        """The number of tokens currently available."""
+        return self._limiter.available_tokens
+
+
+class Event:
+    """An event object for task synchronization.
+
+    An event can be in one of two states: set or unset. When set, tasks waiting
+    on the event are allowed to proceed.
+    """
+
+    def __init__(self):
+        """Initialize a new event in the unset state."""
+        self._event = anyio.Event()
+
+    def is_set(self) -> bool:
+        """Check if the event is set.
+
+        Returns:
+            True if the event is set, False otherwise.
+        """
+        return self._event.is_set()
+
+    def set(self) -> None:
+        """Set the event, allowing all waiting tasks to proceed."""
+        self._event.set()
+
+    async def wait(self) -> None:
+        """Wait until the event is set."""
+        await self._event.wait()
+
+
+class Condition:
+    """A condition variable for task synchronization."""
+
+    def __init__(self, lock: Lock | None = None):
+        """Initialize a new condition.
+
+        Args:
+            lock: The lock to use, or None to create a new one
+        """
+        self._lock = lock or Lock()
+        self._condition = anyio.Condition(self._lock._lock)
+
+    async def __aenter__(self) -> "Condition":
+        """Acquire the underlying lock.
+
+        Returns:
+            The condition instance.
+        """
+        await self._lock.acquire()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Release the underlying lock."""
+        self._lock.release()
+
+    async def wait(self) -> None:
+        """Wait for a notification.
+
+        This releases the underlying lock, waits for a notification, and then
+        reacquires the lock.
+        """
+        await self._condition.wait()
+
+    async def notify(self, n: int = 1) -> None:
+        """Notify waiting tasks.
+
+        Args:
+            n: The number of tasks to notify
+        """
+        await self._condition.notify(n)
+
+    async def notify_all(self) -> None:
+        """Notify all waiting tasks."""
+        await self._condition.notify_all()