zrb 1.15.3__py3-none-any.whl → 2.0.0a4__py3-none-any.whl

zrb/config/helper.py

@@ -0,0 +1,84 @@
+import logging
+import os
+import platform
+
+
+def get_env(env_name: str | list[str], default: str = "", prefix: str = "ZRB") -> str:
+    env_name_list = env_name if isinstance(env_name, list) else [env_name]
+    for name in env_name_list:
+        value = os.getenv(f"{prefix}_{name}", None)
+        if value is not None:
+            return value
+    return default
+
+
+def get_current_shell() -> str:
+    if platform.system() == "Windows":
+        return "PowerShell"
+    current_shell = os.getenv("SHELL", "")
+    if current_shell.endswith("zsh"):
+        return "zsh"
+    return "bash"
+
+
+def get_default_diff_edit_command(editor: str) -> str:
+    if editor in [
+        "code",
+        "vscode",
+        "vscodium",
+        "windsurf",
+        "cursor",
+        "zed",
+        "zeditor",
+        "agy",
+    ]:
+        return f"{editor} --wait --diff {{old}} {{new}}"
+    if editor == "emacs":
+        return 'emacs --eval \'(ediff-files "{old}" "{new}")\''
+    if editor in ["nvim", "vim"]:
+        return (
+            f"{editor} -d {{old}} {{new}} "
+            "-i NONE "
+            '-c "wincmd h | set readonly | wincmd l" '
+            '-c "highlight DiffAdd cterm=bold ctermbg=22 guibg=#005f00 | highlight DiffChange cterm=bold ctermbg=24 guibg=#005f87 | highlight DiffText ctermbg=21 guibg=#0000af | highlight DiffDelete ctermbg=52 guibg=#5f0000" '  # noqa
+            '-c "set showtabline=2 | set tabline=[Instructions]\\ :wqa(save\\ &\\ quit)\\ \\|\\ i/esc(toggle\\ edit\\ mode)" '  # noqa
+            '-c "wincmd h | setlocal statusline=OLD\\ FILE" '
+            '-c "wincmd l | setlocal statusline=%#StatusBold#NEW\\ FILE\\ :wqa(save\\ &\\ quit)\\ \\|\\ i/esc(toggle\\ edit\\ mode)" '  # noqa
+            '-c "autocmd BufWritePost * wqa"'
+        )
+    return 'vimdiff {old} {new} +"setlocal ro" +"wincmd l" +"autocmd BufWritePost <buffer> qa"'  # noqa
+
+
+def get_log_level(level: str) -> int:
+    level = level.upper()
+    log_levels = {
+        "CRITICAL": logging.CRITICAL,  # 50
+        "FATAL": logging.CRITICAL,  # 50
+        "ERROR": logging.ERROR,  # 40
+        "WARN": logging.WARNING,  # 30
+        "WARNING": logging.WARNING,  # 30
+        "INFO": logging.INFO,  # 20
+        "DEBUG": logging.DEBUG,  # 10
+        "NOTSET": logging.NOTSET,  # 0
+    }
+    if level in log_levels:
+        return log_levels[level]
+    return logging.WARNING
+
+
+def get_max_token_threshold(
+    factor: float, max_tokens_per_minute: int, max_tokens_per_request: int
+) -> int:
+    return round(factor * min(max_tokens_per_minute, max_tokens_per_request))
+
+
+def limit_token_threshold(
+    threshold: int,
+    factor: float,
+    max_tokens_per_minute: int,
+    max_tokens_per_request: int,
+) -> int:
+    return min(
+        threshold,
+        get_max_token_threshold(factor, max_tokens_per_minute, max_tokens_per_request),
+    )

zrb/config/web_auth_config.py

@@ -41,58 +41,108 @@ class WebAuthConfig:
             return self._secret_key
         return CFG.WEB_SECRET_KEY
 
