zrb 1.21.29__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
@@ -87,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

@@ -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,11 +122,14 @@ 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
@@ -137,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,6 +39,7 @@ 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__
@@ -108,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/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):
@@ -58,10 +64,32 @@ class OptionInput(BaseInput):
         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)
+            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/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",
+]

zrb/llm/agent/agent.py

@@ -0,0 +1,215 @@
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any, Callable
+
+from zrb.llm.config.config import llm_config as default_llm_config
+from zrb.llm.config.limiter import LLMLimiter
+from zrb.llm.util.attachment import normalize_attachments
+from zrb.llm.util.prompt import expand_prompt
+
+# Context variable to propagate tool confirmation callback to sub-agents
+tool_confirmation_var: ContextVar[Callable[[Any], Any] | None] = ContextVar(
+    "tool_confirmation", default=None
+)
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent, DeferredToolRequests, DeferredToolResults, Tool
+    from pydantic_ai._agent_graph import HistoryProcessor
+    from pydantic_ai.messages import UserPromptPart
+    from pydantic_ai.models import Model
+    from pydantic_ai.output import OutputDataT, OutputSpec
+    from pydantic_ai.settings import ModelSettings
+    from pydantic_ai.tools import ToolFuncEither
+    from pydantic_ai.toolsets import AbstractToolset
+
+
+def create_agent(
+    model: "Model | str | None" = None,
+    system_prompt: str = "",
+    tools: list["Tool | ToolFuncEither"] = [],
+    toolsets: list["AbstractToolset[None]"] = [],
+    model_settings: "ModelSettings | None" = None,
+    history_processors: list["HistoryProcessor"] | None = None,
+    output_type: "OutputSpec[OutputDataT]" = str,
+    retries: int = 1,
+    yolo: bool = False,
+) -> "Agent[None, Any]":
+    from pydantic_ai import Agent, DeferredToolRequests
+    from pydantic_ai.toolsets import FunctionToolset
+
+    # Expand system prompt with references
+    effective_system_prompt = expand_prompt(system_prompt)
+
+    final_output_type = output_type
+    effective_toolsets = list(toolsets)
+    if tools:
+        effective_toolsets.append(FunctionToolset(tools=tools))
+
+    if not yolo:
+        final_output_type = output_type | DeferredToolRequests
+        effective_toolsets = [ts.approval_required() for ts in effective_toolsets]
+
+    if model is None:
+        model = default_llm_config.model
+
+    return Agent(
+        model=model,
+        output_type=final_output_type,
+        instructions=effective_system_prompt,
+        toolsets=effective_toolsets,
+        model_settings=model_settings,
+        history_processors=history_processors,
+        retries=retries,
+    )
+
+
+async def run_agent(
+    agent: "Agent[None, Any]",
+    message: str | None,
+    message_history: list[Any],
+    limiter: LLMLimiter,
+    attachments: list[Any] | None = None,
+    print_fn: Callable[[str], Any] = print,
+    event_handler: Callable[[Any], Any] | None = None,
+    tool_confirmation: Callable[[Any], Any] | None = None,
+) -> tuple[Any, list[Any]]:
+    """
+    Runs the agent with rate limiting, history management, and optional CLI confirmation loop.
+    Returns (result_output, new_message_history).
+    """
+    import asyncio
+
+    from pydantic_ai import AgentRunResultEvent, DeferredToolRequests
+
+    # Resolve tool confirmation callback (Arg > Context > None)
+    effective_tool_confirmation = tool_confirmation
+    if effective_tool_confirmation is None:
+        effective_tool_confirmation = tool_confirmation_var.get()
+
+    # Set context var for sub-agents
+    token = tool_confirmation_var.set(effective_tool_confirmation)
+
+    try:
+        # Expand user message with references
+        effective_message = expand_prompt(message) if message else message
+
+        # Prepare Prompt Content
+        prompt_content = _get_prompt_content(effective_message, attachments, print_fn)
+
+        # 1. Prune & Throttle
+        current_history = await _acquire_rate_limit(
+            limiter, prompt_content, message_history, print_fn
+        )
+        current_message = prompt_content
+        current_results = None
+
+        # 2. Execution Loop
+        while True:
+            result_output = None
+            run_history = []
+
+            async for event in agent.run_stream_events(
+                current_message,
+                message_history=current_history,
+                deferred_tool_results=current_results,
+            ):
+                await asyncio.sleep(0)
+                if isinstance(event, AgentRunResultEvent):
+                    result = event.result
+                    result_output = result.output
+                    run_history = result.all_messages()
+                if event_handler:
+                    await event_handler(event)
+
+            # Handle Deferred Calls
+            if isinstance(result_output, DeferredToolRequests):
+                current_results = await _process_deferred_requests(
+                    result_output, effective_tool_confirmation
+                )
+                if current_results is None:
+                    return result_output, run_history
+                # Prepare next iteration
+                current_message = None
+                current_history = run_history
+                continue
+            return result_output, run_history
+    finally:
+        tool_confirmation_var.reset(token)
+
+
+def _get_prompt_content(
+    message: str | None, attachments: list[Any] | None, print_fn: Callable[[str], Any]
+) -> "list[UserPromptPart] | str | None":
+    from pydantic_ai.messages import UserPromptPart
+
+    prompt_content = message
+    if attachments:
+        attachments = normalize_attachments(attachments, print_fn)
+        parts: list[UserPromptPart] = []
+        if message:
+            parts.append(UserPromptPart(content=message))
+        parts.extend(attachments)
+        prompt_content = parts
+    return prompt_content
+
+
+async def _acquire_rate_limit(
+    limiter: LLMLimiter,
+    message: str | None,
+    message_history: list[Any],
+    print_fn: Callable[[str], Any],
+) -> list[Any]:
+    """Prunes history and waits if rate limits are exceeded."""
+    if not message:
+        return message_history
+
+    # Prune
+    pruned_history = limiter.fit_context_window(message_history, message)
+
+    # Throttle
+    est_tokens = limiter.count_tokens(pruned_history) + limiter.count_tokens(message)
+    await limiter.acquire(
+        est_tokens, notifier=lambda msg: print_fn(msg) if msg else None
+    )
+
+    return pruned_history
+
+
+async def _process_deferred_requests(
+    result_output: "DeferredToolRequests",
+    effective_tool_confirmation: Callable[[Any], Any] | None,
+) -> "DeferredToolResults | None":
+    """Handles tool approvals/denials via callback or CLI fallback."""
+    import asyncio
+    import inspect
+
+    from pydantic_ai import DeferredToolResults, ToolApproved, ToolDenied
+
+    all_requests = (result_output.calls or []) + (result_output.approvals or [])
+    if not all_requests:
+        return None
+
+    current_results = DeferredToolResults()
+
+    for call in all_requests:
+        if effective_tool_confirmation:
+            res = effective_tool_confirmation(call)
+            if inspect.isawaitable(res):
+                result = await res
+            else:
+                result = res
+            current_results.approvals[call.tool_call_id] = result
+        else:
+            # CLI Fallback
+            prompt_text = f"Execute tool '{call.tool_name}' with args {call.args}?"
+            prompt_cli = f"\n[?] {prompt_text} (y/N) "
+
+            # We use asyncio.to_thread(input, ...) to avoid blocking the loop
+            user_input = await asyncio.to_thread(input, prompt_cli)
+            answer = user_input.strip().lower() in ("y", "yes")
+
+            if answer:
+                current_results.approvals[call.tool_call_id] = ToolApproved()
+            else:
+                current_results.approvals[call.tool_call_id] = ToolDenied("User denied")
+
+    return current_results

zrb/llm/agent/summarizer.py

@@ -0,0 +1,20 @@
+from typing import TYPE_CHECKING
+
+from zrb.llm.agent.agent import create_agent
+from zrb.llm.prompt.default import get_summarizer_system_prompt
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.models import Model
+
+
+def create_summarizer_agent(
+    model: "str | None | Model" = None,
+    system_prompt: str | None = None,
+) -> "Agent[None, str]":
+    effective_system_prompt = system_prompt or get_summarizer_system_prompt()
+
+    return create_agent(
+        model=model,
+        system_prompt=effective_system_prompt,
+    )

zrb/llm/app/__init__.py

@@ -0,0 +1,10 @@
+from zrb.llm.app.confirmation.allow_tool import allow_tool_usage
+from zrb.llm.app.confirmation.handler import ConfirmationMiddleware, last_confirmation
+from zrb.llm.app.confirmation.replace_confirmation import replace_confirmation
+
+__all__ = [
+    "allow_tool_usage",
+    "ConfirmationMiddleware",
+    "last_confirmation",
+    "replace_confirmation",
+]