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

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

@@ -917,9 +917,7 @@ class Branch(Element, Communicatable, Relational):
         actions: bool = False,
         reason: bool = False,
         action_kwargs: dict = None,
-        action_strategy: Literal[
-            "sequential", "concurrent", "batch"
-        ] = "concurrent",
+        action_strategy: Literal["sequential", "concurrent"] = "concurrent",
         verbose_action: bool = False,
         field_models: list[FieldModel] = None,
         exclude_fields: list | dict | None = None,
@@ -987,7 +985,7 @@ class Branch(Element, Communicatable, Relational):
                 If `True`, signals that the LLM should provide chain-of-thought or reasoning (where applicable).
             action_kwargs (dict | None, optional):
                 Additional parameters for the `branch.act()` call if tools are invoked.
-            action_strategy (Literal["sequential","concurrent","batch"], optional):
+            action_strategy (Literal["sequential","concurrent"], optional):
                 The strategy for invoking tools (default: "concurrent").
             verbose_action (bool, optional):
                 If `True`, logs detailed information about tool invocation.
@@ -1174,9 +1172,8 @@ class Branch(Element, Communicatable, Relational):
         self,
         action_request: list | ActionRequest | BaseModel | dict,
         *,
-        strategy: Literal["concurrent", "sequential", "batch"] = "concurrent",
+        strategy: Literal["concurrent", "sequential"] = "concurrent",
         verbose_action: bool = False,
-        batch_size: int = None,
         suppress_errors: bool = True,
         sanitize_input: bool = False,
         unique_input: bool = False,
@@ -1186,7 +1183,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,10 +1219,8 @@ 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).
+                Maximum concurrent tasks.
             throttle_period (float|None):
                 Minimum spacing (in seconds) between requests.
             flatten (bool):
@@ -1242,11 +1236,6 @@ class Branch(Element, Communicatable, Relational):
             Any:
                 The result or results from the invoked tool(s).
         """
-        if batch_size and not strategy == "batch":
-            raise ValueError(
-                "Batch size is only applicable for 'batch' strategy."
-            )
-
         match strategy:
             case "concurrent":
                 return await self._concurrent_act(
@@ -1261,7 +1250,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,
@@ -1275,28 +1263,8 @@ class Branch(Element, Communicatable, Relational):
                     verbose_action=verbose_action,
                     suppress_errors=suppress_errors,
                 )
-            case "batch":
-                return await self._batch_act(
-                    action_request,
-                    verbose_action=verbose_action,
-                    batch_size=batch_size or 1,
-                    max_concurrent=max_concurrent,
-                    suppress_errors=suppress_errors,
-                    sanitize_input=sanitize_input,
-                    unique_input=unique_input,
-                    num_retries=num_retries,
-                    initial_delay=initial_delay,
-                    retry_delay=retry_delay,
-                    backoff_factor=backoff_factor,
-                    retry_default=retry_default,
-                    retry_timeout=retry_timeout,
-                    retry_timing=retry_timing,
-                    throttle_period=throttle_period,
-                    flatten=flatten,
-                    dropna=dropna,
-                    unique_output=unique_output,
-                    flatten_tuple_set=flatten_tuple_set,
-                )
+            case _:
+                raise
 
     async def _concurrent_act(
         self,
@@ -1327,19 +1295,6 @@ class Branch(Element, Communicatable, Relational):
             )
         return results
 
-    async def _batch_act(
-        self,
-        action_request: list[ActionRequest | BaseModel | dict],
-        batch_size: int = None,
-        **kwargs,
-    ) -> list:
-        result = []
-        async for i in bcall(
-            action_request, self._act, batch_size=batch_size, **kwargs
-        ):
-            result.extend(i)
-        return result
-
     async def translate(
         self,
         text: str,

lionagi/session/session.py

@@ -2,6 +2,7 @@
 #
 # SPDX-License-Identifier: Apache-2.0
 
+import contextlib
 from collections.abc import Callable
 from typing import Any
 
@@ -46,7 +47,7 @@ class Session(Node, Communicatable, Relational):
         mail_manager (MailManager | None): Manages mail operations.
     """
 