+    @secret_key.setter
+    def secret_key(self, secret_key: str):
+        self._secret_key = secret_key
+
     @property
     def access_token_expire_minutes(self) -> int:
         if self._access_token_expire_minutes is not None:
             return self._access_token_expire_minutes
         return CFG.WEB_AUTH_ACCESS_TOKEN_EXPIRE_MINUTES
 
+    @access_token_expire_minutes.setter
+    def access_token_expire_minutes(self, minutes: int):
+        self._access_token_expire_minutes = minutes
+
     @property
     def refresh_token_expire_minutes(self) -> int:
         if self._refresh_token_expire_minutes is not None:
             return self._refresh_token_expire_minutes
         return CFG.WEB_AUTH_REFRESH_TOKEN_EXPIRE_MINUTES
 
+    @refresh_token_expire_minutes.setter
+    def refresh_token_expire_minutes(self, minutes: int):
+        self._refresh_token_expire_minutes = minutes
+
     @property
     def access_token_cookie_name(self) -> str:
         if self._access_token_cookie_name is not None:
             return self._access_token_cookie_name
         return CFG.WEB_ACCESS_TOKEN_COOKIE_NAME
 
+    @access_token_cookie_name.setter
+    def access_token_cookie_name(self, name: str):
+        self._access_token_cookie_name = name
+
     @property
     def refresh_token_cookie_name(self) -> str:
         if self._refresh_token_cookie_name is not None:
             return self._refresh_token_cookie_name
         return CFG.WEB_REFRESH_TOKEN_COOKIE_NAME
 
+    @refresh_token_cookie_name.setter
+    def refresh_token_cookie_name(self, name: str):
+        self._refresh_token_cookie_name = name
+
     @property
     def enable_auth(self) -> bool:
         if self._enable_auth is not None:
             return self._enable_auth
         return CFG.WEB_ENABLE_AUTH
 
+    @enable_auth.setter
+    def enable_auth(self, enable: bool):
+        self._enable_auth = enable
+
     @property
     def super_admin_username(self) -> str:
         if self._super_admin_username is not None:
             return self._super_admin_username
         return CFG.WEB_SUPER_ADMIN_USERNAME
 
+    @super_admin_username.setter
+    def super_admin_username(self, username: str):
+        self._super_admin_username = username
+
     @property
     def super_admin_password(self) -> str:
         if self._super_admin_password is not None:
             return self._super_admin_password
         return CFG.WEB_SUPER_ADMIN_PASSWORD
 
+    @super_admin_password.setter
+    def super_admin_password(self, password: str):
+        self._super_admin_password = password
+
     @property
     def guest_username(self) -> str:
         if self._guest_username is not None:
             return self._guest_username
         return CFG.WEB_GUEST_USERNAME
 
+    @guest_username.setter
+    def guest_username(self, username: str):
+        self._guest_username = username
+
     @property
     def guest_accessible_tasks(self) -> list[AnyTask | str]:
         return self._guest_accessible_tasks
 
+    @guest_accessible_tasks.setter
+    def guest_accessible_tasks(self, tasks: list[AnyTask | str]):
+        self._guest_accessible_tasks = tasks
+
+    @property
+    def find_user_by_username_callback(self) -> Callable[[str], "User | None"] | None:
+        return self._find_user_by_username
+
+    @find_user_by_username_callback.setter
+    def find_user_by_username_callback(
+        self, find_user_by_username: Callable[[str], "User | None"]
+    ):
+        self._find_user_by_username = find_user_by_username
+
     @property
     def default_user(self) -> "User":
         from zrb.runner.web_schema.user import User
@@ -127,41 +177,6 @@ class WebAuthConfig:
             return [self.default_user]
         return self._user_list + [self.super_admin, self.default_user]
 
