zrb 1.13.1__py3-none-any.whl → 1.21.33__py3-none-any.whl

zrb/builtin/llm/chat_session.py

@@ -1,16 +1,37 @@
-"""
-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:
@@ -18,114 +39,97 @@ 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.
     """
-    _show_info(ctx)
-    final_result = await _handle_initial_message(ctx)
-    if ctx.is_web_mode:
-        return final_result
-    is_tty = ctx.is_tty
-    reader = await _setup_input_reader(is_tty)
+    print_commands(ctx)
+    is_tty: bool = ctx.is_tty
+    reader: PromptSession[Any] | StreamReader = await _setup_input_reader(is_tty)
     multiline_mode = False
-    current_modes = ctx.input.modes
-    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)
-        # Get user input based on mode
-        if not multiline_mode:
-            ctx.print("💬 >>", plain=True)
-        user_input = await _read_next_line(is_tty, reader, ctx)
-        if not multiline_mode:
-            ctx.print("", plain=True)
-        # Handle user input
-        if user_input.strip().lower() in ("/bye", "/quit", "/q", "/exit"):
-            user_prompt = "\n".join(user_inputs)
-            user_inputs = []
-            result = await _trigger_ask_and_wait_for_result(
-                ctx, user_prompt, current_modes
-            )
-            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",):
-            ctx.print("", plain=True)
-            multiline_mode = False
-            user_prompt = "\n".join(user_inputs)
-            user_inputs = []
-            result = await _trigger_ask_and_wait_for_result(
-                ctx, user_prompt, current_modes
-            )
-            if result is not None:
-                final_result = result
-        elif user_input.strip().lower().startswith("/mode"):
-            mode_parts = user_input.split(" ", maxsplit=2)
-            if len(mode_parts) > 1:
-                current_modes = mode_parts[1]
-            ctx.print(f"Current mode: {current_modes}", plain=True)
-            ctx.print("", plain=True)
-            continue
-        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:
-            user_inputs.append(user_input)
-            if multiline_mode:
+            # 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
-            user_prompt = "\n".join(user_inputs)
-            user_inputs = []
-            result = await _trigger_ask_and_wait_for_result(
-                ctx, user_prompt, current_modes
-            )
-            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
 
 
-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(
-            [
-                _format_info_line("/bye", "Quit from chat session"),
-                _format_info_line("/multi", "Start multiline input"),
-                _format_info_line("/end", "End multiline input"),
-                _format_info_line("/modes", "Show current modes"),
-                _format_info_line("/modes <mode1,mode2,..>", "Set current modes"),
-                _format_info_line("/help", "Show this message"),
-            ]
-        ),
-        plain=True,
-    )
-
-
-def _format_info_line(command: str, description: str) -> str:
-    styled_command = stylize_bold_yellow(command.ljust(25))
-    styled_description = stylize_faint(description)
-    return f"  {styled_command} {styled_description}"
-
-
-async def _handle_initial_message(ctx: AnyContext) -> str:
-    """Processes the initial message from the command line."""
-    if not ctx.input.message or ctx.input.message.strip() == "":
-        return ""
-    ctx.print("💬 >>", plain=True)
-    ctx.print(ctx.input.message, plain=True)
-    ctx.print("", plain=True)
-    result = await _trigger_ask_and_wait_for_result(
-        ctx,
-        user_prompt=ctx.input.message,
-        modes=ctx.input.modes,
-        previous_session_name=ctx.input.previous_session,
-        start_new=ctx.input.start_new,
-    )
-    return result if result is not None else ""
-
-
-async def _setup_input_reader(is_interactive: bool):
+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
@@ -139,24 +143,12 @@ async def _setup_input_reader(is_interactive: bool):
     return reader
 
 
