zrb 1.17.2__tar.gz → 1.17.3__tar.gz

{zrb-1.17.2 → zrb-1.17.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zrb
-Version: 1.17.2
+Version: 1.17.3
 Summary: Your Automation Powerhouse
 License: AGPL-3.0-or-later
 Keywords: Automation,Task Runner,Code Generator,Monorepo,Low Code

{zrb-1.17.2 → zrb-1.17.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.17.2"
+version = "1.17.3"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"

{zrb-1.17.2 → zrb-1.17.3}/src/zrb/builtin/llm/tool/cli.py

@@ -4,7 +4,8 @@ from typing import Any
 
 def run_shell_command(command: str) -> dict[str, Any]:
     """
-    Executes a shell command on the user's local machine and returns the output.
+    Executes a non interactive shell command on the user's local machine and returns
+    the output.
 
     This tool is powerful and should be used for tasks that require interacting
     with the command line, such as running scripts, managing system processes,
@@ -15,6 +16,9 @@ def run_shell_command(command: str) -> dict[str, Any]:
     modify files or system state (e.g., `git`, `npm`, `pip`, `docker`), you
     MUST explain what the command does and ask the user for confirmation.
 
+    **Note:** Make sure to run any server or long running process, as background process.
+    (e.g., python -m http.server &)
+
     Args:
         command (str): The exact shell command to execute.
 

{zrb-1.17.2 → zrb-1.17.3}/src/zrb/builtin/llm/tool/file.py

@@ -500,10 +500,13 @@ async def analyze_file(
         tools=[read_from_file, search_files],
     )
     payload = json.dumps(
-        {"instruction": query, "file_path": abs_path, "file_content": file_content}
+        {
+            "instruction": query,
+            "file_path": abs_path,
+            "file_content": llm_rate_limitter.clip_prompt(file_content, token_limit),
+        }
     )
-    clipped_payload = llm_rate_limitter.clip_prompt(payload, token_limit)
-    return await _analyze_file(ctx, clipped_payload)
+    return await _analyze_file(ctx, payload)
 
 
 def read_many_files(paths: list[str]) -> dict[str, str]:
@@ -593,11 +596,9 @@ def write_many_files(files: list[FileToWrite]) -> dict[str, Any]:
     errors = {}
     # 4. Access the data using dictionary key-lookup syntax.
     for file in files:
+        path = file["path"]
+        content = file["content"]
         try:
-            # Use file['path'] and file['content'] instead of file.path
-            path = file["path"]
-            content = file["content"]
-
             abs_path = os.path.abspath(os.path.expanduser(path))
             directory = os.path.dirname(abs_path)
             if directory and not os.path.exists(directory):

{zrb-1.17.2 → zrb-1.17.3}/src/zrb/config/config.py

@@ -313,6 +313,11 @@ class Config:
         """Maximum number of tokens allowed per individual LLM request."""
         return int(self._getenv("LLM_MAX_TOKENS_PER_REQUEST", "100000"))
 
+    @property
+    def LLM_MAX_TOKENS_PER_TOOL_CALL_RESULT(self) -> int:
+        """Maximum number of tokens allowed per tool call result."""
+        return int(self._getenv("LLM_MAX_TOKENS_PER_TOOL_CALL_RESULT", "75000"))
+
     @property
     def LLM_THROTTLE_SLEEP(self) -> float:
         """Number of seconds to sleep when throttling is required."""

{zrb-1.17.2 → zrb-1.17.3}/src/zrb/config/llm_rate_limitter.py

@@ -1,7 +1,8 @@
 import asyncio
+import json
 import time
 from collections import deque
-from typing import Callable
+from typing import Any, Callable
 
 from zrb.config.config import CFG
 
@@ -17,16 +18,18 @@ class LLMRateLimiter:
         max_requests_per_minute: int | None = None,
         max_tokens_per_minute: int | None = None,
         max_tokens_per_request: int | None = None,
+        max_tokens_per_tool_call_result: int | None = None,
         throttle_sleep: float | None = None,
         use_tiktoken: bool | None = None,
-        tiktoken_encodeing_name: str | None = None,
+        tiktoken_encoding_name: str | None = None,
     ):
         self._max_requests_per_minute = max_requests_per_minute
         self._max_tokens_per_minute = max_tokens_per_minute
         self._max_tokens_per_request = max_tokens_per_request
+        self._max_tokens_per_tool_call_result = max_tokens_per_tool_call_result
         self._throttle_sleep = throttle_sleep
         self._use_tiktoken = use_tiktoken
-        self._tiktoken_encoding_name = tiktoken_encodeing_name
+        self._tiktoken_encoding_name = tiktoken_encoding_name
         self.request_times = deque()
         self.token_times = deque()
 
@@ -48,6 +51,12 @@ class LLMRateLimiter:
             return self._max_tokens_per_request
         return CFG.LLM_MAX_TOKENS_PER_REQUEST
 
+    @property
+    def max_tokens_per_tool_call_result(self) -> int:
+        if self._max_tokens_per_tool_call_result is not None:
+            return self._max_tokens_per_tool_call_result
+        return CFG.LLM_MAX_TOKENS_PER_TOOL_CALL_RESULT
+
     @property
     def throttle_sleep(self) -> float:
         if self._throttle_sleep is not None:
@@ -75,48 +84,56 @@ class LLMRateLimiter:
     def set_max_tokens_per_request(self, value: int):
         self._max_tokens_per_request = value
 
+    def set_max_tokens_per_tool_call_result(self, value: int):
+        self._max_tokens_per_tool_call_result = value
+
     def set_throttle_sleep(self, value: float):
         self._throttle_sleep = value
 
-    def count_token(self, prompt: str) -> int:
+    def count_token(self, prompt: Any) -> int:
+        str_prompt = self._prompt_to_str(prompt)
         if not self.use_tiktoken:
-            return self._fallback_count_token(prompt)
+            return self._fallback_count_token(str_prompt)
         try:
             import tiktoken
 
             enc = tiktoken.get_encoding(self.tiktoken_encoding_name)
-            return len(enc.encode(prompt))
+            return len(enc.encode(str_prompt))
         except Exception:
-            return self._fallback_count_token(prompt)
+            return self._fallback_count_token(str_prompt)
 
-    def _fallback_count_token(self, prompt: str) -> int:
-        return len(prompt) // 4
+    def _fallback_count_token(self, str_prompt: str) -> int:
+        return len(str_prompt) // 4
 
-    def clip_prompt(self, prompt: str, limit: int) -> str:
+    def clip_prompt(self, prompt: Any, limit: int) -> str:
+        str_prompt = self._prompt_to_str(prompt)
         if not self.use_tiktoken:
-            return self._fallback_clip_prompt(prompt, limit)
+            return self._fallback_clip_prompt(str_prompt, limit)
         try:
             import tiktoken
 
-            enc = tiktoken.get_encoding("cl100k_base")
-            tokens = enc.encode(prompt)
+            enc = tiktoken.get_encoding(self.tiktoken_encoding_name)
+            tokens = enc.encode(str_prompt)
             if len(tokens) <= limit:
-                return prompt
+                return str_prompt
             truncated = tokens[: limit - 3]
             clipped_text = enc.decode(truncated)
             return clipped_text + "..."
         except Exception:
-            return self._fallback_clip_prompt(prompt, limit)
+            return self._fallback_clip_prompt(str_prompt, limit)
 
-    def _fallback_clip_prompt(self, prompt: str, limit: int) -> str:
+    def _fallback_clip_prompt(self, str_prompt: str, limit: int) -> str:
         char_limit = limit * 4 if limit * 4 <= 10 else limit * 4 - 10
-        return prompt[:char_limit] + "..."
+        return str_prompt[:char_limit] + "..."
 
     async def throttle(
-        self, prompt: str, throttle_notif_callback: Callable | None = None
+        self,
+        prompt: Any,
+        throttle_notif_callback: Callable | None = None,
     ):
         now = time.time()
-        tokens = self.count_token(prompt)
+        str_prompt = self._prompt_to_str(prompt)
+        tokens = self.count_token(str_prompt)
         # Clean up old entries
         while self.request_times and now - self.request_times[0] > 60:
             self.request_times.popleft()
@@ -144,5 +161,11 @@ class LLMRateLimiter:
         self.request_times.append(now)
         self.token_times.append((now, tokens))
 
+    def _prompt_to_str(self, prompt: Any) -> str:
+        try:
+            return json.dumps(prompt)
+        except Exception:
+            return f"{prompt}"
+
 
 llm_rate_limitter = LLMRateLimiter()

{zrb-1.17.2 → zrb-1.17.3}/src/zrb/task/llm/prompt.py

@@ -2,11 +2,12 @@ import os
 import platform
 import re
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Callable
+from typing import TYPE_CHECKING, Any, Callable
 
 from zrb.attr.type import StrAttr, StrListAttr
-from zrb.config.llm_config import llm_config as llm_config
+from zrb.config.llm_config import llm_config
 from zrb.config.llm_context.config import llm_context_config
+from zrb.config.llm_rate_limitter import llm_rate_limitter
 from zrb.context.any_context import AnyContext
 from zrb.context.any_shared_context import AnySharedContext
 from zrb.task.llm.conversation_history_model import ConversationHistory
@@ -173,7 +174,43 @@ def get_system_and_user_prompt(
     render_modes: bool = False,
     conversation_history: ConversationHistory | None = None,
 ) -> tuple[str, str]:
-    """Combines persona, base system prompt, and special instructions."""
+    if conversation_history is None:
+        conversation_history = ConversationHistory()
+    modified_user_message, user_message_context = _extract_user_prompt_components(
+        user_message
+    )
+    new_system_prompt = _construct_system_prompt(
+        ctx=ctx,
+        persona_attr=persona_attr,
+        render_persona=render_persona,
+        system_prompt_attr=system_prompt_attr,
+        render_system_prompt=render_system_prompt,
+        special_instruction_prompt_attr=special_instruction_prompt_attr,
+        render_special_instruction_prompt=render_special_instruction_prompt,
+        modes_attr=modes_attr,
+        render_modes=render_modes,
+        conversation_history=conversation_history,
+    )
+    user_message_context["Token Counter"] = _get_token_usage_info(
+        new_system_prompt, user_message, conversation_history
+    )
+    return new_system_prompt, _contruct_user_message_prompt(
+        modified_user_message, user_message_context
+    )
+
+
+def _construct_system_prompt(
+    ctx: AnyContext,
+    persona_attr: StrAttr | None = None,
+    render_persona: bool = False,
+    system_prompt_attr: StrAttr | None = None,
+    render_system_prompt: bool = False,
+    special_instruction_prompt_attr: StrAttr | None = None,
+    render_special_instruction_prompt: bool = False,
+    modes_attr: StrListAttr | None = None,
+    render_modes: bool = False,
+    conversation_history: ConversationHistory | None = None,
+) -> str:
     persona = get_persona(ctx, persona_attr, render_persona)
     base_system_prompt = get_base_system_prompt(
         ctx, system_prompt_attr, render_system_prompt
@@ -185,10 +222,9 @@ def get_system_and_user_prompt(
     project_context_prompt = get_project_context_prompt()
     if conversation_history is None:
         conversation_history = ConversationHistory()
-    conversation_environment, new_user_message = extract_conversation_environment(
-        user_message
-    )
-    new_system_prompt = "\n".join(
+    current_directory = os.getcwd()
+    directory_tree = _generate_directory_tree(current_directory, max_depth=2)
+    return "\n".join(
         [
             make_prompt_section("Persona", persona),
             make_prompt_section("System Prompt", base_system_prompt),
@@ -229,10 +265,42 @@ def get_system_and_user_prompt(
                 ),
             ),
             make_prompt_section("Project Context", project_context_prompt),
-            make_prompt_section("Conversation Environment", conversation_environment),
+            make_prompt_section(
+                "Conversation Environment",
+                "\n".join(
+                    [
+                        make_prompt_section("Current OS", platform.system()),
+                        make_prompt_section("OS Version", platform.version()),
+                        make_prompt_section(
+                            "Python Version", platform.python_version()
+                        ),
+                        make_prompt_section(
+                            "Directory Tree (depth=2)", directory_tree, as_code=True
+                        ),
+                    ]
+                ),
+            ),
+        ]
+    )
+
+
+def _contruct_user_message_prompt(
+    modified_user_message: str, user_message_context: dict[str, str]
+) -> str:
+    return "\n".join(
+        [
+            make_prompt_section("User Message", modified_user_message),
+            make_prompt_section(
+                "Context",
+                "\n".join(
+                    [
+                        make_prompt_section(key, val)
+                        for key, val in user_message_context.items()
+                    ]
+                ),
+            ),
         ]
     )
-    return new_system_prompt, new_user_message
 
 
 def _generate_directory_tree(
@@ -258,13 +326,13 @@ def _generate_directory_tree(
 
         for i, entry in enumerate(entries):
             if i >= max_children:
-                tree_lines.append(f"{prefix}└───... (more)")
+                tree_lines.append(f"{prefix}└─... (more)")
                 break
             is_last = i == len(entries) - 1
-            connector = "└───" if is_last else "├───"
+            connector = "└─" if is_last else "├─"
             tree_lines.append(f"{prefix}{connector}{entry.name}")
             if entry.is_dir():
-                new_prefix = prefix + ("    " if is_last else "│   ")
+                new_prefix = prefix + ("  " if is_last else "│ ")
                 recurse(entry.path, depth + 1, new_prefix)
 
     tree_lines.append(os.path.basename(dir_path))
@@ -272,7 +340,24 @@ def _generate_directory_tree(
     return "\n".join(tree_lines)
 
 
-def extract_conversation_environment(user_message: str) -> tuple[str, str]:
+def _get_token_usage_info(
+    system_prompt: str, user_message: str, conversation_history: ConversationHistory
+) -> str:
+    system_prompt_token_count = llm_rate_limitter.count_token(system_prompt)
+    user_message_token_count = llm_rate_limitter.count_token(user_message)
+    conversation_history_token_count = llm_rate_limitter.count_token(
+        conversation_history.history
+    )
+    total_token_count = (
+        system_prompt_token_count
+        + user_message_token_count
+        + conversation_history_token_count
+    )
+    max_token = llm_rate_limitter.max_tokens_per_request
+    return f"{total_token_count} tokens of {max_token} tokens"
+
+
+def _extract_user_prompt_components(user_message: str) -> tuple[str, dict[str, Any]]:
     modified_user_message = user_message
     # Match “@” + any non-space/comma sequence that contains at least one “/”
     pattern = r"(?<!\w)@(?=[^,\s]*\/)([^,\?\!\s]+)"
@@ -297,41 +382,18 @@ def extract_conversation_environment(user_message: str) -> tuple[str, str]:
             apendixes.append(
                 make_prompt_section(
                     f"{placeholder} ({ref_type} path: `{resource_path}`)",
-                    content,
+                    "\n".join(content) if isinstance(content, list) else content,
                     as_code=True,
                 )
             )
     current_directory = os.getcwd()
-    directory_tree = _generate_directory_tree(current_directory, max_depth=2)
-    conversation_environment = "\n".join(
-        [
-            make_prompt_section("Current OS", platform.system()),
-            make_prompt_section("OS Version", platform.version()),
-            make_prompt_section("Python Version", platform.python_version()),
-            make_prompt_section(
-                "Directory Tree (depth=2)", directory_tree, as_code=True
-            ),
-        ]
-    )
     iso_date = datetime.now(timezone.utc).astimezone().isoformat()
-    modified_user_message = "\n".join(
-        [
-            make_prompt_section("User Message", modified_user_message),
-            make_prompt_section(
-                "Context",
-                "\n".join(
-                    [
-                        make_prompt_section(
-                            "Current working directory", current_directory
-                        ),
-                        make_prompt_section("Current time", iso_date),
-                        make_prompt_section("Apendixes", "\n".join(apendixes)),
-                    ]
-                ),
-            ),
-        ]
-    )
-    return conversation_environment, modified_user_message
+    user_message_context = {
+        "Current Working Directory": current_directory,
+        "Current Time": iso_date,
+        "Apendixes": "\n".join(apendixes),
+    }
+    return modified_user_message, user_message_context
 
 
 def get_user_message(

{zrb-1.17.2 → zrb-1.17.3}/src/zrb/task/llm/tool_wrapper.py

@@ -7,6 +7,7 @@ from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
 from zrb.config.config import CFG
+from zrb.config.llm_rate_limitter import llm_rate_limitter
 from zrb.context.any_context import AnyContext
 from zrb.task.llm.error import ToolExecutionError
 from zrb.util.callable import get_callable_name
@@ -124,6 +125,7 @@ def _create_wrapper(
                             f"Tool execution cancelled. User disapproving: {reason}"
                         )
             result = await run_async(func(*args, **kwargs))
+            _check_tool_call_result_limit(result)
             if has_ever_edited:
                 return {
                     "tool_call_result": result,
@@ -145,6 +147,14 @@ def _create_wrapper(
     return wrapper
 
 
+def _check_tool_call_result_limit(result: Any):
+    if (
+        llm_rate_limitter.count_token(result)
+        > llm_rate_limitter.max_tokens_per_tool_call_result
+    ):
+        raise ValueError("Result value is too large, please adjust the parameter")
+
+
 async def _handle_user_response(
     ctx: AnyContext,
     func: Callable,