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

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

zrb/llm/app/layout.py

@@ -0,0 +1,117 @@
+from typing import Callable
+
+from prompt_toolkit.formatted_text import HTML, AnyFormattedText
+from prompt_toolkit.layout import HSplit, Layout, Window, WindowAlign
+from prompt_toolkit.layout.containers import Float, FloatContainer
+from prompt_toolkit.layout.controls import FormattedTextControl
+from prompt_toolkit.layout.menus import CompletionsMenu
+from prompt_toolkit.lexers import Lexer
+from prompt_toolkit.widgets import Frame, TextArea
+
+from zrb.llm.app.completion import InputCompleter
+from zrb.llm.history_manager.any_history_manager import AnyHistoryManager
+
+
+def create_input_field(
+    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],
+) -> TextArea:
+    return TextArea(
+        height=4,
+        prompt=HTML('<style color="ansibrightblue"><b>&gt;&gt;&gt; </b></style>'),
+        multiline=True,
+        wrap_lines=True,
+        completer=InputCompleter(
+            history_manager=history_manager,
+            attach_commands=attach_commands,
+            exit_commands=exit_commands,
+            info_commands=info_commands,
+            save_commands=save_commands,
+            load_commands=load_commands,
+            redirect_output_commands=redirect_output_commands,
+            summarize_commands=summarize_commands,
+            exec_commands=exec_commands,
+        ),
+        complete_while_typing=True,
+        focus_on_click=True,
+        style="class:input_field",
+    )
+
+
+def create_output_field(greeting: str, lexer: Lexer) -> TextArea:
+    return TextArea(
+        text=greeting.rstrip() + "\n\n",
+        read_only=True,
+        scrollbar=False,
+        wrap_lines=True,
+        lexer=lexer,
+        focus_on_click=True,
+        focusable=True,
+        style="class:output_field",
+    )
+
+
+def create_layout(
+    title: str,
+    jargon: str,
+    input_field: TextArea,
+    output_field: TextArea,
+    info_bar_text: Callable[[], AnyFormattedText],
+    status_bar_text: Callable[[], AnyFormattedText],
+) -> Layout:
+    title_bar_text = HTML(
+        f" <style bg='ansipurple' color='white'><b> {title} </b></style> "
+        f"<style color='#888888'>| {jargon}</style>"
+    )
+
+    return Layout(
+        FloatContainer(
+            content=HSplit(
+                [
+                    # Title Bar
+                    Window(
+                        height=2,
+                        content=FormattedTextControl(title_bar_text),
+                        style="class:title-bar",
+                        align=WindowAlign.CENTER,
+                    ),
+                    # Info Bar
+                    Window(
+                        height=2,
+                        content=FormattedTextControl(info_bar_text),
+                        style="class:info-bar",
+                        align=WindowAlign.CENTER,
+                    ),
+                    # Chat History
+                    Frame(output_field, title="Conversation", style="class:frame"),
+                    # Input Area
+                    Frame(
+                        input_field,
+                        title="(ENTER to send, CTRL+ENTER for newline, ESC to cancel)",
+                        style="class:input-frame",
+                    ),
+                    # Status Bar
+                    Window(
+                        height=1,
+                        content=FormattedTextControl(status_bar_text),
+                        style="class:bottom-toolbar",
+                    ),
+                ]
+            ),
+            floats=[
+                Float(
+                    xcursor=True,
+                    ycursor=True,
+                    content=CompletionsMenu(max_height=16, scroll_offset=1),
+                ),
+            ],
+        ),
+        focused_element=input_field,
+    )

zrb/llm/app/lexer.py

