zrb 1.8.10__py3-none-any.whl → 1.21.29__py3-none-any.whl

zrb/builtin/llm/chat_completion.py

@@ -0,0 +1,274 @@
+import os
+
+from prompt_toolkit.completion import CompleteEvent, Completer, Completion
+from prompt_toolkit.document import Document
+
+from zrb.builtin.llm.chat_session_cmd import (
+    ADD_SUB_CMD,
+    ATTACHMENT_ADD_SUB_CMD_DESC,
+    ATTACHMENT_CLEAR_SUB_CMD_DESC,
+    ATTACHMENT_CMD,
+    ATTACHMENT_CMD_DESC,
+    ATTACHMENT_SET_SUB_CMD_DESC,
+    CLEAR_SUB_CMD,
+    HELP_CMD,
+    HELP_CMD_DESC,
+    MULTILINE_END_CMD,
+    MULTILINE_END_CMD_DESC,
+    MULTILINE_START_CMD,
+    MULTILINE_START_CMD_DESC,
+    QUIT_CMD,
+    QUIT_CMD_DESC,
+    RUN_CLI_CMD,
+    RUN_CLI_CMD_DESC,
+    SAVE_CMD,
+    SAVE_CMD_DESC,
+    SET_SUB_CMD,
+    WORKFLOW_ADD_SUB_CMD_DESC,
+    WORKFLOW_CLEAR_SUB_CMD_DESC,
+    WORKFLOW_CMD,
+    WORKFLOW_CMD_DESC,
+    WORKFLOW_SET_SUB_CMD_DESC,
+    YOLO_CMD,
+    YOLO_CMD_DESC,
+    YOLO_SET_CMD_DESC,
+    YOLO_SET_FALSE_CMD_DESC,
+    YOLO_SET_TRUE_CMD_DESC,
+)
+
+
+class ChatCompleter(Completer):
+
+    def get_completions(self, document: Document, complete_event: CompleteEvent):
+        # Slash command
+        for completion in self._complete_slash_command(document):
+            yield completion
+        for completion in self._complete_slash_file_command(document):
+            yield completion
+        # Appendix
+        for completion in self._complete_appendix(document):
+            yield completion
+
+    def _complete_slash_file_command(self, document: Document):
+        text = document.text_before_cursor
+        prefixes = []
+        for cmd in ATTACHMENT_CMD:
+            for subcmd in ADD_SUB_CMD:
+                prefixes.append(f"{cmd} {subcmd} ")
+        for prefix in prefixes:
+            if text.startswith(prefix):
+                pattern = text[len(prefix) :]
+                potential_options = self._fuzzy_path_search(pattern, dirs=False)
+                for prefixed_option in [
+                    f"{prefix}{option}" for option in potential_options
+                ]:
+                    yield Completion(
+                        prefixed_option,
+                        start_position=-len(text),
+                    )
+
+    def _complete_slash_command(self, document: Document):
+        text = document.text_before_cursor
+        if not text.startswith("/"):
+            return
+        for command, description in self._get_cmd_options().items():
+            if command.lower().startswith(text.lower()):
+                yield Completion(
+                    command,
+                    start_position=-len(text),
+                    display_meta=description,
+                )
+
+    def _complete_appendix(self, document: Document):
+        token = document.get_word_before_cursor(WORD=True)
+        prefix = "@"
+        if not token.startswith(prefix):
+            return
+        pattern = token[len(prefix) :]
+        potential_options = self._fuzzy_path_search(pattern, dirs=False)
+        for prefixed_option in [f"{prefix}{option}" for option in potential_options]:
+            yield Completion(
+                prefixed_option,
+                start_position=-len(token),
+            )
+
+    def _get_cmd_options(self):
+        cmd_options = {}
+        # Add all commands with their descriptions
+        for cmd in MULTILINE_START_CMD:
+            cmd_options[cmd] = MULTILINE_START_CMD_DESC
+        for cmd in MULTILINE_END_CMD:
+            cmd_options[cmd] = MULTILINE_END_CMD_DESC
+        for cmd in QUIT_CMD:
+            cmd_options[cmd] = QUIT_CMD_DESC
+        for cmd in WORKFLOW_CMD:
+            cmd_options[cmd] = WORKFLOW_CMD_DESC
+            for subcmd in ADD_SUB_CMD:
+                cmd_options[f"{cmd} {subcmd}"] = WORKFLOW_ADD_SUB_CMD_DESC
+            for subcmd in CLEAR_SUB_CMD:
+                cmd_options[f"{cmd} {subcmd}"] = WORKFLOW_CLEAR_SUB_CMD_DESC
+            for subcmd in SET_SUB_CMD:
+                cmd_options[f"{cmd} {subcmd}"] = WORKFLOW_SET_SUB_CMD_DESC
+        for cmd in SAVE_CMD:
+            cmd_options[cmd] = SAVE_CMD_DESC
+        for cmd in ATTACHMENT_CMD:
+            cmd_options[cmd] = ATTACHMENT_CMD_DESC
+            for subcmd in ADD_SUB_CMD:
+                cmd_options[f"{cmd} {subcmd}"] = ATTACHMENT_ADD_SUB_CMD_DESC
+            for subcmd in CLEAR_SUB_CMD:
+                cmd_options[f"{cmd} {subcmd}"] = ATTACHMENT_CLEAR_SUB_CMD_DESC
+            for subcmd in SET_SUB_CMD:
+                cmd_options[f"{cmd} {subcmd}"] = ATTACHMENT_SET_SUB_CMD_DESC
+        for cmd in YOLO_CMD:
+            cmd_options[cmd] = YOLO_CMD_DESC
+            for subcmd in SET_SUB_CMD:
+                cmd_options[f"{cmd} {subcmd} true"] = YOLO_SET_TRUE_CMD_DESC
+                cmd_options[f"{cmd} {subcmd} false"] = YOLO_SET_FALSE_CMD_DESC
+                cmd_options[f"{cmd} {subcmd}"] = YOLO_SET_CMD_DESC
+        for cmd in HELP_CMD:
+            cmd_options[cmd] = HELP_CMD_DESC
+        for cmd in RUN_CLI_CMD:
+            cmd_options[cmd] = RUN_CLI_CMD_DESC
+        return dict(sorted(cmd_options.items()))
+
+    def _fuzzy_path_search(
+        self,
+        pattern: str,
+        root: str | None = None,
+        max_results: int = 20,
+        include_hidden: bool = False,
+        case_sensitive: bool = False,
+        dirs: bool = True,
+        files: bool = True,
+    ) -> list[str]:
+        """
+        Return a list of filesystem paths under `root` that fuzzy-match `pattern`.
+        - pattern: e.g. "./some/x" or "proj util/io"
+        - include_hidden: if False skip files/dirs starting with '.'
+        - dirs/files booleans let you restrict results
+        - returns list of relative paths (from root), sorted best-first
+        """
+        search_pattern = pattern
+        if root is None:
+            # Determine root and adjust pattern if necessary
+            expanded_pattern = os.path.expanduser(pattern)
+            if os.path.isabs(expanded_pattern) or pattern.startswith("~"):
+                # For absolute paths, find the deepest existing directory
+                if os.path.isdir(expanded_pattern):
+                    root = expanded_pattern
+                    search_pattern = ""
+                else:
+                    root = os.path.dirname(expanded_pattern)
+                    while root and not os.path.isdir(root) and len(root) > 1:
+                        root = os.path.dirname(root)
+                    if not os.path.isdir(root):
+                        root = "."  # Fallback
+                        search_pattern = pattern
+                    else:
+                        try:
+                            search_pattern = os.path.relpath(expanded_pattern, root)
+                            if search_pattern == ".":
+                                search_pattern = ""
+                        except ValueError:
+                            search_pattern = os.path.basename(pattern)
+            else:
+                root = "."
+                search_pattern = pattern
+        # Normalize pattern -> tokens split on path separators or whitespace
+        search_pattern = search_pattern.strip()
+        if search_pattern:
+            raw_tokens = [t for t in search_pattern.split(os.path.sep) if t]
+        else:
+            raw_tokens = []
+        # prepare tokens (case)
+        if not case_sensitive:
+            tokens = [t.lower() for t in raw_tokens]
+        else:
+            tokens = raw_tokens
+        # specific ignore list
+        try:
+            is_recursive = os.path.abspath(os.path.expanduser(root)).startswith(
+                os.path.abspath(os.getcwd())
+            )
+        except Exception:
+            is_recursive = False
+        # walk filesystem
+        candidates: list[tuple[float, str]] = []
+        for dirpath, dirnames, filenames in os.walk(root):
+            # Filter directories
+            if not include_hidden:
+                dirnames[:] = [d for d in dirnames if not d.startswith(".")]
+            rel_dir = os.path.relpath(dirpath, root)
+            # treat '.' as empty prefix
+            if rel_dir == ".":
+                rel_dir = ""
+            # build list of entries to test depending on files/dirs flags
+            entries = []
+            if dirs:
+                entries.extend([os.path.join(rel_dir, d) for d in dirnames])
+            if files:
+                entries.extend([os.path.join(rel_dir, f) for f in filenames])
+            if not is_recursive:
+                dirnames[:] = []
+            for ent in entries:
+                # Normalize presentation: use ./ prefix for relative paths
+                display_path = ent if ent else "."
+                # Skip hidden entries unless requested (double check for rel path segments)
+                if not include_hidden:
+                    if any(
+                        seg.startswith(".") for seg in display_path.split(os.sep) if seg
+                    ):
+                        continue
+                cand = display_path.replace(os.sep, "/")  # unify separator
+                cand_cmp = cand if case_sensitive else cand.lower()
+                last_pos = 0
+                score = 0.0
+                matched_all = True
+                for token in tokens:
+                    # try contiguous substring search first
+                    idx = cand_cmp.find(token, last_pos)
+                    if idx != -1:
+                        # good match: reward contiguous early matches
+                        score += idx  # smaller idx preferred
+                        last_pos = idx + len(token)
+                    else:
+                        # fallback to subsequence matching
+                        pos = self._find_subsequence_pos(cand_cmp, token, last_pos)
+                        if pos is None:
+                            matched_all = False
+                            break
+                        # subsequence match is less preferred than contiguous substring
+                        score += pos + 0.5 * len(token)
+                        last_pos = pos + len(token)
+                if matched_all:
+                    # prefer shorter paths when score ties, so include length as tiebreaker
+                    score += 0.01 * len(cand)
+                    out = (
+                        cand
+                        if os.path.abspath(cand) == cand
+                        else os.path.join(root, cand)
+                    )
+                    candidates.append((score, out))
+        # sort by score then lexicographically and return top results
+        candidates.sort(key=lambda x: (x[0], x[1]))
+        return [p for _, p in candidates[:max_results]]
+
+    def _find_subsequence_pos(
+        self, hay: str, needle: str, start: int = 0
+    ) -> int | None:
+        """
+        Try to locate needle in hay as a subsequence starting at `start`.
+        Returns the index of the first matched character of the subsequence or None if not match.
+        """
+        if not needle:
+            return start
+        i = start
+        j = 0
+        first_pos = None
+        while i < len(hay) and j < len(needle):
+            if hay[i] == needle[j]:
+                if first_pos is None:
+                    first_pos = i
+                j += 1
+            i += 1
+        return first_pos if j == len(needle) else None

