zrb 1.5.4__py3-none-any.whl → 1.5.6__py3-none-any.whl

zrb/__init__.py

@@ -43,6 +43,7 @@ from zrb.task.base_task import BaseTask
 from zrb.task.base_trigger import BaseTrigger
 from zrb.task.cmd_task import CmdTask
 from zrb.task.http_check import HttpCheck
+from zrb.task.llm.history import ConversationHistoryData
 from zrb.task.llm_task import LLMTask
 from zrb.task.make_task import make_task
 from zrb.task.rsync_task import RsyncTask
@@ -91,6 +92,7 @@ assert BaseTrigger
 assert RsyncTask
 assert Task
 assert LLMTask
+assert ConversationHistoryData
 assert Session
 assert AnyContext
 assert Context

zrb/__main__.py

@@ -1,26 +1,51 @@
+import logging
 import os
 import sys
 
-from zrb.config import INIT_MODULES, INIT_SCRIPTS
+from zrb.config import INIT_MODULES, INIT_SCRIPTS, LOGGER, LOGGING_LEVEL
 from zrb.runner.cli import cli
-from zrb.util.cli.style import stylize_error, stylize_warning
+from zrb.util.cli.style import stylize_error, stylize_faint, stylize_warning
 from zrb.util.group import NodeNotFoundError
 from zrb.util.load import load_file, load_module
 
 
+# Custom Formatter for faint styling
+class FaintFormatter(logging.Formatter):
+
+    def __init__(self, fmt=None, datefmt=None):
+        default_fmt = "%(asctime)s %(levelname)s: %(message)s"
+        default_datefmt = "%Y-%m-%d %H:%M:%S"
+        super().__init__(fmt=fmt or default_fmt, datefmt=datefmt or default_datefmt)
+
+    def format(self, record):
+        log_msg = super().format(record)
+        return stylize_faint(log_msg)
+
+
 def serve_cli():
+    LOGGER.setLevel(LOGGING_LEVEL)
+    # Remove existing handlers to avoid duplicates/default formatting
+    for handler in LOGGER.handlers[:]:
+        LOGGER.removeHandler(handler)
+    handler = logging.StreamHandler()
+    handler.setFormatter(FaintFormatter())
+    LOGGER.addHandler(handler)
+    # --- End Logging Configuration ---
     try:
         # load init modules
         for init_module in INIT_MODULES:
+            LOGGER.info(f"Loading {init_module}")
             load_module(init_module)
         zrb_init_path_list = _get_zrb_init_path_list()
         # load init scripts
         for init_script in INIT_SCRIPTS:
             abs_init_script = os.path.abspath(os.path.expanduser(init_script))
             if abs_init_script not in zrb_init_path_list:
+                LOGGER.info(f"Loading {abs_init_script}")
                 load_file(abs_init_script, -1)
         # load zrb init
         for zrb_init_path in zrb_init_path_list:
+            LOGGER.info(f"Loading {zrb_init_path}")
             load_file(zrb_init_path)
         # run the CLI
         cli.run(sys.argv[1:])
@@ -45,6 +70,7 @@ def _get_zrb_init_path_list() -> list[str]:
     zrb_init_path_list = []
     for current_path in dir_path_list[::-1]:
         zrb_init_path = os.path.join(current_path, "zrb_init.py")
+        LOGGER.info(f"Finding {zrb_init_path}")
         if os.path.isfile(zrb_init_path):
             zrb_init_path_list.append(zrb_init_path)
     return zrb_init_path_list

zrb/builtin/llm/history.py

