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

zrb/builtin/llm/llm_ask.py

@@ -1,114 +1,257 @@
+import os
+from collections.abc import Callable
+from typing import TYPE_CHECKING
+
 from zrb.builtin.group import llm_group
+from zrb.builtin.llm.attachment import get_media_type
 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.api import (
+    create_get_current_location,
+    create_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,
-    apply_diff,
     list_files,
     read_from_file,
+    replace_in_file,
     search_files,
     write_to_file,
 )
+from zrb.builtin.llm.tool.note import (
+    read_contextual_note,
+    read_long_term_note,
+    write_contextual_note,
+    write_long_term_note,
+)
 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 import CFG
+from zrb.config.config import CFG
+from zrb.config.llm_config import llm_config
+from zrb.context.any_context import AnyContext
+from zrb.input.any_input import AnyInput
 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
+from zrb.util.string.conversion import to_boolean
 
-_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,
-    ),
-    BoolInput(
-        "start-new",
-        description="Start new conversation (LLM will forget everything)",
-        prompt="Start new conversation (LLM will forget everything)",
-        default=False,
-        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
+if TYPE_CHECKING:
+    from pydantic_ai import AbstractToolset, Tool, UserContent
+
+    ToolOrCallable = Tool | Callable
+
+
+def _get_toolset(ctx: AnyContext) -> list["AbstractToolset[None] | str"]:
+    cwd = os.getcwd()
+    toolsets = []
+    for config_path in [
+        os.path.join(cwd, "mcp_config.json"),
+        os.path.join(cwd, "mcp-config.json"),
+    ]:
+        if os.path.isfile(config_path):
+            toolsets.append(config_path)
+    return toolsets
+
+
+def _get_tool(ctx: AnyContext) -> list["ToolOrCallable"]:
+    tools = [
+        read_long_term_note,
+        write_long_term_note,
+        read_contextual_note,
+        write_contextual_note,
+    ]
+    if CFG.LLM_ALLOW_ANALYZE_REPO:
+        tools.append(analyze_repo)
+    if CFG.LLM_ALLOW_ANALYZE_FILE:
+        tools.append(analyze_file)
+    if CFG.LLM_ALLOW_ACCESS_LOCAL_FILE:
+        tools.append(search_files)
+        tools.append(list_files)
+        tools.append(read_from_file)
+        tools.append(replace_in_file)
+        tools.append(write_to_file)
+    if CFG.LLM_ALLOW_ACCESS_SHELL:
+        tools.append(run_shell_command)
+    if CFG.LLM_ALLOW_OPEN_WEB_PAGE:
+        tools.append(open_web_page)
+    if CFG.LLM_ALLOW_GET_CURRENT_LOCATION:
+        tools.append(create_get_current_location())
+    if CFG.LLM_ALLOW_GET_CURRENT_WEATHER:
+        tools.append(create_get_current_weather())
+    if CFG.SERPAPI_KEY != "" and CFG.LLM_ALLOW_SEARCH_INTERNET:
+        tools.append(create_search_internet_tool())
+    return tools
+
+
+def _get_default_yolo_mode(ctx: AnyContext) -> str:
+    default_value = llm_config.default_yolo_mode
+    if isinstance(default_value, list):
+        return ",".join(default_value)
+    return f"{default_value}"
+
+
+def _render_yolo_mode_input(ctx: AnyContext) -> list[str] | bool:
+    if ctx.input.yolo.strip() == "":
+        return []
+    elements = [element.strip() for element in ctx.input.yolo.split(",")]
+    if len(elements) == 1:
+        try:
+            return to_boolean(elements[0])
+        except Exception:
+            pass
+    return elements
+
+
+def _render_attach_input(ctx: AnyContext) -> "list[UserContent]":
+    from pathlib import Path
+
+    from pydantic_ai import BinaryContent
+
+    attachment_paths: list[str] = [
+        attachment_path.strip()
+        for attachment_path in ctx.input.attach.split(",")
+        if attachment_path.strip() != ""
+    ]
+    if len(attachment_paths) == 0:
+        return []
+    attachments = []
+    for attachment_path in attachment_paths:
+        attachment_path = os.path.abspath(os.path.expanduser(attachment_path))
+        media_type = get_media_type(attachment_path)
+        if media_type is None:
+            ctx.log_error(f"Cannot get media type of {attachment_path}")
+            continue
+        data = Path(attachment_path).read_bytes()
+        attachments.append(BinaryContent(data, media_type=media_type))
+    return attachments
+
+
+def _get_inputs(require_message: bool = True) -> list[AnyInput | None]:
+    return [
+        StrInput(
+            "model",
+            description="LLM Model",
+            prompt="LLM Model",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+            allow_empty=True,
         ),
-        model_api_key=lambda ctx: (
-            None if ctx.input.api_key.strip() == "" else ctx.input.api_key
+        StrInput(
+            "base-url",
+            description="LLM API Base URL",
+            prompt="LLM API Base URL",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+            allow_empty=True,
         ),
-        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
+        StrInput(
+            "api-key",
+            description="LLM API Key",
+            prompt="LLM API Key",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+            allow_empty=True,
         ),
-        message="{ctx.input.message}",
-        retries=0,
+        TextInput(
+            "system-prompt",
+            description="System prompt",
+            prompt="System prompt",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+        ),
+        TextInput(
+            "workflows",
+            description="Workflows",
+            prompt="Workflows",
+            default=lambda ctx: ",".join(llm_config.default_workflows),
+            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,
+        ),
+        StrInput(
+            "yolo",
+            description="YOLO mode (LLM Agent will start in YOLO Mode)",
+            prompt="YOLO mode (LLM Agent will start in YOLO Mode)",
+            default=_get_default_yolo_mode,
+            allow_positional_parsing=False,
+            always_prompt=False,
+        ),
+        TextInput(
+            "message",
+            description="User message",
+            prompt="Your message",
+            always_prompt=require_message,
+            allow_empty=not require_message,
+        ),
+        StrInput(
+            name="attach",
+            description="Comma separated attachments",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+        ),
+        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(
+    name="llm-ask",
+    input=_get_inputs(True),
+    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
     ),
-    alias="ask",
+    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
+    ),
+    workflows=lambda ctx: (
+        None if ctx.input.workflows.strip() == "" else ctx.input.workflows.split(",")
+    ),
+    attachment=_render_attach_input,
+    message="{ctx.input.message}",
+    tools=_get_tool,
+    toolsets=_get_toolset,
+    yolo_mode=_render_yolo_mode_input,
+    retries=0,
 )
+llm_group.add_task(llm_ask, alias="ask")
 
 llm_group.add_task(
     BaseTrigger(
         name="llm-chat",
-        input=_llm_ask_inputs,
+        input=_get_inputs(False),
         description="💬 Chat with LLM",
         queue_name="ask_trigger",
         action=read_user_prompt,
@@ -124,23 +267,3 @@ llm_group.add_task(
     ),
     alias="chat",
 )
-
-if CFG.LLM_ALLOW_ACCESS_LOCAL_FILE:
-    llm_ask.append_tool(
-        list_files,
-        read_from_file,
-        write_to_file,
-        search_files,
-        apply_diff,
-        analyze_file,
-        analyze_repo,
-    )
-
-if CFG.LLM_ALLOW_ACCESS_SHELL:
-    llm_ask.append_tool(run_shell_command)
-
-if CFG.LLM_ALLOW_ACCESS_INTERNET:
-    llm_ask.append_tool(open_web_page, search_wikipedia, search_arxiv)
-    if CFG.SERPAPI_KEY != "":
-        llm_ask.append_tool(create_search_internet_tool(CFG.SERPAPI_KEY))
-    llm_ask.append_tool(get_current_location, get_current_weather)

zrb/builtin/llm/tool/api.py

@@ -1,53 +1,75 @@
-import json
-from typing import Literal
-
-
-def get_current_location() -> str:
-    """Get current location (latitude, longitude) based on IP address.
-    Returns:
-        str: JSON string representing latitude and longitude.
-    Raises:
-        requests.RequestException: If the API request 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:
-    """Get current weather for a specific location.
-    Args:
-        latitude (float): Latitude coordinate.
-        longitude (float): Longitude coordinate.
-        temperature_unit (Literal["celsius", "fahrenheit"]): Temperature unit.
-    Returns:
-        str: JSON string with weather data.
-    Raises:
-        requests.RequestException: If the API request 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,
+from typing import Any, Callable, Literal
+
+from zrb.config.llm_config import llm_config
+
+
+def create_get_current_location() -> Callable:
+    if llm_config.default_current_location_tool is not None:
+        return llm_config.default_current_location_tool
+
+    def get_current_location() -> dict[str, float]:
+        """
+        Gets the user's current geographical location (latitude and longitude).
+
+        Example:
+        get_current_location() # Returns {'lat': 48.8584, 'lon': 2.2945}
+
+        Returns:
+            dict: Dictionary with 'lat' and 'lon' keys.
+        """
+        import requests
+
+        try:
+            response = requests.get("http://ip-api.com/json?fields=lat,lon", timeout=5)
+            response.raise_for_status()
+            return dict(response.json())
+        except requests.RequestException as e:
+            raise requests.RequestException(f"Failed to get location: {e}")
+
+    return get_current_location
+
+
+def create_get_current_weather() -> Callable:
+    if llm_config.default_current_weather_tool is not None:
+        return llm_config.default_current_weather_tool
+
+    def get_current_weather(
+        latitude: float,
+        longitude: float,
+        temperature_unit: Literal["celsius", "fahrenheit"],
+    ) -> dict[str, Any]:
+        """
+        Gets current weather conditions for a given location.
+
+        Example:
+        get_current_weather(
+            latitude=34.0522, longitude=-118.2437, temperature_unit='fahrenheit'
         )
-        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
+
+        Args:
+            latitude (float): Latitude of the location.
+            longitude (float): Longitude of the location.
+            temperature_unit (Literal): 'celsius' or 'fahrenheit'.
+
+        Returns:
+            dict: Detailed weather data.
+        """
+        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 dict(response.json())
+        except requests.RequestException as e:
+            raise requests.RequestException(f"Failed to get weather data: {e}")
+
+    return get_current_weather

zrb/builtin/llm/tool/cli.py

@@ -1,23 +1,52 @@
 import subprocess
+import sys
 
+if sys.version_info >= (3, 12):
+    from typing import TypedDict
+else:
+    from typing_extensions import TypedDict
+
+
+class ShellCommandResult(TypedDict):
+    """
+    Result of shell command execution
+
+    Attributes:
+        return_code: The return code, 0 indicating no error
+        stdout: Standard output
+        stderr: Standard error
+    """
+
+    return_code: int
+    stdout: str
+    stderr: str
+
+
+def run_shell_command(command: str) -> ShellCommandResult:
+    """
+    Executes a non-interactive shell command on the user's machine.
+
+    CRITICAL: This tool runs with user-level permissions. Explain commands that modify
+        the system (e.g., `git`, `pip`) and ask for confirmation.
+    IMPORTANT: Long-running processes should be run in the background (e.g., `command &`).
+
+    Example:
+    run_shell_command(command='ls -l')
 
-def run_shell_command(command: str) -> str:
-    """Execute a shell command and return its combined stdout and stderr.
     Args:
-        command (str): Shell command to execute on the user's system.
+        command (str): The exact shell command to be executed.
+
     Returns:
-        str: The command's output (stdout and stderr combined).
-    Raises:
-        subprocess.CalledProcessError: If the command returns a non-zero exit code.
-        subprocess.SubprocessError: If there's an issue with subprocess execution.
+        dict: return_code, stdout, and stderr.
     """
-    try:
-        output = subprocess.check_output(
-            command, shell=True, stderr=subprocess.STDOUT, text=True
-        )
-        return output
-    except subprocess.CalledProcessError as e:
-        # Include the error output in the exception message
-        raise subprocess.CalledProcessError(
-            e.returncode, e.cmd, e.output, e.stderr
-        ) from None
+    result = subprocess.run(
+        command,
+        shell=True,
+        capture_output=True,
+        text=True,
+    )
+    return {
+        "return_code": result.returncode,
+        "stdout": result.stdout,
+        "stderr": result.stderr,
+    }