@@ -0,0 +1,155 @@
+import re
+
+from prompt_toolkit.lexers import Lexer
+
+
+class CLIStyleLexer(Lexer):
+    def lex_document(self, document):
+        lines = document.lines
+        line_tokens = {}  # Cache for tokens per line
+
+        # Global state for the document
+        current_attrs = set()
+        current_fg = ""
+        current_bg = ""
+
+        # Pre-process all lines to handle state across newlines
+        # Regex to find ANSI escape sequences (CSI)
+        ansi_escape = re.compile(r"\x1B\[([0-9;]*)m")
+
+        for lineno, line in enumerate(lines):
+            tokens = []
+            last_end = 0
+
+            def build_style():
+                parts = list(current_attrs)
+                if current_fg:
+                    parts.append(current_fg)
+                if current_bg:
+                    parts.append(current_bg)
+                return " ".join(parts)
+
+            for match in ansi_escape.finditer(line):
+                start, end = match.span()
+
+                # Add text before the escape sequence with current style
+                if start > last_end:
+                    tokens.append((build_style(), line[last_end:start]))
+
+                # Parse codes
+                codes = match.group(1).split(";")
+                if not codes or codes == [""]:
+                    codes = ["0"]
+
+                # Convert to integers
+                int_codes = []
+                for c in codes:
+                    if c.isdigit():
+                        int_codes.append(int(c))
+
+                i = 0
+                while i < len(int_codes):
+                    c = int_codes[i]
+                    i += 1
+
+                    if c == 0:
+                        current_attrs.clear()
+                        current_fg = ""
+                        current_bg = ""
+                    elif c == 1:
+                        current_attrs.add("bold")
+                    elif c == 2:
+                        current_attrs.add("class:faint")
+                    elif c == 3:
+                        current_attrs.add("italic")
+                    elif c == 4:
+                        current_attrs.add("underline")
+                    elif c == 22:
+                        current_attrs.discard("bold")
+                        current_attrs.discard("class:faint")
+                    elif c == 23:
+                        current_attrs.discard("italic")
+                    elif c == 24:
+                        current_attrs.discard("underline")
+                    elif 30 <= c <= 37:
+                        colors = [
+                            "#000000",
+                            "#ff0000",
+                            "#00ff00",
+                            "#ffff00",
+                            "#0000ff",
+                            "#ff00ff",
+                            "#00ffff",
+                            "#ffffff",
+                        ]
+                        current_fg = colors[c - 30]
+                    elif c == 38:
+                        if i < len(int_codes):
+                            mode = int_codes[i]
+                            i += 1
+                            if mode == 5 and i < len(int_codes):
+                                i += 1  # Skip 256 color
+                            elif mode == 2 and i + 2 < len(int_codes):
+                                r, g, b = (
+                                    int_codes[i],
+                                    int_codes[i + 1],
+                                    int_codes[i + 2],
+                                )
+                                i += 3
+                                current_fg = f"#{r:02x}{g:02x}{b:02x}"
+                    elif c == 39:
+                        current_fg = ""
+                    elif 40 <= c <= 47:
+                        colors = [
+                            "#000000",
+                            "#ff0000",
+                            "#00ff00",
+                            "#ffff00",
+                            "#0000ff",
+                            "#ff00ff",
+                            "#00ffff",
+                            "#ffffff",
+                        ]
+                        current_bg = f"bg:{colors[c - 40]}"
+                    elif c == 48:
+                        if i < len(int_codes):
+                            mode = int_codes[i]
+                            i += 1
+                            if mode == 5 and i < len(int_codes):
+                                i += 1
+                            elif mode == 2 and i + 2 < len(int_codes):
+                                r, g, b = (
+                                    int_codes[i],
+                                    int_codes[i + 1],
+                                    int_codes[i + 2],
+                                )
+                                i += 3
+                                current_bg = f"bg:#{r:02x}{g:02x}{b:02x}"
+                    elif c == 49:
+                        current_bg = ""
+                    elif 90 <= c <= 97:
+                        colors = [
+                            "#555555",
+                            "#ff5555",
+                            "#55ff55",
+                            "#ffff55",
+                            "#5555ff",
+                            "#ff55ff",
+                            "#55ffff",
+                            "#ffffff",
+                        ]
+                        current_fg = colors[c - 90]
+
+                last_end = end
+
+            # Add remaining text
+            if last_end < len(line):
+                tokens.append((build_style(), line[last_end:]))
+
+            # Store tokens for this line
+            line_tokens[lineno] = tokens
+
+        def get_line(lineno):
+            return line_tokens.get(lineno, [])
+
+        return get_line

zrb/llm/app/redirection.py

@@ -0,0 +1,28 @@
+import io
+import sys
+
+
+class StreamToUI(io.TextIOBase):
+    """Redirect stdout to UI's append_to_output."""
+
+    def __init__(self, ui_callback):
+        self.ui_callback = ui_callback
+        self.original_stdout = sys.stdout
+        self.original_stderr = sys.stderr
+        self._is_first_write = True
+
+    def write(self, text: str) -> int:
+        from prompt_toolkit.application import get_app
+
+        text = text.expandtabs(4)
+        if text:
+            if self._is_first_write:
+                self.ui_callback("\n", end="")
+                self._is_first_write = False
+            self.ui_callback(text, end="")
+            get_app().invalidate()
+        return len(text)
+
+    def flush(self):
+        self.original_stdout.flush()
+        self.original_stderr.flush()

zrb/llm/app/style.py

@@ -0,0 +1,16 @@
+from prompt_toolkit.styles import Style
+
+
+def create_style() -> Style:
+    return Style.from_dict(
+        {
+            "frame.label": "bg:#000000 #ffff00",
+            "thinking": "ansigreen italic",
+            "faint": "#888888",
+            "output_field": "bg:#000000 #eeeeee",
+            "input_field": "bg:#000000 #eeeeee",
+            "text": "#eeeeee",
+            "status": "reverse",
+            "bottom-toolbar": "bg:#333333 #aaaaaa",
+        }
+    )