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

zrb/llm/app/completion.py

@@ -0,0 +1,281 @@
+import os
+from datetime import datetime
+from typing import Iterable
+
+from prompt_toolkit.completion import (
+    CompleteEvent,
+    Completer,
+    Completion,
+    PathCompleter,
+)
+from prompt_toolkit.document import Document
+
+from zrb.llm.history_manager.any_history_manager import AnyHistoryManager
+from zrb.util.match import fuzzy_match
+
+
+class InputCompleter(Completer):
+    def __init__(
+        self,
+        history_manager: AnyHistoryManager,
+        attach_commands: list[str] = [],
+        exit_commands: list[str] = [],
+        info_commands: list[str] = [],
+        save_commands: list[str] = [],
+        load_commands: list[str] = [],
+        redirect_output_commands: list[str] = [],
+        summarize_commands: list[str] = [],
+        exec_commands: list[str] = [],
+    ):
+        self._history_manager = history_manager
+        self._attach_commands = attach_commands
+        self._exit_commands = exit_commands
+        self._info_commands = info_commands
+        self._save_commands = save_commands
+        self._load_commands = load_commands
+        self._redirect_output_commands = redirect_output_commands
+        self._summarize_commands = summarize_commands
+        self._exec_commands = exec_commands
+        # expanduser=True allows ~/path
+        self._path_completer = PathCompleter(expanduser=True)
+        # Cache for file listing to improve performance
+        self._file_cache: list[str] | None = None
+        self._file_cache_time = 0
+        self._cmd_history = self._get_cmd_history()
+
+    def get_completions(
+        self, document: Document, complete_event: CompleteEvent
+    ) -> Iterable[Completion]:
+        text_before_cursor = document.text_before_cursor.lstrip()
+        word = document.get_word_before_cursor(WORD=True)
+
+        all_commands = (
+            self._exit_commands
+            + self._attach_commands
+            + self._summarize_commands
+            + self._info_commands
+            + self._save_commands
+            + self._load_commands
+            + self._redirect_output_commands
+            + self._exec_commands
+        )
+        command_prefixes = {cmd[0] for cmd in all_commands if cmd}
+
+        # 1. Command and Argument Completion
+        if text_before_cursor and text_before_cursor[0] in command_prefixes:
+            parts = text_before_cursor.split()
+            # Check if we are typing the command itself or arguments
+            is_typing_command = len(parts) == 1 and not text_before_cursor.endswith(" ")
+            is_typing_arg = (len(parts) == 1 and text_before_cursor.endswith(" ")) or (
+                len(parts) >= 2
+            )
+
+            if is_typing_command:
+                lower_word = word.lower()
+                prefix = text_before_cursor[0]
+                for cmd in all_commands:
+                    if cmd.startswith(prefix) and cmd.lower().startswith(lower_word):
+                        yield Completion(cmd, start_position=-len(word))
+                return
+
+            if is_typing_arg:
+                cmd = parts[0]
+                arg_prefix = text_before_cursor[len(cmd) :].lstrip()
+
+                # Exec Command: Suggest History
+                if self._is_command(cmd, self._exec_commands):
+                    # Filter history
+                    matches = [h for h in self._cmd_history if h.startswith(arg_prefix)]
+                    # Sort matches by length (shorter first) as heuristic? Or just recent?
+                    # Since _cmd_history is set (unique), we lose order.
+                    # But Python 3.7+ dicts preserve insertion order, so if we used dict keys, we kept order.
+                    # Let's assume _get_cmd_history returns recent last.
+                    # We reverse to show most recent first.
+                    for h in reversed(matches):
+                        yield Completion(h, start_position=-len(arg_prefix))
+                    return
+
+                # Check if we are typing the second part (argument) strictly
+                # (Re-evaluating logic for other commands which only take 1 arg usually)
+                if not (
+                    (len(parts) == 1 and text_before_cursor.endswith(" "))
+                    or (len(parts) == 2 and not text_before_cursor.endswith(" "))
+                ):
+                    return
+
+                arg_prefix = parts[1] if len(parts) == 2 else ""
+
+                # Save Command: Suggest Timestamp
+                if self._is_command(cmd, self._save_commands):
+                    ts = datetime.now().strftime("%Y-%m-%d-%H-%M")
+                    if ts.startswith(arg_prefix):
+                        yield Completion(ts, start_position=-len(arg_prefix))
+                    return
+
+                # Redirect Command: Suggest Timestamp.txt
+                if self._is_command(cmd, self._redirect_output_commands):
+                    ts = datetime.now().strftime("%Y-%m-%d-%H-%M.txt")
+                    if ts.startswith(arg_prefix):
+                        yield Completion(ts, start_position=-len(arg_prefix))
+                    return
+
+                # Load Command: Search History
+                if self._is_command(cmd, self._load_commands):
+                    results = self._history_manager.search(arg_prefix)
+                    for res in results[:10]:
+                        yield Completion(res, start_position=-len(arg_prefix))
+                    return
+
+                # Attach Command: Suggest Files
+                if self._is_command(cmd, self._attach_commands):
+                    yield from self._get_file_completions(
+                        arg_prefix, complete_event, only_files=True
+                    )
+                    return
+
+                # Other commands (Exit, Info, Summarize) need no completion
+                return
+
+        # 2. File Completion (@)
+        if word.startswith("@"):
+            path_part = word[1:]
+            yield from self._get_file_completions(
+                path_part, complete_event, only_files=False
+            )
+
+    def _get_cmd_history(self) -> list[str]:
+        history_files = [
+            os.path.expanduser("~/.bash_history"),
+            os.path.expanduser("~/.zsh_history"),
+        ]
+        unique_cmds = {}  # Use dict to preserve order (insertion order)
+
+        for hist_file in history_files:
+            if not os.path.exists(hist_file):
+                continue
+            try:
+                with open(hist_file, "r", errors="ignore") as f:
+                    for line in f:
+                        line = line.strip()
+                        if not line:
+                            continue
+                        # Handle zsh timestamp format: : 1612345678:0;command
+                        if line.startswith(": ") and ";" in line:
+                            parts = line.split(";", 1)
+                            if len(parts) == 2:
+                                line = parts[1]
+
+                        if line:
+                            # Remove existing to update position to end (most recent)
+                            if line in unique_cmds:
+                                del unique_cmds[line]
+                            unique_cmds[line] = None
+            except Exception:
+                pass
+
+        return list(unique_cmds.keys())
+
+    def _is_command(self, cmd: str, cmd_list: list[str]) -> bool:
+        return cmd.lower() in [c.lower() for c in cmd_list]
+
+    def _get_file_completions(
+        self, text: str, complete_event: CompleteEvent, only_files: bool = False
+    ) -> Iterable[Completion]:
+        # Logic:
+        # - If text indicates path traversal (/, ., ~), use PathCompleter
+        # - Else, check file count. If < 5000, use Fuzzy. Else use PathCompleter.
+
+        if self._is_path_navigation(text):
+            yield from self._get_path_completions(text, complete_event, only_files)
+            return
+
+        # Count files (cached strategy could be added here if needed)
+        files = self._get_recursive_files(limit=5000)
+        if len(files) < 5000:
+            # Fuzzy Match
+            yield from self._get_fuzzy_completions(text, files, only_files)
+        else:
+            # Fallback to PathCompleter for large repos
+            yield from self._get_path_completions(text, complete_event, only_files)
+
+    def _is_path_navigation(self, text: str) -> bool:
+        return (
+            text.startswith("/")
+            or text.startswith(".")
+            or text.startswith("~")
+            or os.sep in text
+        )
+
+    def _get_path_completions(
+        self, text: str, complete_event: CompleteEvent, only_files: bool
+    ) -> Iterable[Completion]:
+        # PathCompleter needs a document where text represents the path
+        fake_document = Document(text=text, cursor_position=len(text))
+        for c in self._path_completer.get_completions(fake_document, complete_event):
+            if only_files:
+                # Check if the completed path is a directory
+                # Note: 'text' is the prefix. c.text is the completion suffix.
+                # We need to reconstruct full path to check isdir
+                # This is tricky with PathCompleter's internal logic.
+                # A simple heuristic: if it ends with path separator, it's a dir.
+                if c.text.endswith(os.sep):
+                    continue
+            yield c
+
+    def _get_fuzzy_completions(
+        self, text: str, files: list[str], only_files: bool
+    ) -> Iterable[Completion]:
+        matches = []
+        for f in files:
+            if only_files and f.endswith(os.sep):
+                continue
+            is_match, score = fuzzy_match(f, text)
+            if is_match:
+                matches.append((score, f))
+
+        # Sort by score (lower is better)
+        matches.sort(key=lambda x: x[0])
+
+        # Return top 20
+        for _, f in matches[:20]:
+            yield Completion(f, start_position=-len(text))
+
+    def _get_recursive_files(self, root: str = ".", limit: int = 5000) -> list[str]:
+        # Simple walker with exclusions
+        paths = []
+        # Check if current dir is hidden
+        cwd_is_hidden = os.path.basename(os.path.abspath(root)).startswith(".")
+
+        try:
+            for dirpath, dirnames, filenames in os.walk(root):
+                # Exclude hidden directories unless root is hidden
+                if not cwd_is_hidden:
+                    dirnames[:] = [d for d in dirnames if not d.startswith(".")]
+
+                # Exclude common ignores
+                dirnames[:] = [
+                    d
+                    for d in dirnames
+                    if d not in ("node_modules", "__pycache__", "venv", ".venv")
+                ]
+
+                rel_dir = os.path.relpath(dirpath, root)
+                if rel_dir == ".":
+                    rel_dir = ""
+
+                # Add directories
+                for d in dirnames:
+                    paths.append(os.path.join(rel_dir, d) + os.sep)
+                    if len(paths) >= limit:
+                        return paths
+
+                # Add files
+                for f in filenames:
+                    if not cwd_is_hidden and f.startswith("."):
+                        continue
+                    paths.append(os.path.join(rel_dir, f))
+                    if len(paths) >= limit:
+                        return paths
+        except Exception:
+            pass
+        return paths