@@ -0,0 +1,73 @@
+import json
+import os
+from typing import Any
+
+from zrb.config import LLM_HISTORY_DIR
+from zrb.context.any_shared_context import AnySharedContext
+from zrb.task.llm.history import ConversationHistoryData
+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 ConversationHistoryData.
+    """
+    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(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(
+        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: ConversationHistoryData
+):
+    """Writes the conversation history data (including context) to a session file."""
+    os.makedirs(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(
+        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(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

@@ -0,0 +1,27 @@
+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_chat.py

@@ -1,8 +1,6 @@
-import json
-import os
-from typing import Any
-
 from zrb.builtin.group import llm_group
+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.file import (
@@ -22,67 +20,12 @@ from zrb.config import (
     LLM_ALLOW_ACCESS_INTERNET,
     LLM_ALLOW_ACCESS_LOCAL_FILE,
     LLM_ALLOW_ACCESS_SHELL,
-    LLM_HISTORY_DIR,
     SERP_API_KEY,
 )
-from zrb.context.any_shared_context import AnySharedContext
 from zrb.input.bool_input import BoolInput
 from zrb.input.str_input import StrInput
 from zrb.input.text_input import TextInput
 from zrb.task.llm_task import LLMTask
-from zrb.util.file import read_file, write_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>",
-            ]
-        )
-
-
-def _read_chat_conversation(ctx: AnySharedContext) -> list[dict[str, Any]]:
-    if ctx.input.start_new:
-        return []
-    previous_session_name = ctx.input.previous_session
-    if previous_session_name == "" or previous_session_name is None:
-        last_session_file_path = os.path.join(LLM_HISTORY_DIR, "last-session")
-        if os.path.isfile(last_session_file_path):
-            previous_session_name = read_file(last_session_file_path).strip()
-    conversation_file_path = os.path.join(
-        LLM_HISTORY_DIR, f"{previous_session_name}.json"
-    )
-    if not os.path.isfile(conversation_file_path):
-        return []
-    return json.loads(read_file(conversation_file_path))
-
-
-def _write_chat_conversation(
-    ctx: AnySharedContext, conversations: list[dict[str, Any]]
-):
-    os.makedirs(LLM_HISTORY_DIR, exist_ok=True)
-    current_session_name = ctx.session.name
-    conversation_file_path = os.path.join(
-        LLM_HISTORY_DIR, f"{current_session_name}.json"
-    )
-    write_file(conversation_file_path, json.dumps(conversations, indent=2))
-    last_session_file_path = os.path.join(LLM_HISTORY_DIR, "last-session")
-    write_file(last_session_file_path, current_session_name)
-
 
 llm_chat: LLMTask = llm_group.add_task(
     LLMTask(
@@ -148,8 +91,8 @@ llm_chat: LLMTask = llm_group.add_task(
         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,
+        conversation_history_reader=read_chat_conversation,
+        conversation_history_writer=write_chat_conversation,
         description="Chat with LLM",
         system_prompt=lambda ctx: (
             None if ctx.input.system_prompt.strip() == "" else ctx.input.system_prompt

zrb/builtin/llm/tool/api.py

@@ -1,14 +1,22 @@
 import json
-from typing import Annotated, Literal
+from typing import Literal
 
 
-def get_current_location() -> (
-    Annotated[str, "JSON string representing latitude and longitude"]
-):  # noqa
-    """Get the user's current location. This function take no argument."""
+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
 
-    return json.dumps(requests.get("http://ip-api.com/json?fields=lat,lon").json())
+    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(
@@ -16,16 +24,30 @@ def get_current_weather(
     longitude: float,
     temperature_unit: Literal["celsius", "fahrenheit"],
 ) -> str:
-    """Get the current weather in a given location."""
+    """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
 
-    resp = requests.get(
-        "https://api.open-meteo.com/v1/forecast",
-        params={
-            "latitude": latitude,
-            "longitude": longitude,
-            "temperature_unit": temperature_unit,
-            "current_weather": True,
-        },
-    )
-    return json.dumps(resp.json())
+    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

zrb/builtin/llm/tool/cli.py

@@ -2,8 +2,22 @@ import subprocess
 
 
 def run_shell_command(command: str) -> str:
-    """Running an actual shell command on user's computer."""
-    output = subprocess.check_output(
-        command, shell=True, stderr=subprocess.STDOUT, text=True
-    )
-    return output
+    """Execute a shell command and return its combined stdout and stderr.
+    Args:
+        command (str): Shell command to execute on the user's system.
+    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.
+    """
+    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