-async def _read_next_line(is_interactive: bool, reader, ctx: AnyContext) -> str:
-    """Reads one line of input using the provided reader."""
-    if is_interactive:
-        return await reader.prompt_async()
-
-    line_bytes = await reader.readline()
-    if not line_bytes:
-        return "/bye"  # Signal to exit
-
-    user_input = line_bytes.decode().strip()
-    ctx.print(user_input, plain=True)
-    return user_input
-
-
 async def _trigger_ask_and_wait_for_result(
     ctx: AnyContext,
     user_prompt: str,
-    modes: str,
+    attach: str,
+    workflows: str,
+    yolo_mode: bool | str,
     previous_session_name: str | None = None,
     start_new: bool = False,
 ) -> str | None:
@@ -174,29 +166,17 @@ async def _trigger_ask_and_wait_for_result(
     """
     if user_prompt.strip() == "":
         return None
-    await _trigger_ask(ctx, user_prompt, modes, 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)
-    md_result = _render_markdown(result) if result is not None else ""
+    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 _render_markdown(markdown_text: str) -> str:
-    """
-    Renders Markdown to a string, ensuring link URLs are visible.
-    """
-    from rich.console import Console
-    from rich.markdown import Markdown
-
-    console = Console()
-    markdown = Markdown(markdown_text, hyperlinks=False)
-    with console.capture() as capture:
-        console.print(markdown)
-    return capture.get()
-
-
 def get_llm_ask_input_mapping(callback_ctx: AnyContext):
     """
     Generates the input mapping for the LLM ask task from the callback context.
@@ -219,14 +199,18 @@ def get_llm_ask_input_mapping(callback_ctx: AnyContext):
         "start-new": data.get("start_new"),
         "previous-session": data.get("previous_session_name"),
         "message": data.get("message"),
-        "modes": data.get("modes"),
+        "attach": data.get("attach"),
+        "workflows": data.get("workflows"),
+        "yolo": data.get("yolo"),
     }
 
 
 async def _trigger_ask(
     ctx: AnyContext,
     user_prompt: str,
-    modes: str,
+    attach: str,
+    workflows: str,
+    yolo_mode: bool | str,
     previous_session_name: str | None = None,
     start_new: bool = False,
 ):
@@ -246,7 +230,9 @@ async def _trigger_ask(
             "previous_session_name": previous_session_name,
             "start_new": start_new,
             "message": user_prompt,
-            "modes": modes,
+            "attach": attach,
+            "workflows": workflows,
+            "yolo": yolo_mode,
         }
     )
 

zrb/builtin/llm/chat_session_cmd.py