zrb/llm/app/confirmation/allow_tool.py

@@ -0,0 +1,66 @@
+import json
+import re
+from typing import Any, Awaitable, Callable, Dict, Optional
+
+from zrb.llm.app.confirmation.handler import ConfirmationMiddleware, UIProtocol
+
+
+def allow_tool_usage(
+    tool_name: str, kwargs: Optional[Dict[str, str]] = None
+) -> ConfirmationMiddleware:
+    """
+    Creates a confirmation middleware that automatically approves a tool execution
+    if it matches the specified tool_name and argument constraints.
+
+    :param tool_name: The name of the tool to allow.
+    :param kwargs: A dictionary of regex patterns for arguments.
+                   If None or empty, the tool is allowed regardless of arguments.
+                   If provided, arguments in the tool call must match the regex patterns
+                   specified in kwargs (only for arguments present in both).
+    :return: A ConfirmationMiddleware function.
+    """
+    from pydantic_ai import ToolApproved
+
+    async def middleware(
+        ui: UIProtocol,
+        call: Any,
+        response: str,
+        next_handler: Callable[[UIProtocol, Any, str], Awaitable[Any]],
+    ) -> Any:
+        # Check if tool name matches
+        if call.tool_name != tool_name:
+            return await next_handler(ui, call, response)
+
+        # If kwargs is empty or None, approve
+        if not kwargs:
+            ui.append_to_output(f"\n✅ Auto-approved tool: {tool_name}")
+            return ToolApproved()
+
+        # Parse arguments
+        try:
+            args = call.args
+            if isinstance(args, str):
+                args = json.loads(args)
+
+            if not isinstance(args, dict):
+                # If args is not a dict (e.g. primitive), and kwargs is not empty,
+                # we assume it doesn't match complex constraints (or we can't check keys).
+                # So we delegate to the next handler.
+                return await next_handler(ui, call, response)
+
+        except (json.JSONDecodeError, ValueError):
+            return await next_handler(ui, call, response)
+
+        # Check constraints
+        # "all parameter in the call parameter has to match the ones in kwargs (if that parameter defined in the kwargs)"
+        for arg_name, arg_value in args.items():
+            if arg_name in kwargs:
+                pattern = kwargs[arg_name]
+                # Convert arg_value to string for regex matching
+                if not re.search(pattern, str(arg_value)):
+                    return await next_handler(ui, call, response)
+
+        ui.append_to_output(f"\n✅ Auto-approved tool: {tool_name} with matching args")
+        return ToolApproved()
+
+    return middleware

