zrb 1.13.1__py3-none-any.whl → 1.21.17__py3-none-any.whl

zrb/config/llm_rate_limitter.py

@@ -1,29 +1,12 @@
 import asyncio
+import json
 import time
 from collections import deque
-from typing import Callable
-
-import tiktoken
+from typing import Any, Callable
 
 from zrb.config.config import CFG
 
 
-def _estimate_token(text: str) -> int:
-    """
-    Estimates the number of tokens in a given text.
-    Tries to use the 'gpt-4o' model's tokenizer for an accurate count.
-    If the tokenizer is unavailable (e.g., due to network issues),
-    it falls back to a heuristic of 4 characters per token.
-    """
-    try:
-        # Primary method: Use tiktoken for an accurate count
-        enc = tiktoken.encoding_for_model("gpt-4o")
-        return len(enc.encode(text))
-    except Exception:
-        # Fallback method: Heuristic (4 characters per token)
-        return len(text) // 4
-
-
 class LLMRateLimiter:
     """
     Helper class to enforce LLM API rate limits and throttling.
@@ -35,14 +18,18 @@ class LLMRateLimiter:
         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,
-        token_counter_fn: Callable[[str], int] | 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._token_counter_fn = token_counter_fn
+        self._use_tiktoken = use_tiktoken
+        self._tiktoken_encoding_name = tiktoken_encoding_name
         self.request_times = deque()
         self.token_times = deque()
 
@@ -64,6 +51,12 @@ class LLMRateLimiter:
             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:
@@ -71,10 +64,16 @@ class LLMRateLimiter:
         return CFG.LLM_THROTTLE_SLEEP
 
     @property
-    def count_token(self) -> Callable[[str], int]:
-        if self._token_counter_fn is not None:
-            return self._token_counter_fn
-        return _estimate_token
+    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
@@ -85,29 +84,56 @@ class LLMRateLimiter:
     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 set_token_counter_fn(self, fn: Callable[[str], int]):
-        self._token_counter_fn = fn
-
-    def clip_prompt(self, prompt: str, limit: int) -> str:
-        token_count = self.count_token(prompt)
-        if token_count <= limit:
-            return prompt
-        while token_count > limit:
-            prompt_parts = prompt.split(" ")
-            last_part_index = len(prompt_parts) - 2
-            clipped_prompt = " ".join(prompt_parts[:last_part_index])
-            clipped_prompt += "(Content clipped...)"
-            token_count = self.count_token(clipped_prompt)
-            if token_count < limit:
-                return clipped_prompt
-        return prompt[:limit]
-
-    async def throttle(self, prompt: str):
+    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 | None = None,
+    ):
         now = time.time()
-        tokens = self.count_token(prompt)
+        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()
@@ -116,13 +142,25 @@ class LLMRateLimiter:
         # Check per-request token limit
         if tokens > self.max_tokens_per_request:
             raise ValueError(
-                f"Request exceeds max_tokens_per_request ({self.max_tokens_per_request})."
+                (
+                    "Request exceeds max_tokens_per_request "
+                    "({tokens} > {self.max_tokens_per_request})."
+                )
+            )
+        if tokens > self.max_tokens_per_minute:
+            raise ValueError(
+                (
+                    "Request exceeds max_tokens_per_minute "
+                    "({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:
+                throttle_notif_callback()
             await asyncio.sleep(self.throttle_sleep)
             now = time.time()
             while self.request_times and now - self.request_times[0] > 60:
@@ -133,5 +171,11 @@ class LLMRateLimiter:
         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 = LLMRateLimiter()

zrb/context/any_shared_context.py

@@ -29,26 +29,32 @@ class AnySharedContext(ABC):
         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

@@ -63,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)
@@ -119,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

@@ -27,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)
@@ -35,18 +36,15 @@ 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.env.get("_ZRB_IS_WEB_MODE", "0") == "1"
+        return self._is_web_mode
 
     @property
     def is_tty(self) -> bool:
@@ -68,7 +66,7 @@ class SharedContext(AnySharedContext):
         return self._args
 
     @property
-    def xcom(self) -> DotDict[str, Xcom]:
+    def xcom(self) -> DotDict:
         return self._xcom
 
     @property
@@ -83,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/text_input.py

@@ -1,12 +1,9 @@
-import os
-import subprocess
-import tempfile
 from collections.abc import Callable
 
 from zrb.config.config import CFG
 from zrb.context.any_shared_context import AnySharedContext
 from zrb.input.base_input import BaseInput
-from zrb.util.file import read_file
+from zrb.util.cli.text import edit_text
 
 
 class TextInput(BaseInput):
@@ -85,24 +82,10 @@ class TextInput(BaseInput):
         comment_prompt_message = (
             f"{self.comment_start}{prompt_message}{self.comment_end}"
         )
-        comment_prompt_message_eol = f"{comment_prompt_message}\n"
         default_value = self.get_default_str(shared_ctx)
-        with tempfile.NamedTemporaryFile(
-            delete=False, suffix=self._extension
-        ) as temp_file:
-            temp_file_name = temp_file.name
-            temp_file.write(comment_prompt_message_eol.encode())
-            # Pre-fill with default content
-            if default_value:
-                temp_file.write(default_value.encode())
-            temp_file.flush()
-            subprocess.call([self.editor_cmd, temp_file_name])
-            # Read the edited content
-            edited_content = read_file(temp_file_name)
-        parts = [
-            text.strip() for text in edited_content.split(comment_prompt_message, 1)
-        ]
-        edited_content = "\n".join(parts).lstrip()
-        os.remove(temp_file_name)
-        print(f"{prompt_message}: {edited_content}")
-        return edited_content
+        return edit_text(
+            prompt_message=comment_prompt_message,
+            value=default_value,
+            editor=self.editor_cmd,
+            extension=self._extension,
+        )

zrb/runner/cli.py

@@ -7,7 +7,7 @@ from zrb.context.any_context import AnyContext
 from zrb.context.shared_context import SharedContext
 from zrb.group.any_group import AnyGroup
 from zrb.group.group import Group
-from zrb.runner.common_util import get_run_kwargs
+from zrb.runner.common_util import get_task_str_kwargs
 from zrb.session.session import Session
 from zrb.session_state_logger.session_state_logger_factory import session_state_logger
 from zrb.task.any_task import AnyTask
@@ -38,23 +38,25 @@ class Cli(Group):
     def banner(self) -> str:
         return CFG.BANNER
 
-    def run(self, args: list[str] = []):
-        kwargs, args = self._extract_kwargs_from_args(args)
-        node, node_path, args = extract_node_from_args(self, args)
+    def run(self, str_args: list[str] = []):
+        str_kwargs, str_args = self._extract_kwargs_from_args(str_args)
+        node, node_path, str_args = extract_node_from_args(self, str_args)
         if isinstance(node, AnyGroup):
             self._show_group_info(node)
             return
-        if "h" in kwargs or "help" in kwargs:
+        if "h" in str_kwargs or "help" in str_kwargs:
             self._show_task_info(node)
             return
-        run_kwargs = get_run_kwargs(task=node, args=args, kwargs=kwargs, cli_mode=True)
+        task_str_kwargs = get_task_str_kwargs(
+            task=node, str_args=str_args, str_kwargs=str_kwargs, cli_mode=True
+        )
         try:
-            result = self._run_task(node, args, run_kwargs)
+            result = self._run_task(node, str_args, task_str_kwargs)
             if result is not None:
                 print(result)
             return result
         finally:
-            run_command = self._get_run_command(node_path, run_kwargs)
+            run_command = self._get_run_command(node_path, task_str_kwargs)
             self._print_run_command(run_command)
 
     def _print_run_command(self, run_command: str):
@@ -64,11 +66,14 @@ class Cli(Group):
             file=sys.stderr,
         )
 
-    def _get_run_command(self, node_path: list[str], run_kwargs: dict[str, str]) -> str:
+    def _get_run_command(
+        self, node_path: list[str], task_str_kwargs: dict[str, str]
+    ) -> str:
         parts = [self.name] + node_path
-        if len(run_kwargs) > 0:
+        if len(task_str_kwargs) > 0:
             parts += [
-                self._get_run_command_param(key, val) for key, val in run_kwargs.items()
+                self._get_run_command_param(key, val)
+                for key, val in task_str_kwargs.items()
             ]
         return " ".join(parts)
 
@@ -81,13 +86,9 @@ class Cli(Group):
         self, task: AnyTask, args: list[str], run_kwargs: dict[str, str]
     ) -> tuple[Any]:
         shared_ctx = SharedContext(args=args)
-        for task_input in task.inputs:
-            if task_input.name in run_kwargs:
-                task_input.update_shared_context(
-                    shared_ctx, run_kwargs[task_input.name]
-                )
-                continue
-        return task.run(Session(shared_ctx=shared_ctx, root_group=self))
+        return task.run(
+            Session(shared_ctx=shared_ctx, root_group=self), str_kwargs=run_kwargs
+        )
 
     def _show_task_info(self, task: AnyTask):
         description = task.description
@@ -150,11 +151,11 @@ class Cli(Group):
                         kwargs[key] = args[i + 1]
                         i += 1  # Skip the next argument as it's a value
                     else:
-                        kwargs[key] = True
+                        kwargs[key] = "true"
             elif arg.startswith("-"):
                 # Handle short flags like -t or -n
                 key = arg[1:]
-                kwargs[key] = True
+                kwargs[key] = "true"
             else:
                 # Anything else is considered a positional argument
                 residual_args.append(arg)