@@ -0,0 +1,288 @@
+import os
+import subprocess
+
+from zrb.context.any_context import AnyContext
+from zrb.task.llm.workflow import get_available_workflows
+from zrb.util.cli.markdown import render_markdown
+from zrb.util.cli.style import (
+    stylize_blue,
+    stylize_bold_yellow,
+    stylize_error,
+    stylize_faint,
+)
+from zrb.util.file import write_file
+from zrb.util.markdown import make_markdown_section
+from zrb.util.string.conversion import FALSE_STRS, TRUE_STRS, to_boolean
+
+MULTILINE_START_CMD = ["/multi", "/multiline"]
+MULTILINE_END_CMD = ["/end"]
+QUIT_CMD = ["/bye", "/quit", "/q", "/exit"]
+WORKFLOW_CMD = ["/workflow", "/workflows", "/skill", "/skills", "/w"]
+SAVE_CMD = ["/save", "/s"]
+ATTACHMENT_CMD = ["/attachment", "/attachments", "/attach"]
+YOLO_CMD = ["/yolo"]
+HELP_CMD = ["/help", "/info"]
+ADD_SUB_CMD = ["add"]
+SET_SUB_CMD = ["set"]
+CLEAR_SUB_CMD = ["clear"]
+RUN_CLI_CMD = ["/run", "/exec", "/execute", "/cmd", "/cli", "!"]
+
+# Command display constants
+MULTILINE_START_CMD_DESC = "Start multiline input"
+MULTILINE_END_CMD_DESC = "End multiline input"
+QUIT_CMD_DESC = "Quit from chat session"
+WORKFLOW_CMD_DESC = "Show active workflows"
+WORKFLOW_ADD_SUB_CMD_DESC = (
+    "Add active workflow "
+    f"(e.g., `{WORKFLOW_CMD[0]} {ADD_SUB_CMD[0]} coding,researching`)"
+)
+WORKFLOW_SET_SUB_CMD_DESC = (
+    "Set active workflows " f"(e.g., `{WORKFLOW_CMD[0]} {SET_SUB_CMD[0]} coding,`)"
+)
+WORKFLOW_CLEAR_SUB_CMD_DESC = "Deactivate all workflows"
+SAVE_CMD_DESC = f"Save last response to a file (e.g., `{SAVE_CMD[0]} conclusion.md`)"
+ATTACHMENT_CMD_DESC = "Show current attachment"
+ATTACHMENT_ADD_SUB_CMD_DESC = (
+    "Attach a file " f"(e.g., `{ATTACHMENT_CMD[0]} {ADD_SUB_CMD[0]} ./logo.png`)"
+)
+ATTACHMENT_SET_SUB_CMD_DESC = (
+    "Set attachments "
+    f"(e.g., `{ATTACHMENT_CMD[0]} {SET_SUB_CMD[0]} ./logo.png,./diagram.png`)"
+)
+ATTACHMENT_CLEAR_SUB_CMD_DESC = "Clear attachment"
+YOLO_CMD_DESC = "Show/manipulate current YOLO mode"
+YOLO_SET_CMD_DESC = (
+    "Assign YOLO tools "
+    f"(e.g., `{YOLO_CMD[0]} {SET_SUB_CMD[0]} read_from_file,analyze_file`)"
+)
+YOLO_SET_TRUE_CMD_DESC = "Activate YOLO mode for all tools"
+YOLO_SET_FALSE_CMD_DESC = "Deactivate YOLO mode for all tools"
+RUN_CLI_CMD_DESC = "Run a non-interactive CLI command"
+HELP_CMD_DESC = "Show info/help"
+
+
+def print_current_yolo_mode(
+    ctx: AnyContext, current_yolo_mode_value: str | bool
+) -> None:
+    yolo_mode_str = (
+        current_yolo_mode_value if current_yolo_mode_value != "" else "*Not Set*"
+    )
+    ctx.print(render_markdown(f"🎲 Current YOLO mode: {yolo_mode_str}"), plain=True)
+    ctx.print("", plain=True)
+
+
+def print_current_attachments(ctx: AnyContext, current_attachments_value: str) -> None:
+    attachments_str = (
+        current_attachments_value if current_attachments_value != "" else "*Not Set*"
+    )
+    ctx.print(render_markdown(f"📎 Current attachments: {attachments_str}"), plain=True)
+    ctx.print("", plain=True)
+
+
+def print_current_workflows(ctx: AnyContext, current_workflows_value: str) -> None:
+    available_workflows = get_available_workflows()
+    available_workflows_str = (
+        ", ".join(sorted([workflow_name for workflow_name in available_workflows]))
+        if len(available_workflows) > 0
+        else "*No Available Workflow*"
+    )
+    current_workflows_str = (
+        current_workflows_value
+        if current_workflows_value != ""
+        else "*No Active Workflow*"
+    )
+    ctx.print(
+        render_markdown(
+            "\n".join(
+                [
+                    f"- 🔄 Current workflows   : {current_workflows_str}",
+                    f"- 📚 Available workflows : {available_workflows_str}",
+                ]
+            )
+        ),
+        plain=True,
+    )
+    ctx.print("", plain=True)
+
+
+def save_final_result(ctx: AnyContext, user_input: str, final_result: str) -> None:
+    save_path = get_command_param(user_input, SAVE_CMD)
+    save_path = os.path.expanduser(save_path)
+    if os.path.exists(save_path):
+        ctx.print(
+            stylize_error(f"Cannot save to existing file: {save_path}"),
+            plain=True,
+        )
+        return
+    write_file(save_path, final_result)
+    ctx.print(f"Response saved to {save_path}", plain=True)
+
+
+def run_cli_command(ctx: AnyContext, user_input: str) -> None:
+    command = get_command_param(user_input, RUN_CLI_CMD)
+    result = subprocess.run(
+        command,
+        shell=True,
+        capture_output=True,
+        text=True,
+    )
+    ctx.print(
+        render_markdown(
+            make_markdown_section(
+                f"`{command}`",
+                "\n".join(
+                    [
+                        make_markdown_section("📤 Stdout", result.stdout, as_code=True),
+                        make_markdown_section("🚫 Stderr", result.stderr, as_code=True),
+                        make_markdown_section(
+                            "🎯 Return code", f"Return Code: {result.returncode}"
+                        ),
+                    ]
+                ),
+            )
+        ),
+        plain=True,
+    )
+    ctx.print("", plain=True)
+
+
+def get_new_yolo_mode(old_yolo_mode: str | bool, user_input: str) -> str | bool:
+    new_yolo_mode = get_command_param(user_input, YOLO_CMD)
+    if new_yolo_mode != "":
+        if new_yolo_mode in TRUE_STRS or new_yolo_mode in FALSE_STRS:
+            return to_boolean(new_yolo_mode)
+        return new_yolo_mode
+    if isinstance(old_yolo_mode, bool):
+        return old_yolo_mode
+    return _normalize_comma_separated_str(old_yolo_mode)
+
+
+def get_new_attachments(old_attachment: str, user_input: str) -> str:
+    if not is_command_match(user_input, ATTACHMENT_CMD):
+        return _normalize_comma_separated_str(old_attachment)
+    if is_command_match(user_input, ATTACHMENT_CMD, SET_SUB_CMD):
+        return _normalize_comma_separated_str(
+            get_command_param(user_input, ATTACHMENT_CMD, SET_SUB_CMD)
+        )
+    if is_command_match(user_input, ATTACHMENT_CMD, CLEAR_SUB_CMD):
+        return ""
+    if is_command_match(user_input, ATTACHMENT_CMD, ADD_SUB_CMD):
+        new_attachment = get_command_param(user_input, ATTACHMENT_CMD, ADD_SUB_CMD)
+        return _normalize_comma_separated_str(
+            ",".join([old_attachment, new_attachment])
+        )
+    return old_attachment
+
+
+def get_new_workflows(old_workflow: str, user_input: str) -> str:
+    if not is_command_match(user_input, WORKFLOW_CMD):
+        return _normalize_comma_separated_str(old_workflow)
+    if is_command_match(user_input, WORKFLOW_CMD, SET_SUB_CMD):
+        return _normalize_comma_separated_str(
+            get_command_param(user_input, WORKFLOW_CMD, SET_SUB_CMD)
+        )
+    if is_command_match(user_input, WORKFLOW_CMD, CLEAR_SUB_CMD):
+        return ""
+    if is_command_match(user_input, WORKFLOW_CMD, ADD_SUB_CMD):
+        new_workflow = get_command_param(user_input, WORKFLOW_CMD, ADD_SUB_CMD)
+        return _normalize_comma_separated_str(",".join([old_workflow, new_workflow]))
+    return _normalize_comma_separated_str(old_workflow)
+
+
+def _normalize_comma_separated_str(comma_separated_str: str) -> str:
+    return ",".join(
+        [
+            workflow_name.strip()
+            for workflow_name in comma_separated_str.split(",")
+            if workflow_name.strip() != ""
+        ]
+    )
+
+
+def get_command_param(user_input: str, *cmd_patterns: list[str]) -> str:
+    if not is_command_match(user_input, *cmd_patterns):
+        return ""
+    parts = [part for part in user_input.split(" ") if part.strip() != ""]
+    if len(parts) <= len(cmd_patterns):
+        return ""
+    params = parts[len(cmd_patterns) :]
+    return " ".join(params)
+
+
+def is_command_match(user_input: str, *cmd_patterns: list[str]) -> bool:
+    parts = [part for part in user_input.split(" ") if part.strip() != ""]
+    if len(cmd_patterns) > len(parts):
+        return False
+    for index, cmd_pattern in enumerate(cmd_patterns):
+        part = parts[index]
+        if part.lower() not in cmd_pattern:
+            return False
+    return True
+
+
+def print_commands(ctx: AnyContext):
+    """
+    Displays the available chat session commands to the user.
+    Args:
+        ctx: The context object for the task.
+    """
+    ctx.print(
+        "\n".join(
+            [
+                _show_command(QUIT_CMD[0], QUIT_CMD_DESC),
+                _show_command(MULTILINE_START_CMD[0], MULTILINE_START_CMD_DESC),
+                _show_command(MULTILINE_END_CMD[0], MULTILINE_END_CMD_DESC),
+                _show_command(ATTACHMENT_CMD[0], ATTACHMENT_CMD_DESC),
+                _show_subcommand(
+                    ADD_SUB_CMD[0], "<file-path>", ATTACHMENT_ADD_SUB_CMD_DESC
+                ),
+                _show_subcommand(
+                    SET_SUB_CMD[0],
+                    "<file1-path,file2-path,...>",
+                    ATTACHMENT_SET_SUB_CMD_DESC,
+                ),
+                _show_subcommand(CLEAR_SUB_CMD[0], "", ATTACHMENT_CLEAR_SUB_CMD_DESC),
+                _show_command(WORKFLOW_CMD[0], WORKFLOW_CMD_DESC),
+                _show_subcommand(
+                    ADD_SUB_CMD[0], "<workflow>", WORKFLOW_ADD_SUB_CMD_DESC
+                ),
+                _show_subcommand(
+                    SET_SUB_CMD[0],
+                    "<workflow1,workflow2,..>",
+                    WORKFLOW_SET_SUB_CMD_DESC,
+                ),
+                _show_subcommand(CLEAR_SUB_CMD[0], "", WORKFLOW_CLEAR_SUB_CMD_DESC),
+                _show_command(f"{SAVE_CMD[0]}", SAVE_CMD_DESC),
+                _show_command(YOLO_CMD[0], YOLO_CMD_DESC),
+                _show_subcommand(SET_SUB_CMD[0], "true", YOLO_SET_TRUE_CMD_DESC),
+                _show_subcommand(SET_SUB_CMD[0], "false", YOLO_SET_FALSE_CMD_DESC),
+                _show_subcommand(
+                    SET_SUB_CMD[0], "<tool1,tool2,tool2>", YOLO_SET_CMD_DESC
+                ),
+                _show_command(RUN_CLI_CMD[0], ""),
+                _show_command_param("<cli-command>", RUN_CLI_CMD_DESC),
+                _show_command(HELP_CMD[0], HELP_CMD_DESC),
+            ]
+        ),
+        plain=True,
+    )
+    ctx.print("", plain=True)
+
+
+def _show_command(command: str, description: str) -> str:
+    styled_command = stylize_bold_yellow(command.ljust(37))
+    styled_description = stylize_faint(description)
+    return f"  {styled_command} {styled_description}"
+
+
+def _show_subcommand(subcommand: str, param: str, description: str) -> str:
+    styled_subcommand = stylize_bold_yellow(f"    {subcommand}")
+    styled_param = stylize_blue(param.ljust(32 - len(subcommand)))
+    styled_description = stylize_faint(description)
+    return f"  {styled_subcommand} {styled_param} {styled_description}"
+
+
+def _show_command_param(param: str, description: str) -> str:
+    styled_param = stylize_blue(f"    {param}".ljust(37))
+    styled_description = stylize_faint(description)
+    return f"  {styled_param} {styled_description}"