zrb 1.8.10__py3-none-any.whl → 1.21.29__py3-none-any.whl

zrb/config/llm_rate_limitter.py

@@ -0,0 +1,190 @@
+import asyncio
+import json
+import time
+from collections import deque
+from typing import Any, Callable
+
+from zrb.config.config import CFG
+
+
+class LLMRateLimitter:
+    """
+    Helper class to enforce LLM API rate limits and throttling.
+    Tracks requests and tokens in a rolling 60-second window.
+    """
+
+    def __init__(
+        self,
+        max_requests_per_minute: int | None = None,
+        max_tokens_per_minute: int | None = None,
+        max_tokens_per_request: int | None = None,
+        max_tokens_per_tool_call_result: int | None = None,
+        throttle_sleep: float | None = None,
+        use_tiktoken: bool | None = None,
+        tiktoken_encoding_name: str | None = None,
+    ):
+        self._max_requests_per_minute = max_requests_per_minute
+        self._max_tokens_per_minute = max_tokens_per_minute
+        self._max_tokens_per_request = max_tokens_per_request
+        self._max_tokens_per_tool_call_result = max_tokens_per_tool_call_result
+        self._throttle_sleep = throttle_sleep
+        self._use_tiktoken = use_tiktoken
+        self._tiktoken_encoding_name = tiktoken_encoding_name
+        self.request_times = deque()
+        self.token_times = deque()
+
+    @property
+    def max_requests_per_minute(self) -> int:
+        if self._max_requests_per_minute is not None:
+            return self._max_requests_per_minute
+        return CFG.LLM_MAX_REQUESTS_PER_MINUTE
+
+    @property
+    def max_tokens_per_minute(self) -> int:
+        if self._max_tokens_per_minute is not None:
+            return self._max_tokens_per_minute
+        return CFG.LLM_MAX_TOKENS_PER_MINUTE
+
+    @property
+    def max_tokens_per_request(self) -> int:
+        if self._max_tokens_per_request is not None:
+            return self._max_tokens_per_request
+        return CFG.LLM_MAX_TOKENS_PER_REQUEST
+
+    @property
+    def max_tokens_per_tool_call_result(self) -> int:
+        if self._max_tokens_per_tool_call_result is not None:
+            return self._max_tokens_per_tool_call_result
+        return CFG.LLM_MAX_TOKENS_PER_TOOL_CALL_RESULT
+
+    @property
+    def throttle_sleep(self) -> float:
+        if self._throttle_sleep is not None:
+            return self._throttle_sleep
+        return CFG.LLM_THROTTLE_SLEEP
+
+    @property
+    def use_tiktoken(self) -> bool:
+        if self._use_tiktoken is not None:
+            return self._use_tiktoken
+        return CFG.USE_TIKTOKEN
+
+    @property
+    def tiktoken_encoding_name(self) -> str:
+        if self._tiktoken_encoding_name is not None:
+            return self._tiktoken_encoding_name
+        return CFG.TIKTOKEN_ENCODING_NAME
+
+    def set_max_requests_per_minute(self, value: int):
+        self._max_requests_per_minute = value
+
+    def set_max_tokens_per_minute(self, value: int):
+        self._max_tokens_per_minute = value
+
+    def set_max_tokens_per_request(self, value: int):
+        self._max_tokens_per_request = value
+
+    def set_max_tokens_per_tool_call_result(self, value: int):
+        self._max_tokens_per_tool_call_result = value
+
+    def set_throttle_sleep(self, value: float):
+        self._throttle_sleep = value
+
+    def count_token(self, prompt: Any) -> int:
+        str_prompt = self._prompt_to_str(prompt)
+        if not self.use_tiktoken:
+            return self._fallback_count_token(str_prompt)
+        try:
+            import tiktoken
+
+            enc = tiktoken.get_encoding(self.tiktoken_encoding_name)
+            return len(enc.encode(str_prompt))
+        except Exception:
+            return self._fallback_count_token(str_prompt)
+
+    def _fallback_count_token(self, str_prompt: str) -> int:
+        return len(str_prompt) // 4
+
+    def clip_prompt(self, prompt: Any, limit: int) -> str:
+        str_prompt = self._prompt_to_str(prompt)
+        if not self.use_tiktoken:
+            return self._fallback_clip_prompt(str_prompt, limit)
+        try:
+            import tiktoken
+
+            enc = tiktoken.get_encoding(self.tiktoken_encoding_name)
+            tokens = enc.encode(str_prompt)
+            if len(tokens) <= limit:
+                return str_prompt
+            truncated = tokens[: limit - 3]
+            clipped_text = enc.decode(truncated)
+            return clipped_text + "..."
+        except Exception:
+            return self._fallback_clip_prompt(str_prompt, limit)
+
+    def _fallback_clip_prompt(self, str_prompt: str, limit: int) -> str:
+        char_limit = limit * 4 if limit * 4 <= 10 else limit * 4 - 10
+        return str_prompt[:char_limit] + "..."
+
+    async def throttle(
+        self,
+        prompt: Any,
+        throttle_notif_callback: Callable[[str], Any] | None = None,
+    ):
+        now = time.time()
+        str_prompt = self._prompt_to_str(prompt)
+        tokens = self.count_token(str_prompt)
+        # Clean up old entries
+        while self.request_times and now - self.request_times[0] > 60:
+            self.request_times.popleft()
+        while self.token_times and now - self.token_times[0][0] > 60:
+            self.token_times.popleft()
+        # Check per-request token limit
+        if tokens > self.max_tokens_per_request:
+            raise ValueError(
+                (
+                    "Request exceeds max_tokens_per_request "
+                    f"({tokens} > {self.max_tokens_per_request})."
+                )
+            )
+        if tokens > self.max_tokens_per_minute:
+            raise ValueError(
+                (
+                    "Request exceeds max_tokens_per_minute "
+                    f"({tokens} > {self.max_tokens_per_minute})."
+                )
+            )
+        # Wait if over per-minute request or token limit
+        while (
+            len(self.request_times) >= self.max_requests_per_minute
+            or sum(t for _, t in self.token_times) + tokens > self.max_tokens_per_minute
+        ):
+            if throttle_notif_callback is not None:
+                if len(self.request_times) >= self.max_requests_per_minute:
+                    rpm = len(self.request_times)
+                    throttle_notif_callback(
+                        f"Max request per minute exceeded: {rpm} of {self.max_requests_per_minute}"
+                    )
+                else:
+                    tpm = sum(t for _, t in self.token_times) + tokens
+                    throttle_notif_callback(
+                        f"Max token per minute exceeded: {tpm} of {self.max_tokens_per_minute}"
+                    )
+            await asyncio.sleep(self.throttle_sleep)
+            now = time.time()
+            while self.request_times and now - self.request_times[0] > 60:
+                self.request_times.popleft()
+            while self.token_times and now - self.token_times[0][0] > 60:
+                self.token_times.popleft()
+        # Record this request
+        self.request_times.append(now)
+        self.token_times.append((now, tokens))
+
+    def _prompt_to_str(self, prompt: Any) -> str:
+        try:
+            return json.dumps(prompt)
+        except Exception:
+            return f"{prompt}"
+
+
+llm_rate_limitter = LLMRateLimitter()

