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

zrb/builtin/llm/chat_session.py

@@ -1,311 +0,0 @@
-"""
-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 zrb.config.llm_config import llm_config
-from zrb.context.any_context import AnyContext
-from zrb.util.cli.style import stylize_blue, stylize_bold_yellow, stylize_faint
-from zrb.util.string.conversion import to_boolean
-
-
-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)
-    multiline_mode = False
-    current_modes = ctx.input.modes
-    current_yolo_mode = ctx.input.yolo
-    user_inputs = []
-    while True:
-        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(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, current_yolo_mode
-            )
-            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, current_yolo_mode
-            )
-            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().startswith("/yolo"):
-            yolo_mode_parts = user_input.split(" ", maxsplit=2)
-            if len(yolo_mode_parts) > 1:
-                current_yolo_mode = to_boolean(yolo_mode_parts[1])
-            ctx.print(f"Current_yolo mode: {current_yolo_mode}", plain=True)
-            ctx.print("", plain=True)
-            continue
-        elif user_input.strip().lower() in ("/help", "/info"):
-            _show_info(ctx)
-            continue
-        else:
-            user_inputs.append(user_input)
-            if multiline_mode:
-                continue
-            user_prompt = "\n".join(user_inputs)
-            user_inputs = []
-            result = await _trigger_ask_and_wait_for_result(
-                ctx, user_prompt, current_modes, current_yolo_mode
-            )
-            if result is not None:
-                final_result = 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(
-            [
-                _show_command("/bye", "Quit from chat session"),
-                _show_command("/multi", "Start multiline input"),
-                _show_command("/end", "End multiline input"),
-                _show_command("/modes", "Show current modes"),
-                _show_subcommand("<mode1,mode2,..>", "Set current modes"),
-                _show_command("/yolo", "Get current YOLO mode"),
-                _show_subcommand("<true|false>", "Set YOLO mode to true/false"),
-                _show_command("/help", "Show this message"),
-            ]
-        ),
-        plain=True,
-    )
-    ctx.print("", plain=True)
-
-
-def _show_command(command: str, description: str) -> str:
-    styled_command = stylize_bold_yellow(command.ljust(25))
-    styled_description = stylize_faint(description)
-    return f"  {styled_command} {styled_description}"
-
-
-def _show_subcommand(command: str, description: str) -> str:
-    styled_command = stylize_blue(f"    {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,
-        yolo_mode=ctx.input.yolo,
-        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):
-    """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 _read_next_line(reader, ctx: AnyContext) -> str:
-    """Reads one line of input using the provided reader."""
-    from prompt_toolkit import PromptSession
-
-    if isinstance(reader, PromptSession):
-        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,
-    yolo_mode: bool,
-    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, modes, 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 _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.
-
-    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"),
-        "modes": data.get("modes"),
-        "yolo": data.get("yolo"),
-    }
-
-
-async def _trigger_ask(
-    ctx: AnyContext,
-    user_prompt: str,
-    modes: str,
-    yolo_mode: bool,
-    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,
-            "modes": modes,
-            "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()

zrb/builtin/llm/history.py

@@ -1,71 +0,0 @@
-import json
-import os
-from typing import Any
-
-from zrb.config.config import CFG
-from zrb.context.any_shared_context import AnySharedContext
-from zrb.task.llm.conversation_history_model import ConversationHistory
-from zrb.util.file import read_file, write_file
-
-
-def read_chat_conversation(ctx: AnySharedContext) -> dict[str, Any] | list | None:
-    """Reads conversation history from the session file.
-    Returns the raw dictionary or list loaded from JSON, or None if not found/empty.
-    The LLMTask will handle parsing this into ConversationHistory.
-    """
-    if ctx.input.start_new:
-        return None  # Indicate no history to load
-    previous_session_name = ctx.input.previous_session
-    if not previous_session_name:  # Check for empty string or None
-        last_session_file_path = os.path.join(CFG.LLM_HISTORY_DIR, "last-session")
-        if os.path.isfile(last_session_file_path):
-            previous_session_name = read_file(last_session_file_path).strip()
-            if not previous_session_name:  # Handle empty last-session file
-                return None
-        else:
-            return None  # No previous session specified and no last session found
-    conversation_file_path = os.path.join(
-        CFG.LLM_HISTORY_DIR, f"{previous_session_name}.json"
-    )
-    if not os.path.isfile(conversation_file_path):
-        ctx.log_warning(f"History file not found: {conversation_file_path}")
-        return None
-    try:
-        content = read_file(conversation_file_path)
-        if not content.strip():
-            ctx.log_warning(f"History file is empty: {conversation_file_path}")
-            return None
-        # Return the raw loaded data (dict or list)
-        return json.loads(content)
-    except json.JSONDecodeError:
-        ctx.log_warning(
-            f"Could not decode JSON from history file '{conversation_file_path}'. "
-            "Treating as empty history."
-        )
-        return None
-    except Exception as e:
-        ctx.log_warning(
-            f"Error reading history file '{conversation_file_path}': {e}. "
-            "Treating as empty history."
-        )
-        return None
-
-
-def write_chat_conversation(ctx: AnySharedContext, history_data: ConversationHistory):
-    """Writes the conversation history data (including context) to a session file."""
-    os.makedirs(CFG.LLM_HISTORY_DIR, exist_ok=True)
-    current_session_name = ctx.session.name
-    if not current_session_name:
-        ctx.log_warning("Cannot write history: Session name is empty.")
-        return
-    conversation_file_path = os.path.join(
-        CFG.LLM_HISTORY_DIR, f"{current_session_name}.json"
-    )
-    try:
-        # Use model_dump_json to serialize the Pydantic model
-        write_file(conversation_file_path, history_data.model_dump_json(indent=2))
-        # Update the last-session pointer
-        last_session_file_path = os.path.join(CFG.LLM_HISTORY_DIR, "last-session")
-        write_file(last_session_file_path, current_session_name)
-    except Exception as e:
-        ctx.log_error(f"Error writing history file '{conversation_file_path}': {e}")

zrb/builtin/llm/input.py

@@ -1,27 +0,0 @@
-import os
-
-from zrb.context.any_shared_context import AnySharedContext
-from zrb.input.str_input import StrInput
-from zrb.util.file import read_file
-from zrb.util.string.conversion import to_pascal_case
-
-
-class PreviousSessionInput(StrInput):
-
-    def to_html(self, ctx: AnySharedContext) -> str:
-        name = self.name
-        description = self.description
-        default = self.get_default_str(ctx)
-        script = read_file(
-            file_path=os.path.join(os.path.dirname(__file__), "previous-session.js"),
-            replace_map={
-                "CURRENT_INPUT_NAME": name,
-                "CurrentPascalInputName": to_pascal_case(name),
-            },
-        )
-        return "\n".join(
-            [
-                f'<input name="{name}" placeholder="{description}" value="{default}" />',
-                f"<script>{script}</script>",
-            ]
-        )

zrb/builtin/llm/llm_ask.py

@@ -1,187 +0,0 @@
-from zrb.builtin.group import llm_group
-from zrb.builtin.llm.chat_session import get_llm_ask_input_mapping, read_user_prompt
-from zrb.builtin.llm.history import read_chat_conversation, write_chat_conversation
-from zrb.builtin.llm.input import PreviousSessionInput
-from zrb.builtin.llm.tool.api import get_current_location, get_current_weather
-from zrb.builtin.llm.tool.cli import run_shell_command
-from zrb.builtin.llm.tool.code import analyze_repo
-from zrb.builtin.llm.tool.file import (
-    analyze_file,
-    list_files,
-    read_from_file,
-    read_many_files,
-    replace_in_file,
-    search_files,
-    write_many_files,
-    write_to_file,
-)
-from zrb.builtin.llm.tool.web import (
-    create_search_internet_tool,
-    open_web_page,
-    search_arxiv,
-    search_wikipedia,
-)
-from zrb.callback.callback import Callback
-from zrb.config.config import CFG
-from zrb.config.llm_config import llm_config
-from zrb.input.bool_input import BoolInput
-from zrb.input.str_input import StrInput
-from zrb.input.text_input import TextInput
-from zrb.task.base_trigger import BaseTrigger
-from zrb.task.llm_task import LLMTask
-
-_llm_ask_inputs = [
-    StrInput(
-        "model",
-        description="LLM Model",
-        prompt="LLM Model",
-        default="",
-        allow_positional_parsing=False,
-        always_prompt=False,
-        allow_empty=True,
-    ),
-    StrInput(
-        "base-url",
-        description="LLM API Base URL",
-        prompt="LLM API Base URL",
-        default="",
-        allow_positional_parsing=False,
-        always_prompt=False,
-        allow_empty=True,
-    ),
-    StrInput(
-        "api-key",
-        description="LLM API Key",
-        prompt="LLM API Key",
-        default="",
-        allow_positional_parsing=False,
-        always_prompt=False,
-        allow_empty=True,
-    ),
-    TextInput(
-        "system-prompt",
-        description="System prompt",
-        prompt="System prompt",
-        default="",
-        allow_positional_parsing=False,
-        always_prompt=False,
-    ),
-    TextInput(
-        "modes",
-        description="Modes",
-        prompt="Modes",
-        default="coding",
-        allow_positional_parsing=False,
-        always_prompt=False,
-    ),
-    BoolInput(
-        "start-new",
-        description="Start new session (LLM Agent will forget past conversation)",
-        prompt="Start new session (LLM Agent will forget past conversation)",
-        default=False,
-        allow_positional_parsing=False,
-        always_prompt=False,
-    ),
-    BoolInput(
-        "yolo",
-        description="YOLO mode (LLM Agent will start in YOLO Mode)",
-        prompt="YOLO mode (LLM Agent will start in YOLO Mode)",
-        default=lambda ctx: llm_config.default_yolo_mode,
-        allow_positional_parsing=False,
-        always_prompt=False,
-    ),
-    TextInput("message", description="User message", prompt="Your message"),
-    PreviousSessionInput(
-        "previous-session",
-        description="Previous conversation session",
-        prompt="Previous conversation session (can be empty)",
-        allow_positional_parsing=False,
-        allow_empty=True,
-        always_prompt=False,
-    ),
-]
-
-llm_ask: LLMTask = llm_group.add_task(
-    LLMTask(
-        name="llm-ask",
-        input=_llm_ask_inputs,
-        description="❓ Ask LLM",
-        model=lambda ctx: None if ctx.input.model.strip() == "" else ctx.input.model,
-        model_base_url=lambda ctx: (
-            None if ctx.input.base_url.strip() == "" else ctx.input.base_url
-        ),
-        model_api_key=lambda ctx: (
-            None if ctx.input.api_key.strip() == "" else ctx.input.api_key
-        ),
-        conversation_history_reader=read_chat_conversation,
-        conversation_history_writer=write_chat_conversation,
-        system_prompt=lambda ctx: (
-            None if ctx.input.system_prompt.strip() == "" else ctx.input.system_prompt
-        ),
-        modes=lambda ctx: (
-            None if ctx.input.modes.strip() == "" else ctx.input.modes.split(",")
-        ),
-        message="{ctx.input.message}",
-        is_yolo_mode="{ctx.input.yolo}",
-        retries=0,
-    ),
-    alias="ask",
-)
-
-llm_group.add_task(
-    BaseTrigger(
-        name="llm-chat",
-        input=_llm_ask_inputs,
-        description="💬 Chat with LLM",
-        queue_name="ask_trigger",
-        action=read_user_prompt,
-        callback=Callback(
-            task=llm_ask,
-            input_mapping=get_llm_ask_input_mapping,
-            result_queue="ask_result",
-            error_queue="ask_error",
-            session_name_queue="ask_session_name",
-        ),
-        retries=0,
-        cli_only=True,
-    ),
-    alias="chat",
-)
-
-if CFG.LLM_ALLOW_ANALYZE_REPO:
-    llm_ask.append_tool(analyze_repo)
-
-if CFG.LLM_ALLOW_ANALYZE_FILE:
-    llm_ask.append_tool(analyze_file)
-
-if CFG.LLM_ALLOW_ACCESS_LOCAL_FILE:
-    llm_ask.append_tool(
-        search_files,
-        list_files,
-        read_from_file,
-        read_many_files,
-        replace_in_file,
-        write_to_file,
-        write_many_files,
-    )
-
-if CFG.LLM_ALLOW_ACCESS_SHELL:
-    llm_ask.append_tool(run_shell_command)
-
-if CFG.LLM_ALLOW_OPEN_WEB_PAGE:
-    llm_ask.append_tool(open_web_page)
-
-if CFG.LLM_ALLOW_SEARCH_WIKIPEDIA:
-    llm_ask.append_tool(search_wikipedia)
-
-if CFG.LLM_ALLOW_SEARCH_ARXIV:
-    llm_ask.append_tool(search_arxiv)
-
-if CFG.LLM_ALLOW_GET_CURRENT_LOCATION:
-    llm_ask.append_tool(get_current_location)
-
-if CFG.LLM_ALLOW_GET_CURRENT_WEATHER:
-    llm_ask.append_tool(get_current_weather)
-
-if CFG.SERPAPI_KEY != "" and CFG.LLM_ALLOW_SEARCH_INTERNET:
-    llm_ask.append_tool(create_search_internet_tool(CFG.SERPAPI_KEY))

zrb/builtin/llm/previous-session.js

@@ -1,21 +0,0 @@
-async function updatePreviousSession(event) {
-    const currentInput = event.target;
-    if (currentInput.name === "CURRENT_INPUT_NAME") {
-        return
-    }
-    const previousSessionInput = submitTaskForm.querySelector('[name="CURRENT_INPUT_NAME"]');
-    if (previousSessionInput) {
-        const currentSessionName = cfg.SESSION_NAME
-        previousSessionInput.value = currentSessionName;
-    }
-}
-
-document.getElementById("submit-task-form").querySelectorAll("input[name], textarea[name]").forEach((element) => {
-    element.addEventListener("input", updatePreviousSession);
-    element.addEventListener("keyup", updatePreviousSession);
-});
-
-document.getElementById("submit-task-form").querySelectorAll("select[name]").forEach((element) => {
-    element.addEventListener("change", updatePreviousSession);
-});
-

zrb/builtin/llm/tool/api.py

@@ -1,71 +0,0 @@
-import json
-from typing import Literal
-
-
-def get_current_location() -> str:
-    """
-    Fetches the user's current geographical location based on their IP address.
-
-    Use this tool when the user asks "Where am I?", "What is my current
-    location?", or has a query that requires knowing their location to be
-    answered.
-
-    Returns:
-        str: A JSON string containing the 'lat' and 'lon' of the current
-            location. Example: '{"lat": 48.8584, "lon": 2.2945}'
-    Raises:
-        requests.RequestException: If the API request to the location service
-            fails.
-    """
-    import requests
-
-    try:
-        response = requests.get("http://ip-api.com/json?fields=lat,lon", timeout=5)
-        response.raise_for_status()
-        return json.dumps(response.json())
-    except requests.RequestException as e:
-        raise requests.RequestException(f"Failed to get location: {e}") from None
-
-
-def get_current_weather(
-    latitude: float,
-    longitude: float,
-    temperature_unit: Literal["celsius", "fahrenheit"],
-) -> str:
-    """
-    Retrieves the current weather conditions for a given geographical location.
-
-    Use this tool when the user asks about the weather. If the user does not
-    provide a location, first use the `get_current_location` tool to
-    determine their location.
-
-    Args:
-        latitude (float): The latitude of the location.
-        longitude (float): The longitude of the location.
-        temperature_unit (Literal["celsius", "fahrenheit"]): The desired unit
-            for the temperature reading.
-
-    Returns:
-        str: A JSON string containing detailed weather data, including
-            temperature, wind speed, and weather code.
-    Raises:
-        requests.RequestException: If the API request to the weather service
-            fails.
-    """
-    import requests
-
-    try:
-        response = requests.get(
-            "https://api.open-meteo.com/v1/forecast",
-            params={
-                "latitude": latitude,
-                "longitude": longitude,
-                "temperature_unit": temperature_unit,
-                "current_weather": True,
-            },
-            timeout=5,
-        )
-        response.raise_for_status()
-        return json.dumps(response.json())
-    except requests.RequestException as e:
-        raise requests.RequestException(f"Failed to get weather data: {e}") from None