zrb/builtin/llm/chat_session.py

@@ -1,84 +1,154 @@
-"""
-This module provides functions for managing interactive chat sessions with an LLM.
-
-It handles reading user input, triggering the LLM task, and managing the
-conversation flow via XCom.
-"""
-
 import asyncio
+import sys
+from typing import TYPE_CHECKING, Any
 
+from zrb.builtin.llm.chat_session_cmd import (
+    ATTACHMENT_CMD,
+    HELP_CMD,
+    MULTILINE_END_CMD,
+    MULTILINE_START_CMD,
+    QUIT_CMD,
+    RUN_CLI_CMD,
+    SAVE_CMD,
+    WORKFLOW_CMD,
+    YOLO_CMD,
+    get_new_attachments,
+    get_new_workflows,
+    get_new_yolo_mode,
+    is_command_match,
+    print_commands,
+    print_current_attachments,
+    print_current_workflows,
+    print_current_yolo_mode,
+    run_cli_command,
+    save_final_result,
+)
+from zrb.builtin.llm.chat_trigger import llm_chat_trigger
+from zrb.config.llm_config import llm_config
 from zrb.context.any_context import AnyContext
-from zrb.util.cli.style import stylize_bold_yellow, stylize_faint
+from zrb.util.cli.markdown import render_markdown
+
+if TYPE_CHECKING:
+    from asyncio import StreamReader
+
+    from prompt_toolkit import PromptSession
 
 
 async def read_user_prompt(ctx: AnyContext) -> str:
     """
     Reads user input from the CLI for an interactive chat session.
-
-    Handles special commands like /bye, /multi, /end, and /help.
-    Triggers the LLM task and waits for the result.
-
-    Args:
-        ctx: The context object for the task.
-
-    Returns:
-        The final result from the LLM session.
+    Orchestrates the session by calling helper functions.
     """