zrb/{runner → config}/web_auth_config.py

@@ -1,14 +1,15 @@
-from typing import Callable
+from typing import TYPE_CHECKING, Callable
 
-from zrb.config import CFG
-from zrb.runner.web_schema.user import User
+from zrb.config.config import CFG
 from zrb.task.any_task import AnyTask
 
+if TYPE_CHECKING:
+    from zrb.runner.web_schema.user import User
+
 
 class WebAuthConfig:
     def __init__(
         self,
-        port: int | None = None,
         secret_key: str | None = None,
         access_token_expire_minutes: int | None = None,
         refresh_token_expire_minutes: int | None = None,
@@ -19,9 +20,8 @@ class WebAuthConfig:
         super_admin_password: str | None = None,
         guest_username: str | None = None,
         guest_accessible_tasks: list[AnyTask | str] = [],
-        find_user_by_username: Callable[[str], User | None] | None = None,
+        find_user_by_username: Callable[[str], "User | None"] | None = None,
     ):
-        self._port = port
         self._secret_key = secret_key
         self._access_token_expire_minutes = access_token_expire_minutes
         self._refresh_token_expire_minutes = refresh_token_expire_minutes
@@ -31,16 +31,10 @@ class WebAuthConfig:
         self._super_admin_username = super_admin_username
         self._super_admin_password = super_admin_password
         self._guest_username = guest_username
