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

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",
+        }
+    )