lionagi 0.14.4__py3-none-any.whl → 0.14.5__py3-none-any.whl

lionagi/protocols/generic/pile.py

@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import asyncio
-import threading
 from collections import deque
 from collections.abc import (
     AsyncIterator,
@@ -25,6 +24,7 @@ from pydapter import Adaptable, AsyncAdaptable
 from typing_extensions import Self, override
 
 from lionagi._errors import ItemExistsError, ItemNotFoundError
+from lionagi.libs.concurrency import Lock as ConcurrencyLock
 from lionagi.utils import UNDEFINED, is_same_dtype, to_list
 
 from .._concepts import Observable
@@ -38,19 +38,10 @@ T = TypeVar("T", bound=E)
 __all__ = ("Pile",)
 
 
-def synchronized(func: Callable):
-    @wraps(func)
-    def wrapper(self: Pile, *args, **kwargs):
-        with self.lock:
-            return func(self, *args, **kwargs)
-
-    return wrapper
-
-
 def async_synchronized(func: Callable):
     @wraps(func)
     async def wrapper(self: Pile, *args, **kwargs):
-        async with self.async_lock:
+        async with self.lock:
             return await func(self, *args, **kwargs)
 
     return wrapper
@@ -91,8 +82,7 @@ class Pile(Element, Collective[E], Generic[E], Adaptable, AsyncAdaptable):
 
     def __pydantic_extra__(self) -> dict[str, FieldInfo]:
         return {
-            "_lock": Field(default_factory=threading.Lock),
-            "_async": Field(default_factory=asyncio.Lock),
+            "_lock": Field(default_factory=ConcurrencyLock),
         }
 
     def __pydantic_private__(self) -> dict[str, FieldInfo]:
@@ -172,7 +162,6 @@ class Pile(Element, Collective[E], Generic[E], Adaptable, AsyncAdaptable):
         """
         self._setitem(key, item)
 
-    @synchronized
     def pop(
         self,
         key: ID.Ref | ID.RefSeq | int | slice,
@@ -235,7 +224,6 @@ class Pile(Element, Collective[E], Generic[E], Adaptable, AsyncAdaptable):
         """
         self._exclude(item)
 
-    @synchronized
     def clear(self) -> None:
         """Remove all items."""
         self._clear()
@@ -255,7 +243,6 @@ class Pile(Element, Collective[E], Generic[E], Adaptable, AsyncAdaptable):
         """
         self._update(other)
 
-    @synchronized
     def insert(self, index: int, item: T, /) -> None:
         """Insert item at position.
 
@@ -269,7 +256,6 @@ class Pile(Element, Collective[E], Generic[E], Adaptable, AsyncAdaptable):
         """
         self._insert(index, item)
 
-    @synchronized
     def append(self, item: T, /) -> None:
         """Append item to end (alias for include).
 
@@ -281,7 +267,6 @@ class Pile(Element, Collective[E], Generic[E], Adaptable, AsyncAdaptable):
         """
         self.update(item)
 
-    @synchronized
     def get(
         self,
         key: ID.Ref | ID.RefSeq | int | slice,
@@ -321,11 +306,12 @@ class Pile(Element, Collective[E], Generic[E], Adaptable, AsyncAdaptable):
 
     def __iter__(self) -> Iterator[T]:
         """Iterate over items safely."""
-        with self.lock:
-            current_order = list(self.progression)
+        # Take a snapshot of the current order to avoid holding lock during iteration
+        current_order = list(self.progression)
 
         for key in current_order:
-            yield self.collections[key]
+            if key in self.collections:
+                yield self.collections[key]
 
     def __next__(self) -> T:
         """Get next item."""
@@ -478,29 +464,20 @@ class Pile(Element, Collective[E], Generic[E], Adaptable, AsyncAdaptable):
         """Prepare for pickling."""
         state = self.__dict__.copy()
         state["_lock"] = None
-        state["_async_lock"] = None
         return state
 
     def __setstate__(self, state):
         """Restore after unpickling."""
         self.__dict__.update(state)
-        self._lock = threading.Lock()
-        self._async_lock = asyncio.Lock()
+        self._lock = ConcurrencyLock()
 
     @property
     def lock(self):
-        """Thread lock."""
+        """Unified concurrency lock for both sync and async operations."""
         if not hasattr(self, "_lock") or self._lock is None:
-            self._lock = threading.Lock()
+            self._lock = ConcurrencyLock()
         return self._lock
 
-    @property
-    def async_lock(self):
-        """Async lock."""
-        if not hasattr(self, "_async_lock") or self._async_lock is None:
-            self._async_lock = asyncio.Lock()
-        return self._async_lock
-
     # Async Interface methods
     @async_synchronized
     async def asetitem(
@@ -577,12 +554,13 @@ class Pile(Element, Collective[E], Generic[E], Adaptable, AsyncAdaptable):
 
     async def __aiter__(self) -> AsyncIterator[T]:
         """Async iterate over items."""
-        async with self.async_lock:
+        async with self.lock:
             current_order = list(self.progression)
 
         for key in current_order:
-            yield self.collections[key]
-            await asyncio.sleep(0)  # Yield control to the event loop
+            if key in self.collections:
+                yield self.collections[key]
+                await asyncio.sleep(0)  # Yield control to the event loop
 
     async def __anext__(self) -> T:
         """Async get next item."""
@@ -915,7 +893,7 @@ class Pile(Element, Collective[E], Generic[E], Adaptable, AsyncAdaptable):
 
     async def __aenter__(self) -> Self:
         """Enter async context."""
-        await self.async_lock.acquire()
+        await self.lock.__aenter__()
         return self
 
     async def __aexit__(
@@ -925,7 +903,7 @@ class Pile(Element, Collective[E], Generic[E], Adaptable, AsyncAdaptable):
         exc_tb: Any,
     ) -> None:
         """Exit async context."""
-        self.async_lock.release()
+        await self.lock.__aexit__(exc_type, exc_val, exc_tb)
 
     def is_homogenous(self) -> bool:
         """Check if all items are same type."""

lionagi/protocols/generic/processor.py

@@ -5,6 +5,9 @@
 import asyncio
 from typing import Any, ClassVar
 
+from lionagi.libs.concurrency import Event as ConcurrencyEvent
+from lionagi.libs.concurrency import Semaphore, create_task_group
+
 from .._concepts import Observer
 from .element import ID
 from .event import Event, EventStatus
@@ -56,9 +59,9 @@ class Processor(Observer):
         self.queue = asyncio.Queue()
         self._available_capacity = queue_capacity
         self._execution_mode = False
-        self._stop_event = asyncio.Event()
+        self._stop_event = ConcurrencyEvent()
         if concurrency_limit:
-            self._concurrency_sem = asyncio.Semaphore(concurrency_limit)
+            self._concurrency_sem = Semaphore(concurrency_limit)
         else:
             self._concurrency_sem = None
 
@@ -106,7 +109,9 @@ class Processor(Observer):
 
     async def start(self) -> None:
         """Clears the stop signal, allowing event processing to resume."""
-        self._stop_event.clear()
+        # Create a new event since ConcurrencyEvent doesn't have clear()
+        if self._stop_event.is_set():
+            self._stop_event = ConcurrencyEvent()
 
     def is_stopped(self) -> bool:
         """Checks whether the processor is in a stopped state.
@@ -136,30 +141,52 @@ class Processor(Observer):
         for tasks to complete. Resets capacity afterward if any events
         were processed.
         """
-        tasks = set()
         prev_event: Event | None = None
-
-        while self.available_capacity > 0 and not self.queue.empty():
-            next_event = None
-            if prev_event and prev_event.status == EventStatus.PENDING:
-                # Wait if previous event is still pending
-                await asyncio.sleep(self.capacity_refresh_time)
-                next_event = prev_event
-            else:
-                next_event = await self.dequeue()
-
-            if await self.request_permission(**next_event.request):
-                if next_event.streaming:
-                    task = asyncio.create_task(next_event.stream())
+        events_processed = 0
+
+        async with create_task_group() as tg:
+            while self.available_capacity > 0 and not self.queue.empty():
+                next_event = None
+                if prev_event and prev_event.status == EventStatus.PENDING:
+                    # Wait if previous event is still pending
+                    await asyncio.sleep(self.capacity_refresh_time)
+                    next_event = prev_event
                 else:
-                    task = asyncio.create_task(next_event.invoke())
-                tasks.add(task)
+                    next_event = await self.dequeue()
+
+                if await self.request_permission(**next_event.request):
+                    if next_event.streaming:
+                        # For streaming, we need to consume the async generator
+                        async def consume_stream(event):
+                            async for _ in event.stream():
+                                pass
+
+                        if self._concurrency_sem:
+
+                            async def stream_with_sem(event):
+                                async with self._concurrency_sem:
+                                    await consume_stream(event)
+
+                            await tg.start_soon(stream_with_sem, next_event)
+                        else:
+                            await tg.start_soon(consume_stream, next_event)
+                    else:
+                        # For non-streaming, just invoke
+                        if self._concurrency_sem:
+
+                            async def invoke_with_sem(event):
+                                async with self._concurrency_sem:
+                                    await event.invoke()
+
+                            await tg.start_soon(invoke_with_sem, next_event)
+                        else:
+                            await tg.start_soon(next_event.invoke)
+                    events_processed += 1
 
-            prev_event = next_event
-            self._available_capacity -= 1
+                prev_event = next_event
+                self._available_capacity -= 1
 
-        if tasks:
-            await asyncio.wait(tasks)
+        if events_processed > 0:
             self.available_capacity = self.queue_capacity
 
     async def request_permission(self, **kwargs: Any) -> bool:
@@ -270,9 +297,9 @@ class Executor(Observer):
         Args:
             event (Event): The event to add.
         """
-        async with self.pile:
-            self.pile.include(event)
-            self.pending.include(event)
+        # Use async methods to avoid deadlock between sync/async locks
+        await self.pile.ainclude(event)
+        self.pending.include(event)
 
     @property
     def completed_events(self) -> Pile[Event]:

lionagi/service/rate_limited_processor.py

@@ -4,9 +4,11 @@
 
 import asyncio
 import logging
+from typing import Any
 
 from typing_extensions import Self, override
 
+from lionagi.libs.concurrency import CapacityLimiter, Lock, move_on_after
 from lionagi.protocols.types import Executor, Processor
 
 from .connections.api_calling import APICalling
@@ -40,24 +42,40 @@ class RateLimitedAPIProcessor(Processor):
         self.available_request = self.limit_requests
         self.available_token = self.limit_tokens
         self._rate_limit_replenisher_task: asyncio.Task | None = None
-        self._lock: asyncio.Lock = asyncio.Lock()
-        self._concurrency_sem = asyncio.Semaphore(
-            concurrency_limit or queue_capacity
-        )
+        self._lock = Lock()
+
+        # Use CapacityLimiter for better token management
+        if self.limit_tokens:
+            self._token_limiter = CapacityLimiter(self.limit_tokens)
+        else:
+            self._token_limiter = None
+
+        if self.limit_requests:
+            self._request_limiter = CapacityLimiter(self.limit_requests)
+        else:
+            self._request_limiter = None
 
     async def start_replenishing(self):
         """Start replenishing rate limit capacities at regular intervals."""
         await self.start()
         try:
             while not self.is_stopped():
-                await asyncio.sleep(delay=self.interval)
-                async with self._lock:
-                    if self.limit_requests is not None:
-                        self.available_request = (
-                            self.limit_requests - self.queue.qsize()
+                await asyncio.sleep(self.interval)
+
+                # Reset capacity limiters to their original values
+                if self._request_limiter and self.limit_requests:
+                    # Adjust total tokens to reset capacity
+                    current_borrowed = self._request_limiter.borrowed_tokens
+                    if current_borrowed < self.limit_requests:
+                        self._request_limiter.total_tokens = (
+                            self.limit_requests
                         )
-                    if self.limit_tokens is not None:
-                        self.available_token = self.limit_tokens
+
+                if self._token_limiter and self.limit_tokens:
+                    # Reset token limiter capacity
+                    current_borrowed = self._token_limiter.borrowed_tokens
+                    if current_borrowed < self.limit_tokens:
+                        self._token_limiter.total_tokens = self.limit_tokens
 
         except asyncio.CancelledError:
             logging.info("Rate limit replenisher task cancelled.")
@@ -98,31 +116,31 @@ class RateLimitedAPIProcessor(Processor):
 
     @override
     async def request_permission(
-        self, required_tokens: int = None, **kwargs
+        self, required_tokens: int = None, **kwargs: Any
     ) -> bool:
-        async with self._lock:
-            if self.limit_requests is None and self.limit_tokens is None:
-                if self.queue.qsize() < self.queue_capacity:
-                    return True
-
-            if self.limit_requests is not None:
-                if self.available_request > 0:
-                    self.available_request -= 1
-                if required_tokens is None:
-                    return True
-                else:
-                    if self.limit_tokens >= required_tokens:
-                        self.limit_tokens -= required_tokens
-                        return True
-
-            if self.limit_tokens is not None:
-                if required_tokens is None:
-                    return True
-                if self.limit_tokens >= required_tokens:
-                    self.limit_tokens -= required_tokens
-                    return True
-
-            return False
+        # No limits configured, just check queue capacity
+        if self._request_limiter is None and self._token_limiter is None:
+            return self.queue.qsize() < self.queue_capacity
+
+        # Check request limit
+        if self._request_limiter:
+            # Try to acquire with timeout
+            with move_on_after(0.1) as scope:
+                await self._request_limiter.acquire()
+                if scope.cancelled_caught:
+                    return False
+
+        # Check token limit if required
+        if self._token_limiter and required_tokens:
+            # For token-based limiting, we need to acquire multiple tokens
+            # This is a simplified approach - in production you might want
+            # a more sophisticated token bucket algorithm
+            if self._token_limiter.available_tokens < required_tokens:
+                if self._request_limiter:
+                    self._request_limiter.release()
+                return False
+
+        return True
 
 
 class RateLimitedAPIExecutor(Executor):

lionagi/session/branch.py

@@ -1186,7 +1186,6 @@ class Branch(Element, Communicatable, Relational):
         backoff_factor: float = 1,
         retry_default: Any = UNDEFINED,
         retry_timeout: float | None = None,
-        retry_timing: bool = False,
         max_concurrent: int | None = None,
         throttle_period: float | None = None,
         flatten: bool = True,
@@ -1223,8 +1222,6 @@ class Branch(Element, Communicatable, Relational):
                 Fallback value if all retries fail (if suppressing errors).
             retry_timeout (float|None):
                 Overall timeout for all attempts (None = no limit).
-            retry_timing (bool):
-                If True, track time used for retries.
             max_concurrent (int|None):
                 Maximum concurrent tasks (if batching).
             throttle_period (float|None):
@@ -1261,7 +1258,6 @@ class Branch(Element, Communicatable, Relational):
                     backoff_factor=backoff_factor,
                     retry_default=retry_default,
                     retry_timeout=retry_timeout,
-                    retry_timing=retry_timing,
                     max_concurrent=max_concurrent,
                     throttle_period=throttle_period,
                     flatten=flatten,
@@ -1290,7 +1286,6 @@ class Branch(Element, Communicatable, Relational):
                     backoff_factor=backoff_factor,
                     retry_default=retry_default,
                     retry_timeout=retry_timeout,
-                    retry_timing=retry_timing,
                     throttle_period=throttle_period,
                     flatten=flatten,
                     dropna=dropna,