-        self._user_list = []
+        self._user_list: list["User"] = []
         self._guest_accessible_tasks = guest_accessible_tasks
         self._find_user_by_username = find_user_by_username
 
-    @property
-    def port(self) -> int:
-        if self._port is not None:
-            return self._port
-        return CFG.WEB_HTTP_PORT
-
     @property
     def secret_key(self) -> str:
         if self._secret_key is not None:
@@ -100,7 +94,9 @@ class WebAuthConfig:
         return self._guest_accessible_tasks
 
     @property
-    def default_user(self) -> User:
+    def default_user(self) -> "User":
+        from zrb.runner.web_schema.user import User
+
         if self.enable_auth:
             return User(
                 username=self.guest_username,
@@ -116,7 +112,9 @@ class WebAuthConfig:
         )
 
     @property
-    def super_admin(self) -> User:
+    def super_admin(self) -> "User":
+        from zrb.runner.web_schema.user import User
+
         return User(
             username=self.super_admin_username,
             password=self.super_admin_password,
@@ -124,14 +122,11 @@ class WebAuthConfig:
         )
 
     @property
-    def user_list(self) -> list[User]:
+    def user_list(self) -> list["User"]:
         if not self.enable_auth:
             return [self.default_user]
         return self._user_list + [self.super_admin, self.default_user]
 
-    def set_port(self, port: int):
-        self._port = port
-
     def set_secret_key(self, secret_key: str):
         self._secret_key = secret_key
 
@@ -163,11 +158,11 @@ class WebAuthConfig:
         self._guest_accessible_tasks = tasks
 
     def set_find_user_by_username(
-        self, find_user_by_username: Callable[[str], User | None]
+        self, find_user_by_username: Callable[[str], "User | None"]
     ):
         self._find_user_by_username = find_user_by_username
 
