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

zrb/builtin/llm/chat_completion.py

@@ -1,274 +0,0 @@
-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,270 +0,0 @@
-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.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.
-    Orchestrates the session by calling helper functions.
-    """
-    print_commands(ctx)
-    is_tty: bool = ctx.is_tty
-    reader: PromptSession[Any] | StreamReader = await _setup_input_reader(is_tty)
-    multiline_mode = False
-    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)
-        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:
-            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
-            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:
-    """
-    Triggers the LLM ask task and waits for the result via XCom.
-
-    Args:
-        ctx: The context object for the task.
-        user_prompt: The user's message to send to the LLM.
-        previous_session_name: The name of the previous chat session (optional).
-        start_new: Whether to start a new conversation (optional).
-
-    Returns:
-        The result from the LLM task, or None if the user prompt is empty.
-    """
-    if user_prompt.strip() == "":
-        return None
-    await _trigger_ask(
-        ctx, user_prompt, attach, workflows, yolo_mode, previous_session_name, start_new
-    )
-    result = await _wait_ask_result(ctx)
-    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
-
-
-def get_llm_ask_input_mapping(callback_ctx: AnyContext):
-    """
-    Generates the input mapping for the LLM ask task from the callback context.
-
-    Args:
-        callback_ctx: The context object for the callback.
-
-    Returns:
-        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": 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,
-):
-    """
-    Triggers the LLM ask task by pushing data to the 'ask_trigger' XCom queue.
-
-    Args:
-        ctx: The context object for the task.
-        user_prompt: The user's message to send to the LLM.
-        previous_session_name: The name of the previous chat session (optional).
-        start_new: Whether to start a new conversation (optional).
-    """
-    if previous_session_name is None:
-        previous_session_name = await _wait_ask_session_name(ctx)
-    ctx.xcom["ask_trigger"].push(
-        {
-            "previous_session_name": previous_session_name,
-            "start_new": start_new,
-            "message": user_prompt,
-            "attach": attach,
-            "workflows": workflows,
-            "yolo": yolo_mode,
-        }
-    )
-
-
-async def _wait_ask_result(ctx: AnyContext) -> str | None:
-    """
-    Waits for and retrieves the LLM task result from the 'ask_result' XCom queue.
-
-    Args:
-        ctx: The context object for the task.
-
-    Returns:
-        The result string from the LLM task.
-    """
-    while "ask_result" not in ctx.xcom or len(ctx.xcom.ask_result) == 0:
-        await asyncio.sleep(0.1)
-        if "ask_error" in ctx.xcom and len(ctx.xcom.ask_error) > 0:
-            ctx.xcom.ask_error.pop()
-            return None
-    return ctx.xcom.ask_result.pop()
-
-
-async def _wait_ask_session_name(ctx: AnyContext) -> str:
-    """
-    Waits for and retrieves the LLM chat session name from the 'ask_session_name' XCom queue.
-
-    Args:
-        ctx: The context object for the task.
-
-    Returns:
-        The session name string.
-    """
-    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()