-    _show_info(ctx)
-    final_result = ""
-    ctx.print(stylize_faint("🧑 >> ") + f"{ctx.input.message}", plain=True)
-    result = await _trigger_ask_and_wait_for_result(
-        ctx,
-        user_prompt=ctx.input.message,
-        previous_session_name=ctx.input.previous_session,
-        start_new=ctx.input.start_new,
-    )
-    if result is not None:
-        final_result = result
-    if ctx.env.get("_ZRB_WEB_ENV", "0") != "0":
-        # On web environment this won't be interactive
-        return final_result
+    print_commands(ctx)
+    is_tty: bool = ctx.is_tty
+    reader: PromptSession[Any] | StreamReader = await _setup_input_reader(is_tty)
     multiline_mode = False
-    user_inputs = []
-    while True:
+    is_first_time = True
+    current_workflows: str = ctx.input.workflows
+    current_yolo_mode: bool | str = ctx.input.yolo
+    current_attachments: str = ctx.input.attach
+    user_inputs: list[str] = []
+    final_result: str = ""
+    should_end = False
+    while not should_end:
         await asyncio.sleep(0.01)
-        ctx.print(stylize_faint("🧑 >> "), end="", plain=True)
-        user_input = input()
-        # Handle special input
-        if user_input.strip().lower() in ("/bye", "/quit"):
-            user_prompt = "\n".join(user_inputs)
-            user_inputs = []
-            result = await _trigger_ask_and_wait_for_result(ctx, user_prompt)
-            if result is not None:
-                final_result = result
-            break
-        elif user_input.strip().lower() in ("/multi"):
-            multiline_mode = True
-        elif user_input.strip().lower() in ("/end"):
-            multiline_mode = False
-            user_prompt = "\n".join(user_inputs)
-            user_inputs = []
-            result = await _trigger_ask_and_wait_for_result(ctx, user_prompt)
-            if result is not None:
-                final_result = result
-        elif user_input.strip().lower() in ("/help", "/info"):
-            _show_info(ctx)
-            continue
+        previous_session_name: str | None = (
+            ctx.input.previous_session if is_first_time else ""
+        )
+        start_new: bool = ctx.input.start_new if is_first_time else False
+        if is_first_time and ctx.input.message.strip() != "":
+            user_input = ctx.input.message
+        else:
+            # Get user input based on mode
+            if not multiline_mode:
+                ctx.print("💬 >>", plain=True)
+            user_input = await llm_chat_trigger.wait(reader, ctx)
+            if not multiline_mode:
+                ctx.print("", plain=True)
+        # At this point, is_first_time has to be False
+        if is_first_time:
+            is_first_time = False
+        # Handle user input (including slash commands)
+        if multiline_mode:
+            if is_command_match(user_input, MULTILINE_END_CMD):
+                ctx.print("", plain=True)
+                multiline_mode = False
+            else:
+                user_inputs.append(user_input)
+                continue
         else:
-            user_inputs.append(user_input)
-            if multiline_mode:
+            if is_command_match(user_input, QUIT_CMD):
+                should_end = True
+            elif is_command_match(user_input, MULTILINE_START_CMD):
+                multiline_mode = True
+                ctx.print("", plain=True)
+                continue
+            elif is_command_match(user_input, WORKFLOW_CMD):
+                current_workflows = get_new_workflows(current_workflows, user_input)
+                print_current_workflows(ctx, current_workflows)
+                continue
+            elif is_command_match(user_input, SAVE_CMD):
+                save_final_result(ctx, user_input, final_result)
+                continue
+            elif is_command_match(user_input, ATTACHMENT_CMD):
+                current_attachments = get_new_attachments(
+                    current_attachments, user_input
+                )
+                print_current_attachments(ctx, current_attachments)
                 continue
-            user_prompt = "\n".join(user_inputs)
-            user_inputs = []
-            result = await _trigger_ask_and_wait_for_result(ctx, user_prompt)
-            if result is not None:
-                final_result = result
+            elif is_command_match(user_input, YOLO_CMD):
+                current_yolo_mode = get_new_yolo_mode(current_yolo_mode, user_input)
+                print_current_yolo_mode(ctx, current_yolo_mode)
+                continue
+            elif is_command_match(user_input, RUN_CLI_CMD):
+                run_cli_command(ctx, user_input)
+                continue
+            elif is_command_match(user_input, HELP_CMD):
+                print_commands(ctx)
+                continue
+            else:
+                user_inputs.append(user_input)
+        # Trigger LLM
+        user_prompt = "\n".join(user_inputs)
+        user_inputs = []
+        result = await _trigger_ask_and_wait_for_result(
+            ctx=ctx,
+            user_prompt=user_prompt,
+            attach=current_attachments,
+            workflows=current_workflows,
+            yolo_mode=current_yolo_mode,
+            previous_session_name=previous_session_name,
+            start_new=start_new,
+        )
+        current_attachments = ""
+        final_result = final_result if result is None else result
+        if ctx.is_web_mode or not is_tty:
+            return final_result
     return final_result
 
 