-    branches: Pile[Any] = Field(
+    branches: Pile[Branch] = Field(
         default_factory=lambda: Pile(item_type={Branch}, strict_type=False)
     )
     default_branch: Any = Field(default=None, exclude=True)
@@ -57,21 +58,38 @@ class Session(Node, Communicatable, Relational):
     name: str = Field(default="Session")
 
     @model_validator(mode="after")
-    def _initialize_default_branch(self) -> Self:
+    def _add_mail_sources(self) -> Self:
         if self.default_branch is None:
-            from .branch import Branch
-
             self.default_branch = Branch()
-        return self
-
-    @model_validator(mode="after")
-    def _add_mail_sources(self) -> Self:
         if self.default_branch not in self.branches:
             self.branches.include(self.default_branch)
         if self.branches:
             self.mail_manager.add_sources(self.branches)
         return self
 
+    def _lookup_branch_by_name(self, name: str) -> Branch | None:
+        for branch in self.branches:
+            if branch.name == name:
+                return branch
+        return None
+
+    def get_branch(
+        self, branch: ID.Ref | str, default: Any = ..., /
+    ) -> Branch:
+        """Get a branch by its ID or name."""
+
+        with contextlib.suppress(ItemNotFoundError, ValueError):
+            id = ID.get_id(branch)
+            return self.branches[id]
+
+        if isinstance(branch, str):
+            if b := self._lookup_branch_by_name(branch):
+                return b
+
+        if default is ...:
+            raise ItemNotFoundError(f"Branch '{branch}' not found.")
+        return default
+
     def new_branch(
         self,
         system: System | JsonValue = None,

lionagi/utils.py

@@ -43,9 +43,12 @@ from typing import (
     overload,
 )
 
+import anyio
 from pydantic import BaseModel, model_validator
 from pydantic_core import PydanticUndefinedType
 
+from .libs.concurrency import Lock as ConcurrencyLock
+from .libs.concurrency import Semaphore, create_task_group
 from .settings import Settings
 
 R = TypeVar("R")
@@ -612,32 +615,6 @@ class CallParams(Params):
         )
 
 
-class LCallParams(CallParams):
-    func: Any = None
-    sanitize_input: bool = False
-    unique_input: bool = False
-    flatten: bool = False
-    dropna: bool = False
-    unique_output: bool = False
-    flatten_tuple_set: bool = False
-
-    def __call__(self, input_: Any, func=None):
-        if self.func is None and func is None:
-            raise ValueError("a sync func must be provided")
-        return lcall(
-            input_,
-            func or self.func,
-            *self.args,
-            sanitize_input=self.sanitize_input,
-            unique_input=self.unique_input,
-            flatten=self.flatten,
-            dropna=self.dropna,
-            unique_output=self.unique_output,
-            flatten_tuple_set=self.flatten_tuple_set,
-            **self.kwargs,
-        )
-
-
 async def alcall(
     input_: list[Any],
     func: Callable[..., T],
@@ -651,7 +628,6 @@ async def alcall(
     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 = False,
@@ -659,7 +635,7 @@ async def alcall(
     unique_output: bool = False,
     flatten_tuple_set: bool = False,
     **kwargs: Any,
-) -> list[T] | list[tuple[T, float]]:
+) -> list[T]:
     """
     Asynchronously apply a function to each element of a list, with optional input sanitization,
     retries, timeout, and output processing.
@@ -675,7 +651,6 @@ async def alcall(
         backoff_factor (float): Multiplier for delay after each retry.
         retry_default (Any): Default value if all retries fail.
         retry_timeout (float | None): Timeout for each function call.
-        retry_timing (bool): If True, return (result, duration) tuples.
         max_concurrent (int | None): Maximum number of concurrent operations.
         throttle_period (float | None): Delay after each completed operation.
         flatten (bool): Flatten the final result if True.
@@ -685,7 +660,7 @@ async def alcall(
         **kwargs: Additional arguments passed to func.
 
     Returns:
-        list[T] or list[tuple[T, float]]: The processed results, or results with timing if retry_timing is True.
+        list[T]: The processed results.
 
     Raises:
         asyncio.TimeoutError: If a call times out and no default is provided.
@@ -734,9 +709,9 @@ async def alcall(
 
     # Optional initial delay before processing
     if initial_delay:
-        await asyncio.sleep(initial_delay)
+        await anyio.sleep(initial_delay)
 
-    semaphore = asyncio.Semaphore(max_concurrent) if max_concurrent else None
+    semaphore = Semaphore(max_concurrent) if max_concurrent else None
     throttle_delay = throttle_period or 0
     coro_func = is_coro_func(func)
 
@@ -744,137 +719,92 @@ async def alcall(
         if coro_func:
             # Async function
             if retry_timeout is not None:
-                return await asyncio.wait_for(
-                    func(item, **kwargs), timeout=retry_timeout
-                )
+                with anyio.move_on_after(retry_timeout) as cancel_scope:
+                    result = await func(item, **kwargs)
+                if cancel_scope.cancelled_caught:
+                    raise asyncio.TimeoutError(
+                        f"Function call timed out after {retry_timeout}s"
+                    )
+                return result
             else:
                 return await func(item, **kwargs)
         else:
             # Sync function
             if retry_timeout is not None:
-                return await asyncio.wait_for(
-                    asyncio.to_thread(func, item, **kwargs),
-                    timeout=retry_timeout,
-                )
+                with anyio.move_on_after(retry_timeout) as cancel_scope:
+                    result = await anyio.to_thread.run_sync(
+                        func, item, **kwargs
+                    )
+                if cancel_scope.cancelled_caught:
+                    raise asyncio.TimeoutError(
+                        f"Function call timed out after {retry_timeout}s"
+                    )
+                return result
             else:
-                return func(item, **kwargs)
+                return await anyio.to_thread.run_sync(func, item, **kwargs)
 
     async def execute_task(i: Any, index: int) -> Any:
-        start_time = asyncio.get_running_loop().time()
         attempts = 0
         current_delay = retry_delay
         while True:
             try:
                 result = await call_func(i)
-                if retry_timing:
-                    end_time = asyncio.get_running_loop().time()
-                    return index, result, end_time - start_time
-                else:
-                    return index, result
-            except asyncio.CancelledError as e:
-                raise e
+                return index, result
+            except anyio.get_cancelled_exc_class():
+                raise
 
             except Exception:
                 attempts += 1
                 if attempts <= num_retries:
                     if current_delay:
-                        await asyncio.sleep(current_delay)
+                        await anyio.sleep(current_delay)
                         current_delay *= backoff_factor
                     # Retry loop continues
                 else:
                     # Exhausted retries
                     if retry_default is not UNDEFINED:
-                        # Return default if provided
-                        if retry_timing:
-                            end_time = asyncio.get_running_loop().time()
-                            duration = end_time - (start_time or end_time)
-                            return index, retry_default, duration
-                        else:
-                            return index, retry_default
+                        return index, retry_default
                     # No default, re-raise
                     raise
 
     async def task_wrapper(item: Any, idx: int) -> Any:
         if semaphore:
             async with semaphore:
-                return await execute_task(item, idx)
+                result = await execute_task(item, idx)
         else:
-            return await execute_task(item, idx)
+            result = await execute_task(item, idx)
 
-    # Create tasks
-    tasks = [task_wrapper(item, idx) for idx, item in enumerate(input_)]
+        return result
 
-    # Collect results as they complete
+    # Use task group for structured concurrency
     results = []
-    for coro in asyncio.as_completed(tasks):
-        res = await coro
-        results.append(res)
-        if throttle_delay:
-            await asyncio.sleep(throttle_delay)
+    results_lock = ConcurrencyLock()  # Protect results list
 
-    # Sort by original index
-    results.sort(key=lambda x: x[0])
+    async def run_and_store(item: Any, idx: int):
+        result = await task_wrapper(item, idx)
+        async with results_lock:
+            results.append(result)
 
-    if retry_timing:
-        # (index, result, duration)
-        filtered = [
-            (r[1], r[2]) for r in results if not dropna or r[1] is not None
-        ]
-        return filtered
-    else:
-        # (index, result)
-        output_list = [r[1] for r in results]
-        return to_list(
-            output_list,
-            flatten=flatten,
-            dropna=dropna,
-            unique=unique_output,
-            flatten_tuple_set=flatten_tuple_set,
-        )
+    # Execute all tasks using task group
+    async with create_task_group() as tg:
+        for idx, item in enumerate(input_):
+            await tg.start_soon(run_and_store, item, idx)
+            # Apply throttle delay between starting tasks
+            if throttle_delay and idx < len(input_) - 1:
+                await anyio.sleep(throttle_delay)
 
+    # Sort by original index
+    results.sort(key=lambda x: x[0])
 
-class ALCallParams(CallParams):
-    func: Any = None
-    sanitize_input: bool = False
-    unique_input: bool = False
-    num_retries: int = 0
-    initial_delay: float = 0
-    retry_delay: float = 0
-    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 = False
-    dropna: bool = False
-    unique_output: bool = False
-    flatten_tuple_set: bool = False
-
-    async def __call__(self, input_: Any, func=None):
-        if self.func is None and func is None:
-            raise ValueError("a sync/async func must be provided")
-        return await alcall(
-            input_,
-            func or self.func,
-            *self.args,
-            sanitize_input=self.sanitize_input,
-            unique_input=self.unique_input,
-            num_retries=self.num_retries,
-            initial_delay=self.initial_delay,
-            retry_delay=self.retry_delay,
-            backoff_factor=self.backoff_factor,
-            retry_default=self.retry_default,
-            retry_timeout=self.retry_timeout,
-            retry_timing=self.retry_timing,
-            max_concurrent=self.max_concurrent,
-            throttle_period=self.throttle_period,
-            flatten=self.flatten,
-            dropna=self.dropna,
-            unique_output=self.unique_output,
-            flatten_tuple_set=self.flatten_tuple_set,
-            **self.kwargs,
-        )
+    # (index, result)
+    output_list = [r[1] for r in results]
+    return to_list(
+        output_list,
+        flatten=flatten,
+        dropna=dropna,
+        unique=unique_output,
+        flatten_tuple_set=flatten_tuple_set,
+    )
 
 
 async def bcall(
@@ -891,7 +821,6 @@ async def bcall(
     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 = False,
@@ -915,7 +844,6 @@ async def bcall(
             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,
@@ -926,52 +854,6 @@ async def bcall(
         )
 
 
-class BCallParams(CallParams):
-    func: Any = None
-    batch_size: int
-    sanitize_input: bool = False
-    unique_input: bool = False
-    num_retries: int = 0
-    initial_delay: float = 0
-    retry_delay: float = 0
-    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 = False
-    dropna: bool = False
-    unique_output: bool = False
-    flatten_tuple_set: bool = False
-
-    async def __call__(self, input_, func=None):
-        if self.func is None and func is None:
-            raise ValueError("a sync/async func must be provided")
-        return await bcall(
-            input_,
-            func or self.func,
-            *self.args,
-            batch_size=self.batch_size,
-            sanitize_input=self.sanitize_input,
-            unique_input=self.unique_input,
-            num_retries=self.num_retries,
-            initial_delay=self.initial_delay,
-            retry_delay=self.retry_delay,
-            backoff_factor=self.backoff_factor,
-            retry_default=self.retry_default,
-            retry_timeout=self.retry_timeout,
-            retry_timing=self.retry_timing,
-            max_concurrent=self.max_concurrent,
-            throttle_period=self.throttle_period,
-            flatten=self.flatten,
-            dropna=self.dropna,
-            unique_output=self.unique_output,
-            flatten_tuple_set=self.flatten_tuple_set,
-            **self.kwargs,
-        )
-
-
 def create_path(
     directory: Path | str,
     filename: str,

lionagi/version.py

@@ -1 +1 @@
-__version__ = "0.14.4"
+__version__ = "0.14.6"