zrb/llm/app/confirmation/handler.py

@@ -0,0 +1,178 @@
+import json
+import os
+import subprocess
+import tempfile
+from typing import Any, Awaitable, Callable, Protocol, TextIO
+
+import yaml
+
+from zrb.config.config import CFG
+from zrb.util.yaml import yaml_dump
+
+
+class UIProtocol(Protocol):
+    async def ask_user(self, prompt: str) -> str: ...
+
+    def append_to_output(
+        self,
+        *values: object,
+        sep: str = " ",
+        end: str = "\n",
+        file: TextIO | None = None,
+        flush: bool = False,
+    ): ...
+
+
+ConfirmationMiddleware = Callable[
+    [UIProtocol, Any, str, Callable[[UIProtocol, Any, str], Awaitable[Any]]],
+    Awaitable[Any],
+]
+
+
+class ConfirmationHandler:
+    def __init__(self, middlewares: list[ConfirmationMiddleware]):
+        self._middlewares = middlewares
+
+    def add_middleware(self, *middleware: ConfirmationMiddleware):
+        self.prepend_middleware(*middleware)
+
+    def prepend_middleware(self, *middleware: ConfirmationMiddleware):
+        self._middlewares = list(middleware) + self._middlewares
+
+    async def handle(self, ui: UIProtocol, call: Any) -> Any:
+        while True:
+            message = self._get_confirm_user_message(call)
+            ui.append_to_output(f"\n\n{message}", end="")
+            # Wait for user input
+            user_input = await ui.ask_user("")
+            user_response = user_input.strip()
+
+            # Build the chain
+            async def _next(
+                ui: UIProtocol, call: Any, response: str, index: int
+            ) -> Any:
+                if index >= len(self._middlewares):
+                    # Default if no middleware handles it
+                    return None
+                middleware = self._middlewares[index]
+                return await middleware(
+                    ui,
+                    call,
+                    response,
+                    lambda u, c, r: _next(u, c, r, index + 1),
+                )
+
+            result = await _next(ui, call, user_response, 0)
+            if result is None:
+                continue
+            return result
+
+    def _get_confirm_user_message(self, call: Any) -> str:
+        args_section = ""
+        if f"{call.args}" != "{}":
+            args_str = self._format_args(call.args)
+            args_section = f"       Arguments:\n{args_str}\n"
+        return (
+            f"  🎰 Executing tool '{call.tool_name}'\n"
+            f"{args_section}"
+            "  ❓ Allow tool Execution? (✅ Y | 🛑 n | ✏️ e)? "
+        )
+
+    def _format_args(self, args: Any) -> str:
+        indent = " " * 7
+        try:
+            if isinstance(args, str):
+                try:
+                    args = json.loads(args)
+                except json.JSONDecodeError:
+                    pass
+            args_str = yaml_dump(args)
+            # Indent nicely for display
+            return "\n".join([f"{indent}{line}" for line in args_str.splitlines()])
+        except Exception:
+            return f"{indent}{args}"
+
+
+async def last_confirmation(
+    ui: UIProtocol,
+    call: Any,
+    user_response: str,
+    next_handler: Callable[[UIProtocol, Any, str], Awaitable[Any]],
+) -> Any:
+    from pydantic_ai import ToolApproved, ToolDenied
+
+    if user_response.lower() in ("y", "yes", "ok", "okay", ""):
+        ui.append_to_output("\n✅ Execution approved.")
+        return ToolApproved()
+    elif user_response.lower() in ("n", "no"):
+        ui.append_to_output("\n🛑 Execution denied.")
+        return ToolDenied("User denied execution")
+    elif user_response.lower() in ("e", "edit"):
+        # Edit logic
+        try:
+            args = call.args
+            if isinstance(args, str):
+                try:
+                    args = json.loads(args)
+                except json.JSONDecodeError:
+                    pass
+
+            # YAML for editing
+            is_yaml_edit = True
+            try:
+                content = yaml_dump(args)
+                extension = ".yaml"
+            except Exception:
+                # Fallback to JSON
+                content = json.dumps(args, indent=2)
+                extension = ".json"
+                is_yaml_edit = False
+
+            new_content = await wait_edit_content(
+                text_editor=CFG.DEFAULT_EDITOR,
+                content=content,
+                extension=extension,
+            )
+
+            # Compare content
+            if new_content == content:
+                ui.append_to_output("\nℹ️ No changes made.")
+                return None
+
+            try:
+                if is_yaml_edit:
+                    new_args = yaml.safe_load(new_content)
+                else:
+                    new_args = json.loads(new_content)
+                ui.append_to_output("\n✅ Execution approved (with modification).")
+                return ToolApproved(override_args=new_args)
+            except Exception as e:
+                ui.append_to_output(f"\n❌ Invalid format: {e}. ", end="")
+                # Return None to signal loop retry
+                return None
+
+        except Exception as e:
+            ui.append_to_output(f"\n❌ Error editing: {e}. ", end="")
+            return None
+    else:
+        ui.append_to_output("\n🛑 Execution denied.")
+        return ToolDenied(f"User denied execution with message: {user_response}")
+
+
+async def wait_edit_content(
+    text_editor: str, content: str, extension: str = ".txt"
+) -> str:
+    from prompt_toolkit.application import run_in_terminal
+
+    # Write temporary file
+    with tempfile.NamedTemporaryFile(suffix=extension, mode="w+", delete=False) as tf:
+        tf.write(content)
+        tf_path = tf.name
+
+    # Edit and wait
+    await run_in_terminal(lambda: subprocess.call([text_editor, tf_path]))
+    with open(tf_path, "r") as tf:
+        new_content = tf.read()
+    os.remove(tf_path)
+
+    return new_content