-    def append_user(self, user: User):
+    def append_user(self, user: "User"):
         duplicates = [
             existing_user
             for existing_user in self.user_list
@@ -177,7 +172,7 @@ class WebAuthConfig:
             raise ValueError(f"User already exists {user.username}")
         self._user_list.append(user)
 
-    def find_user_by_username(self, username: str) -> User | None:
+    def find_user_by_username(self, username: str) -> "User | None":
         user = None
         if self._find_user_by_username is not None:
             user = self._find_user_by_username(username)

zrb/context/any_shared_context.py

@@ -19,26 +19,42 @@ class AnySharedContext(ABC):
     """
 
     @property
+    @abstractmethod
+    def is_web_mode(self) -> bool:
+        pass
+
+    @property
+    @abstractmethod
+    def is_tty(self) -> bool:
+        pass
+
+    @property
+    @abstractmethod
     def input(self) -> DotDict:
         pass
 
     @property
+    @abstractmethod
     def env(self) -> DotDict:
         pass
 
     @property
+    @abstractmethod
     def args(self) -> list[Any]:
         pass
 
     @property
-    def xcom(self) -> DotDict[str, Xcom]:
+    @abstractmethod
+    def xcom(self) -> DotDict:
         pass
 
     @property
+    @abstractmethod
     def shared_log(self) -> list[str]:
         pass
 
     @property
+    @abstractmethod
     def session(self) -> any_session.AnySession | None:
         pass
 

zrb/context/context.py

@@ -33,6 +33,14 @@ class Context(AnyContext):
         class_name = self.__class__.__name__
         return f"<{class_name} shared_ctx={self._shared_ctx}>"
 
+    @property
+    def is_web_mode(self) -> bool:
+        return self._shared_ctx.is_web_mode
+
+    @property
+    def is_tty(self) -> bool:
+        return self._shared_ctx.is_tty
+
     @property
     def input(self) -> DotDict:
         return self._shared_ctx.input
@@ -55,7 +63,7 @@ class Context(AnyContext):
 
     @property
     def session(self) -> AnySession | None:
-        return self._shared_ctx._session
+        return self._shared_ctx.session
 
     def update_task_env(self, task_env: dict[str, str]):
         self._env.update(task_env)
@@ -111,7 +119,13 @@ class Context(AnyContext):
             return
         color = self._color
         icon = self._icon
-        max_name_length = max(len(name) + len(icon) for name in self.session.task_names)
+        # Handle case where session is None (e.g., in tests)
+        if self.session is None:
+            max_name_length = len(self._task_name) + len(icon)
+        else:
+            max_name_length = max(
+                len(name) + len(icon) for name in self.session.task_names
+            )
         styled_task_name = f"{icon} {self._task_name}"
         padded_styled_task_name = styled_task_name.rjust(max_name_length + 1)
         if self._attempt == 0:

zrb/context/shared_context.py

@@ -1,7 +1,8 @@
 import datetime
+import sys
 from typing import Any
 
-from zrb.config import CFG
+from zrb.config.config import CFG
 from zrb.context.any_shared_context import AnySharedContext
 from zrb.dot_dict.dot_dict import DotDict
 from zrb.session.any_session import AnySession
@@ -26,6 +27,7 @@ class SharedContext(AnySharedContext):
         env: dict[str, str] = {},
         xcom: dict[str, Xcom] = {},
         logging_level: int | None = None,
+        is_web_mode: bool = False,
     ):
         self.__logging_level = logging_level
         self._input = DotDict(input)
@@ -34,14 +36,22 @@ class SharedContext(AnySharedContext):
         self._xcom = DotDict(xcom)
         self._session: AnySession | None = None
         self._log = []
+        self._is_web_mode = is_web_mode
 
     def __repr__(self):
         class_name = self.__class__.__name__
-        input = self._input
-        args = self._args
-        env = self._env
-        xcom = self._xcom
-        return f"<{class_name} input={input} args={args} xcom={xcom} env={env}>"
+        return f"<{class_name}>"
+
+    @property
+    def is_web_mode(self) -> bool:
+        return self._is_web_mode
+
+    @property
+    def is_tty(self) -> bool:
+        try:
+            return sys.stdin.isatty()
+        except Exception:
+            return False
 
     @property
     def input(self) -> DotDict:
@@ -56,7 +66,7 @@ class SharedContext(AnySharedContext):
         return self._args
 
     @property
-    def xcom(self) -> DotDict[str, Xcom]:
+    def xcom(self) -> DotDict:
         return self._xcom
 
     @property
@@ -71,7 +81,7 @@ class SharedContext(AnySharedContext):
         self._log.append(message)
         session = self.session
         if session is not None:
-            session_parent: AnySession = session.parent
+            session_parent: AnySession | None = session.parent
             if session_parent is not None:
                 session_parent.shared_ctx.append_to_shared_log(message)
 

zrb/group/any_group.py

@@ -1,5 +1,4 @@
 from abc import ABC, abstractmethod
-from typing import Optional, Union
 
 from zrb.task.any_task import AnyTask
 
@@ -31,16 +30,24 @@ class AnyGroup(ABC):
 
     @property
     @abstractmethod
-    def subgroups(self) -> dict[str, "AnyGroup"]:
+    def subgroups(self) -> "dict[str, AnyGroup]":
         """Group subgroups"""
         pass
 
     @abstractmethod
-    def add_group(self, group: Union["AnyGroup", str]) -> "AnyGroup":
+    def add_group(self, group: "AnyGroup", alias: str | None = None) -> "AnyGroup":
         pass
 
     @abstractmethod
-    def add_task(self, task: AnyTask, alias: str | None = None) -> AnyTask:
+    def add_task(self, task: "AnyTask", alias: str | None = None) -> "AnyTask":
+        pass
+
+    @abstractmethod
+    def remove_group(self, group: "AnyGroup | str"):
+        pass
+
+    @abstractmethod
+    def remove_task(self, task: "AnyTask | str"):
         pass
 
     @abstractmethod
@@ -48,5 +55,5 @@ class AnyGroup(ABC):
         pass
 
     @abstractmethod
-    def get_group_by_alias(self, name: str) -> Optional["AnyGroup"]:
+    def get_group_by_alias(self, alias: str) -> "AnyGroup | None":
         pass

zrb/group/group.py

@@ -33,15 +33,15 @@ class Group(AnyGroup):
     def subgroups(self) -> dict[str, AnyGroup]:
         names = list(self._groups.keys())
         names.sort()
-        return {name: self._groups.get(name) for name in names}
+        return {name: self._groups[name] for name in names}
 
     @property
     def subtasks(self) -> dict[str, AnyTask]:
         alias = list(self._tasks.keys())
         alias.sort()
-        return {name: self._tasks.get(name) for name in alias}
+        return {name: self._tasks[name] for name in alias}
 
-    def add_group(self, group: AnyGroup | str, alias: str | None = None) -> AnyGroup:
+    def add_group(self, group: AnyGroup, alias: str | None = None) -> AnyGroup:
         real_group = Group(group) if isinstance(group, str) else group
         alias = alias if alias is not None else real_group.name
         self._groups[alias] = real_group
@@ -52,6 +52,70 @@ class Group(AnyGroup):
         self._tasks[alias] = task
         return task
 
+    def remove_group(self, group: "AnyGroup | str"):
+        original_groups_len = len(self._groups)
+        if isinstance(group, AnyGroup):
+            new_groups = {
+                alias: existing_group
+                for alias, existing_group in self._groups.items()
+                if group != existing_group
+            }
+            if len(new_groups) == original_groups_len:
+                raise ValueError(f"Cannot remove group {group} from {self}")
+            self._groups = new_groups
+            return
+        # group is string, try to remove by alias
+        new_groups = {
+            alias: existing_group
+            for alias, existing_group in self._groups.items()
+            if alias != group
+        }
+        if len(new_groups) < original_groups_len:
+            self._groups = new_groups
+            return
+        # if alias removal didn't work, try to remove by name
+        new_groups = {
+            alias: existing_group
+            for alias, existing_group in self._groups.items()
+            if existing_group.name != group
+        }
+        if len(new_groups) < original_groups_len:
+            self._groups = new_groups
+            return
+        raise ValueError(f"Cannot remove group {group} from {self}")
+
+    def remove_task(self, task: "AnyTask | str"):
+        original_tasks_len = len(self._tasks)
+        if isinstance(task, AnyTask):
+            new_tasks = {
+                alias: existing_task
+                for alias, existing_task in self._tasks.items()
+                if task != existing_task
+            }
+            if len(new_tasks) == original_tasks_len:
+                raise ValueError(f"Cannot remove task {task} from {self}")
+            self._tasks = new_tasks
+            return
+        # task is string, try to remove by alias
+        new_tasks = {
+            alias: existing_task
+            for alias, existing_task in self._tasks.items()
+            if alias != task
+        }
+        if len(new_tasks) < original_tasks_len:
+            self._tasks = new_tasks
+            return
+        # if alias removal didn't work, try to remove by name
+        new_tasks = {
+            alias: existing_task
+            for alias, existing_task in self._tasks.items()
+            if existing_task.name != task
+        }
+        if len(new_tasks) < original_tasks_len:
+            self._tasks = new_tasks
+            return
+        raise ValueError(f"Cannot remove task {task} from {self}")
+
     def get_task_by_alias(self, alias: str) -> AnyTask | None:
         return self._tasks.get(alias)
 

zrb/input/any_input.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from typing import Any
 
 from zrb.context.any_shared_context import AnySharedContext
 
@@ -35,7 +36,10 @@ class AnyInput(ABC):
 
     @abstractmethod
     def update_shared_context(
-        self, shared_ctx: AnySharedContext, str_value: str | None = None
+        self,
+        shared_ctx: AnySharedContext,
+        str_value: str | None = None,
+        value: Any = None,
     ):
         pass
 

zrb/input/base_input.py

@@ -58,11 +58,15 @@ class BaseInput(AnyInput):
         return f'<input name="{name}" placeholder="{description}" value="{default}" />'
 
     def update_shared_context(
-        self, shared_ctx: AnySharedContext, str_value: str | None = None
+        self,
+        shared_ctx: AnySharedContext,
+        str_value: str | None = None,
+        value: Any = None,
     ):
-        if str_value is None:
-            str_value = self.get_default_str(shared_ctx)
-        value = self._parse_str_value(str_value)
+        if value is None:
+            if str_value is None:
+                str_value = self.get_default_str(shared_ctx)
+            value = self._parse_str_value(str_value)
         if self.name in shared_ctx.input:
             raise ValueError(f"Input already defined in the context: {self.name}")
         shared_ctx.input[self.name] = value
@@ -91,12 +95,20 @@ class BaseInput(AnyInput):
         default_str = self.get_default_str(shared_ctx)
         if default_str != "":
             prompt_message = f"{prompt_message} [{default_str}]"
-        print(f"{prompt_message}: ", end="")
-        value = input()
+        value = self._read_line(shared_ctx, prompt_message)
         if value.strip() == "":
             value = default_str
         return value
 
+    def _read_line(self, shared_ctx: AnySharedContext, prompt_message: str) -> str:
+        if not shared_ctx.is_tty:
+            print(f"{prompt_message}: ", end="")
+            return input()
+        from prompt_toolkit import PromptSession
+
+        reader = PromptSession()
+        return reader.prompt(f"{prompt_message}: ")
+
     def get_default_str(self, shared_ctx: AnySharedContext) -> str:
         """Get default value as str"""
         default_value = get_attr(

zrb/input/option_input.py

@@ -47,9 +47,21 @@ class OptionInput(BaseInput):
         option_str = ", ".join(options)
         if default_value != "":
             prompt_message = f"{prompt_message} ({option_str}) [{default_value}]"
-        value = input(f"{prompt_message}: ")
+        value = self._get_value_from_user_input(shared_ctx, prompt_message, options)
         if value.strip() != "" and value.strip() not in options:
             value = self._prompt_cli_str(shared_ctx)
         if value.strip() == "":
             value = default_value
         return value
+
+    def _get_value_from_user_input(
+        self, shared_ctx: AnySharedContext, prompt_message: str, options: list[str]
+    ) -> str:
+        from prompt_toolkit import PromptSession
+        from prompt_toolkit.completion import WordCompleter
+
+        if shared_ctx.is_tty:
+            reader = PromptSession()
+            option_completer = WordCompleter(options, ignore_case=True)
+            return reader.prompt(f"{prompt_message}: ", completer=option_completer)
+        return input(f"{prompt_message}: ")