+async def _setup_input_reader(
+    is_interactive: bool,
+) -> "PromptSession[Any] | StreamReader":
+    """Sets up and returns the appropriate asynchronous input reader."""
+    if is_interactive:
+        from prompt_toolkit import PromptSession
+
+        return PromptSession()
+
+    loop = asyncio.get_event_loop()
+    reader = asyncio.StreamReader(loop=loop)
+    protocol = asyncio.StreamReaderProtocol(reader)
+    await loop.connect_read_pipe(lambda: protocol, sys.stdin)
+    return reader
+
+
 async def _trigger_ask_and_wait_for_result(
     ctx: AnyContext,
     user_prompt: str,
+    attach: str,
+    workflows: str,
+    yolo_mode: bool | str,
     previous_session_name: str | None = None,
     start_new: bool = False,
 ) -> str | None:
@@ -96,9 +166,14 @@ async def _trigger_ask_and_wait_for_result(
     """
     if user_prompt.strip() == "":
         return None
-    await _trigger_ask(ctx, user_prompt, previous_session_name, start_new)
+    await _trigger_ask(
+        ctx, user_prompt, attach, workflows, yolo_mode, previous_session_name, start_new
+    )
     result = await _wait_ask_result(ctx)
-    ctx.print(f"{stylize_faint('🤖 >>')} {result}", plain=True)
+    md_result = render_markdown(result) if result is not None else ""
+    ctx.print("\n🤖 >>", plain=True)
+    ctx.print(md_result, plain=True)
+    ctx.print("", plain=True)
     return result
 
 
@@ -113,20 +188,29 @@ def get_llm_ask_input_mapping(callback_ctx: AnyContext):
         A dictionary containing the input mapping for the LLM ask task.
     """
     data = callback_ctx.xcom.ask_trigger.pop()
+    system_prompt = callback_ctx.input.system_prompt
+    if system_prompt is None or system_prompt.strip() == "":
+        system_prompt = llm_config.default_interactive_system_prompt
     return {
         "model": callback_ctx.input.model,
         "base-url": callback_ctx.input.base_url,
         "api-key": callback_ctx.input.api_key,
-        "system-prompt": callback_ctx.input.system_prompt,
+        "system-prompt": system_prompt,
         "start-new": data.get("start_new"),
         "previous-session": data.get("previous_session_name"),
         "message": data.get("message"),
+        "attach": data.get("attach"),
+        "workflows": data.get("workflows"),
+        "yolo": data.get("yolo"),
     }
 
 
 async def _trigger_ask(
     ctx: AnyContext,
     user_prompt: str,
+    attach: str,
+    workflows: str,
+    yolo_mode: bool | str,
     previous_session_name: str | None = None,
     start_new: bool = False,
 ):
@@ -146,6 +230,9 @@ async def _trigger_ask(
             "previous_session_name": previous_session_name,
             "start_new": start_new,
             "message": user_prompt,
+            "attach": attach,
+            "workflows": workflows,
+            "yolo": yolo_mode,
         }
     )
 
@@ -181,23 +268,3 @@ async def _wait_ask_session_name(ctx: AnyContext) -> str:
     while "ask_session_name" not in ctx.xcom or len(ctx.xcom.ask_session_name) == 0:
         await asyncio.sleep(0.1)
     return ctx.xcom.ask_session_name.pop()
-
-
-def _show_info(ctx: AnyContext):
-    """
-    Displays the available chat session commands to the user.
-
-    Args:
-        ctx: The context object for the task.
-    """
-    ctx.print(
-        "\n".join(
-            [
-                f"{stylize_bold_yellow('/bye')}   {stylize_faint('Quit from chat session')}",
-                f"{stylize_bold_yellow('/multi')} {stylize_faint('Start multiline input')}",
-                f"{stylize_bold_yellow('/end')}   {stylize_faint('End multiline input')}",
-                f"{stylize_bold_yellow('/help')}  {stylize_faint('Show this message')}",
-            ]
-        ),
-        plain=True,
-    )