-    def set_secret_key(self, secret_key: str):
-        self._secret_key = secret_key
-
-    def set_access_token_expire_minutes(self, minutes: int):
-        self._access_token_expire_minutes = minutes
-
-    def set_refresh_token_expire_minutes(self, minutes: int):
-        self._refresh_token_expire_minutes = minutes
-
-    def set_access_token_cookie_name(self, name: str):
-        self._access_token_cookie_name = name
-
-    def set_refresh_token_cookie_name(self, name: str):
-        self._refresh_token_cookie_name = name
-
-    def set_enable_auth(self, enable: bool):
-        self._enable_auth = enable
-
-    def set_super_admin_username(self, username: str):
-        self._super_admin_username = username
-
-    def set_super_admin_password(self, password: str):
-        self._super_admin_password = password
-
-    def set_guest_username(self, username: str):
-        self._guest_username = username
-
-    def set_guest_accessible_tasks(self, tasks: list[AnyTask | str]):
-        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 = find_user_by_username
-
     def append_user(self, user: "User"):
         duplicates = [
             existing_user

zrb/context/any_shared_context.py

@@ -1,10 +1,10 @@
 from __future__ import annotations  # Enables forward references
 
+import sys
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, TextIO
 
 from zrb.dot_dict.dot_dict import DotDict
-from zrb.xcom.xcom import Xcom
 
 if TYPE_CHECKING:
     from zrb.session import any_session
@@ -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
 
@@ -81,3 +87,14 @@ class AnySharedContext(ABC):
             str: The rendered template as a string.
         """
         pass
+
+    @abstractmethod
+    def shared_print(
+        self,
+        *values: object,
+        sep: str = " ",
+        end: str = "\n",
+        file: TextIO | None = sys.stderr,
+        flush: bool = True,
+    ):
+        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)
@@ -101,6 +101,18 @@ class Context(AnyContext):
             return template
         return float(self.render(template))
 
+    def shared_print(
+        self,
+        *values: object,
+        sep: str = " ",
+        end: str = "\n",
+        file: TextIO | None = sys.stderr,
+        flush: bool = True,
+    ):
+        return self._shared_ctx.shared_print(
+            *values, sep=sep, end=end, file=file, flush=flush
+        )
+
     def print(
         self,
         *values: object,
@@ -110,16 +122,25 @@ class Context(AnyContext):
         flush: bool = True,
         plain: bool = False,
     ):
-        sep = " " if sep is None else sep
+        if sep is None:
+            sep = " "
+        if end is None:
+            end = "\n"
         message = sep.join([f"{value}" for value in values])
         if plain:
             # self.append_to_shared_log(remove_style(message))
-            print(message, sep=sep, end=end, file=file, flush=flush)
+            self.shared_print(message, sep=sep, end=end, file=file, flush=flush)
             self.append_to_shared_log(remove_style(f"{message}{end}"))
             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:
@@ -131,7 +152,20 @@ class Context(AnyContext):
         prefix = f"{formatted_time}{attempt_status} {padded_styled_task_name} ⬤ "
         self.append_to_shared_log(remove_style(f"{prefix} {message}{end}"))
         stylized_prefix = stylize(prefix, color=color)
-        print(f"{stylized_prefix} {message}", sep=sep, end=end, file=file, flush=flush)
+        self.shared_print(
+            f"{stylized_prefix} {message}", sep=sep, end=end, file=file, flush=flush
+        )
+
+    def print_err(
+        self,
+        *values: object,
+        sep: str | None = " ",
+        end: str | None = "\n",
+        file: TextIO | None = sys.stderr,
+        flush: bool = True,
+        plain: bool = False,
+    ):
+        self.print(*values, sep=sep, end=end, file=file, flush=flush, plain=plain)
 
     def log_debug(
         self,

zrb/context/print_fn.py

@@ -0,0 +1,13 @@
+from typing import Protocol, TextIO
+
+
+class PrintFn(Protocol):
+    def __call__(
+        self,
+        *values: object,
+        sep: str = " ",
+        end: str = "\n",
+        file: TextIO | None = None,
+        flush: bool = False,
+    ) -> None:
+        pass

zrb/context/shared_context.py

@@ -1,9 +1,10 @@
 import datetime
 import sys
-from typing import Any
+from typing import Any, TextIO
 
 from zrb.config.config import CFG
 from zrb.context.any_shared_context import AnySharedContext
+from zrb.context.print_fn import PrintFn
 from zrb.dot_dict.dot_dict import DotDict
 from zrb.session.any_session import AnySession
 from zrb.util.string.conversion import (
@@ -28,6 +29,7 @@ class SharedContext(AnySharedContext):
         xcom: dict[str, Xcom] = {},
         logging_level: int | None = None,
         is_web_mode: bool = False,
+        print_fn: PrintFn | None = None,
     ):
         self.__logging_level = logging_level
         self._input = DotDict(input)
@@ -37,14 +39,11 @@ class SharedContext(AnySharedContext):
         self._session: AnySession | None = None
         self._log = []
         self._is_web_mode = is_web_mode
+        self._print_fn = print_fn if print_fn is not None else print
 
     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:
@@ -70,7 +69,7 @@ class SharedContext(AnySharedContext):
         return self._args
 
     @property
-    def xcom(self) -> DotDict[str, Xcom]:
+    def xcom(self) -> DotDict:
         return self._xcom
 
     @property
@@ -85,7 +84,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)
 
@@ -112,3 +111,13 @@ class SharedContext(AnySharedContext):
                 "double_quote": double_quote,
             },
         )
+
+    def shared_print(
+        self,
+        *values: object,
+        sep: str = " ",
+        end: str = "\n",
+        file: TextIO | None = sys.stderr,
+        flush: bool = True,
+    ):
+        return self._print_fn(*values, sep=sep, end=end, file=file, flush=flush)

zrb/group/any_group.py

@@ -35,11 +35,11 @@ class AnyGroup(ABC):
         pass
 
     @abstractmethod
-    def add_group(self, group: "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
@@ -55,5 +55,5 @@ class AnyGroup(ABC):
         pass
 
     @abstractmethod
-    def get_group_by_alias(self, name: str) -> "AnyGroup | None":
+    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

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

@@ -1,7 +1,13 @@
+from typing import TYPE_CHECKING
+
 from zrb.attr.type import StrAttr, StrListAttr
 from zrb.context.any_shared_context import AnySharedContext
 from zrb.input.base_input import BaseInput
 from zrb.util.attr import get_str_list_attr
+from zrb.util.match import fuzzy_match
+
+if TYPE_CHECKING:
+    from prompt_toolkit.completion import Completer
 
 
 class OptionInput(BaseInput):
@@ -47,9 +53,43 @@ 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
+
+        if shared_ctx.is_tty:
+            reader = PromptSession()
+            option_completer = self._get_option_completer(options)
+            return reader.prompt(f"{prompt_message}: ", completer=option_completer)
+        return input(f"{prompt_message}: ")
+
+    def _get_option_completer(self, options: list[str]) -> "Completer":
+        from prompt_toolkit.completion import CompleteEvent, Completer, Completion
+        from prompt_toolkit.document import Document
+
+        class OptionCompleter(Completer):
+            def __init__(self, options: list[str]):
+                self._options = options
+
+            def get_completions(
+                self, document: Document, complete_event: CompleteEvent
+            ):
+                search_pattern = document.get_word_before_cursor(WORD=True)
+                candidates = []
+                for option in self._options:
+                    matched, score = fuzzy_match(option, search_pattern)
+                    if matched:
+                        candidates.append((score, option))
+                candidates.sort(key=lambda x: (x[0], x[1]))
+                for _, option in candidates:
+                    yield Completion(option, start_position=-len(search_pattern))
+
+        return OptionCompleter(options)

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/llm/agent/__init__.py

@@ -0,0 +1,9 @@
+from zrb.llm.agent.agent import create_agent, run_agent, tool_confirmation_var
+from zrb.llm.agent.summarizer import create_summarizer_agent
+
+__all__ = [
+    "create_agent",
+    "run_agent",
+    "tool_confirmation_var",
+    "create_summarizer_agent",
+]