zrb/llm/app/confirmation/replace_confirmation.py

@@ -0,0 +1,77 @@
+import json
+import os
+import subprocess
+import tempfile
+from typing import Any, Awaitable, Callable
+
+from prompt_toolkit.application import run_in_terminal
+
+from zrb.config.config import CFG
+from zrb.llm.app.confirmation.handler import UIProtocol
+
+
+async def replace_confirmation(
+    ui: UIProtocol,
+    call: Any,
+    response: str,
+    next_handler: Callable[[UIProtocol, Any, str], Awaitable[Any]],
+) -> Any:
+    from pydantic_ai import ToolApproved
+
+    if call.tool_name != "replace_in_file":
+        return await next_handler(ui, call, response)
+
+    if response.lower() not in ("e", "edit"):
+        return await next_handler(ui, call, response)
+
+    # It is replace_in_file and user wants to edit
+    args = call.args
+    if isinstance(args, str):
+        try:
+            args = json.loads(args)
+        except json.JSONDecodeError:
+            pass
+
+    old_text = args.get("old_text", "")
+    new_text = args.get("new_text", "")
+
+    # Create temporary files
+    with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".old") as tf_old:
+        tf_old.write(old_text)
+        old_path = tf_old.name
+
+    with tempfile.NamedTemporaryFile(mode="w+", delete=False, suffix=".new") as tf_new:
+        tf_new.write(new_text)
+        new_path = tf_new.name
+
+    try:
+        # Prepare command
+        cmd_tpl = CFG.DEFAULT_DIFF_EDIT_COMMAND_TPL
+        cmd = cmd_tpl.format(old=old_path, new=new_path)
+
+        # Run command
+        await run_in_terminal(lambda: subprocess.call(cmd, shell=True))
+
+        # Read back new content
+        with open(new_path, "r") as f:
+            edited_new_text = f.read()
+
+        if edited_new_text != new_text:
+            # Update args
+            new_args = dict(args)
+            new_args["new_text"] = edited_new_text
+            ui.append_to_output("\n✅ Replacement modified.")
+            return ToolApproved(override_args=new_args)
+        else:
+            ui.append_to_output("\nℹ️ No changes made.")
+            return None
+
+    except Exception as e:
+        ui.append_to_output(f"\n❌ Error during diff edit: {e}")
+        return None
+    finally:
+        # Cleanup
+        if os.path.exists(old_path):
+            os.remove(old_path)
+        if os.path.exists(new_path):
+            os.remove(new_path)

zrb/llm/app/keybinding.py

@@ -0,0 +1,34 @@
+import string
+
+from prompt_toolkit.application import get_app
+from prompt_toolkit.key_binding import KeyBindings
+from prompt_toolkit.widgets import TextArea
+
+
+def create_output_keybindings(input_field: TextArea) -> KeyBindings:
+    kb = KeyBindings()
+
+    @kb.add("escape")
+    def _(event):
+        get_app().layout.focus(input_field)
+
+    @kb.add("c-c")
+    def _(event):
+        # Copy selection to clipboard
+        if event.current_buffer.selection_state:
+            data = event.current_buffer.copy_selection()
+            event.app.clipboard.set_data(data)
+        get_app().layout.focus(input_field)
+
+    def redirect_focus(event):
+        get_app().layout.focus(input_field)
+        input_field.buffer.insert_text(event.data)
+
+    for char in string.printable:
+        # Skip control characters (Tab, Newline, etc.)
+        #  to preserve navigation/standard behavior
+        if char in "\t\n\r\x0b\x0c":
+            continue
+        kb.add(char)(redirect_focus)
+
+    return kb