zrb 1.8.3__tar.gz → 1.8.5__tar.gz

{zrb-1.8.3 → zrb-1.8.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.8.3
+Version: 1.8.5
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -30,7 +30,7 @@ Requires-Dist: pydantic-ai (>=0.2.7,<0.3.0)
 Requires-Dist: pyjwt (>=2.10.1,<3.0.0)
 Requires-Dist: python-dotenv (>=1.1.0,<2.0.0)
 Requires-Dist: python-jose[cryptography] (>=3.4.0,<4.0.0)
-Requires-Dist: requests (>=2.32.3,<3.0.0)
+Requires-Dist: requests (>=2.32.4,<3.0.0)
 Requires-Dist: ulid-py (>=1.1.0,<2.0.0)
 Project-URL: Documentation, https://github.com/state-alchemists/zrb
 Project-URL: Repository, https://github.com/state-alchemists/zrb

{zrb-1.8.3 → zrb-1.8.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.8.3"
+version = "1.8.5"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"
@@ -41,7 +41,7 @@ python = ">=3.10.0,<4.0.0"
 python-dotenv = "^1.1.0"
 black = "~25.1.0"
 isort = "~6.0.1"
-requests = "^2.32.3"
+requests = "^2.32.4"
 libcst = "^1.7.0"
 chromadb = {version = "^0.6.3", optional = true}
 pdfplumber = {version = "^0.11.6", optional = true}

{zrb-1.8.3 → zrb-1.8.5}/src/zrb/__init__.py

@@ -34,6 +34,7 @@ from zrb.input.password_input import PasswordInput
 from zrb.input.str_input import StrInput
 from zrb.input.text_input import TextInput
 from zrb.llm_config import llm_config
+from zrb.llm_rate_limitter import llm_rate_limitter
 from zrb.runner.cli import cli
 from zrb.runner.web_auth_config import web_auth_config
 from zrb.runner.web_schema.user import User
@@ -105,6 +106,7 @@ assert Scaffolder
 assert Scheduler
 assert cli
 assert llm_config
+assert llm_rate_limitter
 assert Xcom
 assert web_auth_config
 assert User

{zrb-1.8.3 → zrb-1.8.5}/src/zrb/builtin/llm/chat_session.py

@@ -33,10 +33,10 @@ async def read_user_prompt(ctx: AnyContext) -> str:
         previous_session_name=ctx.input.previous_session,
         start_new=ctx.input.start_new,
     )
+    if result is not None:
+        final_result = result
     if ctx.env.get("_ZRB_WEB_ENV", "0") != "0":
-        # Don't run in web environment
-        if result is not None:
-            final_result = result
+        # On web environment this won't be interactive
         return final_result
     multiline_mode = False
     user_inputs = []

zrb-1.8.5/src/zrb/builtin/llm/tool/code.py

@@ -0,0 +1,233 @@
+import os
+
+from zrb.builtin.llm.tool.file import DEFAULT_EXCLUDED_PATTERNS, is_excluded
+from zrb.builtin.llm.tool.sub_agent import create_sub_agent_tool
+from zrb.context.any_context import AnyContext
+
+_EXTRACT_INFO_SYSTEM_PROMPT = """
+You are an extraction info agent.
+Your goal is to help to extract relevant information to help the main LLM Agent.
+You write your output is in markdown format containing path and relevant information.
+Extract only information that relevant to main LLM Agent's goal.
+
+Extracted Information format (Use this as reference, extract relevant information only):
+# <file-name>
+## imports
+- <imported-package>
+- ...
+## variables
+- <variable-type> <variable-name>: <the-purpose-of-the-variable>
+- ...
+## functions
+- <function-name>:
+  - parameters: <parameters>
+  - logic/description: <what-the-function-do-and-how-it-works>
+...
+# <other-file-name>
+...
+""".strip()
+
+
+_SUMMARIZE_INFO_SYSTEM_PROMPT = """
+You are an information summarization agent.
+Your goal is to summarize information to help the main LLM Agent.
+The summarization result should contains all necessary details
+to help main LLM Agent achieve the goal.
+"""
+
+_DEFAULT_EXTENSIONS = [
+    "py",
+    "go",
+    "java",
+    "ts",
+    "js",
+    "rs",
+    "rb",
+    "php",
+    "sh",
+    "bash",
+    "c",
+    "cpp",
+    "h",
+    "hpp",
+    "cs",
+    "swift",
+    "kt",
+    "scala",
+    "m",
+    "pl",
+    "lua",
+    "sql",
+    "html",
+    "css",
+    "scss",
+    "less",
+    "json",
+    "yaml",
+    "yml",
+    "toml",
+    "ini",
+    "xml",
+    "md",
+    "rst",
+    "txt",
+]
+
+
+async def analyze_repo(
+    ctx: AnyContext,
+    path: str,
+    goal: str,
+    extensions: list[str] = _DEFAULT_EXTENSIONS,
+    exclude_patterns: list[str] = DEFAULT_EXCLUDED_PATTERNS,
+    extraction_char_limit: int = 150000,
+    summarization_char_limit: int = 150000,
+) -> str:
+    """
+    Extract and summarize information from a directory that probably
+    contains a large resources.
+    You should state the goal specifically so that the tool can return relevant informations.
+    Use this tool for:
+    - summarization
+    - outline/structure extraction
+    - code review
+    - create diagram as code
+    - other tasks
+    Args:
+        path (str): File path to be analyze. Pass exactly as provided, including '~'.
+        goal(str): Goal of extracting information (for example creating C4 diagram)
+        extensions(Optional[list[str]]): List of extension to be included
+            while reading resources. Defaults to common programming languages and config files.
+        exclude_patterns(Optional[list[str]]): List of patterns to exclude from analysis.
+            Common patterns like '.venv', 'node_modules' should be excluded by default.
+        extraction_char_limit(Optional[int]): Max resource content char length
+            the extraction assistant able to handle. Defaults to 150000
+        summarization_char_limit(Optional[int]): Max resource content char length
+            the summarization assistant able to handle. Defaults to 150000
+    Returns:
+        str: The analysis result
+    Raises:
+        Exception: If an error occurs.
+    """
+    abs_path = os.path.abspath(os.path.expanduser(path))
+    file_metadatas = _get_file_metadatas(abs_path, extensions, exclude_patterns)
+    ctx.print("Extraction")
+    extracted_infos = await _extract_info(
+        ctx,
+        file_metadatas=file_metadatas,
+        goal=goal,
+        char_limit=extraction_char_limit,
+    )
+    ctx.print("Summarization")
+    summarized_infos = await _summarize_info(
+        ctx,
+        extracted_infos=extracted_infos,
+        goal=goal,
+        char_limit=summarization_char_limit,
+    )
+    while len(summarized_infos) > 1:
+        ctx.print("Summarization")
+        summarized_infos = await _summarize_info(
+            ctx,
+            extracted_infos=summarized_infos,
+            goal=goal,
+            char_limit=summarization_char_limit,
+        )
+    return summarized_infos[0]
+
+
+def _get_file_metadatas(
+    dir_path: str,
+    extensions: list[str],
+    exclude_patterns: list[str],
+) -> list[dict[str, str]]:
+    metadata_list = []
+    for root, _, files in os.walk(dir_path):
+        files.sort()
+        for file in files:
+            if not any(file.endswith(f".{ext}") for ext in extensions):
+                continue
+            file_path = os.path.join(root, file)
+            if is_excluded(file_path, exclude_patterns):
+                continue
+            try:
+                with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+                    rel_path = os.path.relpath(file_path, dir_path)
+                    metadata_list.append({"path": rel_path, "content": f.read()})
+            except Exception as e:
+                print(f"Error reading file {file_path}: {e}")
+    metadata_list.sort(key=lambda m: m["path"])
+    return metadata_list
+
+
+async def _extract_info(
+    ctx: AnyContext,
+    file_metadatas: list[dict[str, str]],
+    goal: str,
+    char_limit: int,
+) -> list[str]:
+    extract = create_sub_agent_tool(
+        tool_name="extract",
+        tool_description="extract",
+        system_prompt=_EXTRACT_INFO_SYSTEM_PROMPT,
+    )
+    extracted_infos = []
+    content_buffer = ""
+    for metadata in file_metadatas:
+        path = metadata.get("path", "")
+        content = metadata.get("content", "")
+        metadata_str = f"Path: {path}\nContent: {content}"
+        if len(content_buffer) + len(metadata_str) > char_limit:
+            if content_buffer:
+                prompt = _create_extract_info_prompt(goal, content_buffer)
+                extracted_info = await extract(ctx, prompt)
+                extracted_infos.append(extracted_info)
+            content_buffer = metadata_str
+        else:
+            content_buffer += metadata_str + "\n"
+
+    # Process any remaining content in the buffer
+    if content_buffer:
+        prompt = _create_extract_info_prompt(goal, content_buffer)
+        extracted_info = await extract(ctx, prompt)
+        extracted_infos.append(extracted_info)
+    return extracted_infos
+
+
+def _create_extract_info_prompt(goal: str, content_buffer: str) -> str:
+    return f"# Main LLM Agent Goal\n{goal}\n# Files\n{content_buffer}"
+
+
+async def _summarize_info(
+    ctx: AnyContext,
+    extracted_infos: list[str],
+    goal: str,
+    char_limit: int,
+) -> list[str]:
+    summarize = create_sub_agent_tool(
+        tool_name="extract",
+        tool_description="extract",
+        system_prompt=_SUMMARIZE_INFO_SYSTEM_PROMPT,
+    )
+    summarized_infos = []
+    content_buffer = ""
+    for extracted_info in extracted_infos:
+        if len(content_buffer) + len(extracted_info) > char_limit:
+            if content_buffer:
+                prompt = _create_summarize_info_prompt(goal, content_buffer)
+                summarized_info = await summarize(ctx, prompt)
+                summarized_infos.append(summarized_info)
+            content_buffer = extracted_info
+        else:
+            content_buffer += extracted_info + "\n"
+
+    # Process any remaining content in the buffer
+    if content_buffer:
+        prompt = _create_summarize_info_prompt(goal, content_buffer)
+        summarized_info = await summarize(ctx, prompt)
+        summarized_infos.append(summarized_info)
+    return summarized_infos
+
+
+def _create_summarize_info_prompt(goal: str, content_buffer: str) -> str:
+    return f"# Main LLM Agent Goal\n{goal}\n# Extracted Info\n{content_buffer}"

{zrb-1.8.3 → zrb-1.8.5}/src/zrb/builtin/llm/tool/file.py

@@ -71,6 +71,9 @@ DEFAULT_EXCLUDED_PATTERNS = [
     "*.so",
     "*.dylib",
     "*.dll",
+    # Minified files
+    "*.min.css",
+    "*.min.js",
 ]
 
 
@@ -112,17 +115,17 @@ def list_files(
                     d
                     for d in dirs
                     if (include_hidden or not _is_hidden(d))
-                    and not _is_excluded(d, patterns_to_exclude)
+                    and not is_excluded(d, patterns_to_exclude)
                 ]
                 # Process files
                 for filename in files:
-                    if (
-                        include_hidden or not _is_hidden(filename)
-                    ) and not _is_excluded(filename, patterns_to_exclude):
+                    if (include_hidden or not _is_hidden(filename)) and not is_excluded(
+                        filename, patterns_to_exclude
+                    ):
                         full_path = os.path.join(root, filename)
                         # Check rel path for patterns like '**/node_modules/*'
                         rel_full_path = os.path.relpath(full_path, abs_path)
-                        is_rel_path_excluded = _is_excluded(
+                        is_rel_path_excluded = is_excluded(
                             rel_full_path, patterns_to_exclude
                         )
                         if not is_rel_path_excluded:
@@ -132,7 +135,7 @@ def list_files(
             for item in os.listdir(abs_path):
                 full_path = os.path.join(abs_path, item)
                 # Include both files and directories if not recursive
-                if (include_hidden or not _is_hidden(item)) and not _is_excluded(
+                if (include_hidden or not _is_hidden(item)) and not is_excluded(
                     item, patterns_to_exclude
                 ):
                     all_files.append(full_path)
@@ -169,7 +172,7 @@ def _is_hidden(path: str) -> bool:
     return basename.startswith(".")
 
 
-def _is_excluded(name: str, patterns: list[str]) -> bool:
+def is_excluded(name: str, patterns: list[str]) -> bool:
     """Check if a name/path matches any exclusion patterns."""
     for pattern in patterns:
         if fnmatch.fnmatch(name, pattern):
@@ -461,14 +464,14 @@ async def analyze_file(ctx: AnyContext, path: str, query: str) -> str:
     _analyze_file = create_sub_agent_tool(
         tool_name="analyze_file",
         tool_description="analyze file with LLM capability",
-        sub_agent_system_prompt="\n".join(
+        system_prompt="\n".join(
             [
                 "You are a file analyzer assistant",
                 "Your goal is to help the main asisstant by reading file",
                 "and perform necessary indepth analysis required by the main assistant",
             ]
         ),
-        sub_agent_tools=[read_from_file, search_files],
+        tools=[read_from_file, search_files],
     )
     return await _analyze_file(
         ctx,

{zrb-1.8.3 → zrb-1.8.5}/src/zrb/builtin/llm/tool/sub_agent.py

@@ -28,11 +28,11 @@ else:
 def create_sub_agent_tool(
     tool_name: str,
     tool_description: str,
-    sub_agent_system_prompt: str | None = None,  # Make optional
-    sub_agent_model: str | Model | None = None,
-    sub_agent_model_settings: ModelSettings | None = None,
-    sub_agent_tools: list[ToolOrCallable] = [],
-    sub_agent_mcp_servers: list[MCPServer] = [],
+    system_prompt: str | None = None,
+    model: str | Model | None = None,
+    model_settings: ModelSettings | None = None,
+    tools: list[ToolOrCallable] = [],
+    mcp_servers: list[MCPServer] = [],
 ) -> Callable[[AnyContext, str], str]:
     """
     Create an LLM "sub-agent" tool function for use by a main LLM agent.
@@ -63,7 +63,7 @@ def create_sub_agent_tool(
         # Resolve parameters, falling back to llm_config defaults if None
         resolved_model = get_model(
             ctx=ctx,
-            model_attr=sub_agent_model,
+            model_attr=model,
             render_model=True,  # Assuming we always want to render model string attributes
             model_base_url_attr=None,
             # Sub-agent tool doesn't have separate base_url/api_key params
@@ -73,10 +73,10 @@ def create_sub_agent_tool(
         )
         resolved_model_settings = get_model_settings(
             ctx=ctx,
-            model_settings_attr=sub_agent_model_settings,
+            model_settings_attr=model_settings,
         )
 
-        if sub_agent_system_prompt is None:
+        if system_prompt is None:
             resolved_system_prompt = get_combined_system_prompt(
                 ctx=ctx,
                 persona_attr=None,
@@ -87,20 +87,16 @@ def create_sub_agent_tool(
                 render_special_instruction_prompt=False,
             )
         else:
-            resolved_system_prompt = sub_agent_system_prompt
-
+            resolved_system_prompt = system_prompt
         # Create the sub-agent instance
         sub_agent_agent = create_agent_instance(
             ctx=ctx,
             model=resolved_model,
             system_prompt=resolved_system_prompt,
             model_settings=resolved_model_settings,
-            tools_attr=sub_agent_tools,  # Pass tools from factory closure
-            additional_tools=[],  # No additional tools added after factory creation
-            mcp_servers_attr=sub_agent_mcp_servers,  # Pass servers from factory closure
-            additional_mcp_servers=[],  # No additional servers added after factory creation
+            tools=tools,
+            mcp_servers=mcp_servers,
         )
-
         # Run the sub-agent iteration
         # Start with an empty history for the sub-agent
         sub_agent_run = await run_agent_iteration(

{zrb-1.8.3 → zrb-1.8.5}/src/zrb/callback/callback.py

@@ -1,3 +1,4 @@
+import traceback
 from typing import Any
 
 from zrb.attr.type import StrDictAttr
@@ -5,6 +6,7 @@ from zrb.callback.any_callback import AnyCallback
 from zrb.session.any_session import AnySession
 from zrb.task.any_task import AnyTask
 from zrb.util.attr import get_str_dict_attr
+from zrb.util.cli.style import stylize_faint
 from zrb.util.string.conversion import to_snake_case
 from zrb.xcom.xcom import Xcom
 
@@ -67,6 +69,8 @@ class Callback(AnyCallback):
             self._maybe_publish_result_to_parent_session(parent_session, result)
             return result
         except BaseException as e:
+            ctx = session.get_ctx(self._task)
+            ctx.print(traceback.format_exc())
             self._maybe_publish_error_to_parent_session(parent_session, e)
 
     def _maybe_publish_session_name_to_parent_session(

{zrb-1.8.3 → zrb-1.8.5}/src/zrb/config.py

@@ -237,6 +237,26 @@ class Config:
     def LLM_SUMMARIZATION_PROMPT(self) -> str | None:
         return os.getenv("ZRB_LLM_SUMMARIZATION_PROMPT", None)
 
+    @property
+    def LLM_MAX_REQUESTS_PER_MINUTE(self) -> int:
+        """Maximum number of LLM requests allowed per minute."""
+        return int(os.getenv("LLM_MAX_REQUESTS_PER_MINUTE", "60"))
+
+    @property
+    def LLM_MAX_TOKENS_PER_MINUTE(self) -> int:
+        """Maximum number of LLM tokens allowed per minute."""
+        return int(os.getenv("LLM_MAX_TOKENS_PER_MINUTE", "120000"))
+
+    @property
+    def LLM_MAX_TOKENS_PER_REQUEST(self) -> int:
+        """Maximum number of tokens allowed per individual LLM request."""
+        return int(os.getenv("LLM_MAX_TOKENS_PER_REQUEST", "4096"))
+
+    @property
+    def LLM_THROTTLE_SLEEP(self) -> float:
+        """Number of seconds to sleep when throttling is required."""
+        return float(os.getenv("LLM_THROTTLE_SLEEP", "1.0"))
+
     @property
     def LLM_CONTEXT_ENRICHMENT_PROMPT(self) -> str | None:
         return os.getenv("ZRB_LLM_CONTEXT_ENRICHMENT_PROMPT", None)

zrb-1.8.5/src/zrb/llm_rate_limitter.py

@@ -0,0 +1,105 @@
+import asyncio
+import time
+from collections import deque
+from typing import Callable
+
+from zrb.config import CFG
+
+
+class LLMRateLimiter:
+    """
+    Helper class to enforce LLM API rate limits and throttling.
+    Tracks requests and tokens in a rolling 60-second window.
+    """
+
+    def __init__(
+        self,
+        max_requests_per_minute: int | None = None,
+        max_tokens_per_minute: int | None = None,
+        max_tokens_per_request: int | None = None,
+        throttle_sleep: float | None = None,
+        token_counter_fn: Callable[[str], int] | 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._throttle_sleep = throttle_sleep
+        self._token_counter_fn = token_counter_fn
+        self.request_times = deque()
+        self.token_times = deque()
+
+    @property
+    def max_requests_per_minute(self) -> int:
+        if self._max_requests_per_minute is not None:
+            return self._max_requests_per_minute
+        return CFG.LLM_MAX_REQUESTS_PER_MINUTE
+
+    @property
+    def max_tokens_per_minute(self) -> int:
+        if self._max_tokens_per_minute is not None:
+            return self._max_tokens_per_minute
+        return CFG.LLM_MAX_TOKENS_PER_MINUTE
+
+    @property
+    def max_tokens_per_request(self) -> int:
+        if self._max_tokens_per_request is not None:
+            return self._max_tokens_per_request
+        return CFG.LLM_MAX_TOKENS_PER_REQUEST
+
+    @property
+    def throttle_sleep(self) -> float:
+        if self._throttle_sleep is not None:
+            return self._throttle_sleep
+        return CFG.LLM_THROTTLE_SLEEP
+
+    @property
+    def token_counter_fn(self) -> Callable[[str], int]:
+        if self._token_counter_fn is not None:
+            return self._token_counter_fn
+        return lambda x: len(x.split())
+
+    def set_max_requests_per_minute(self, value: int):
+        self._max_requests_per_minute = value
+
+    def set_max_tokens_per_minute(self, value: int):
+        self._max_tokens_per_minute = value
+
+    def set_max_tokens_per_request(self, value: int):
+        self._max_tokens_per_request = value
+
+    def set_throttle_sleep(self, value: float):
+        self._throttle_sleep = value
+
+    def set_token_counter_fn(self, fn: Callable[[str], int]):
+        self._token_counter_fn = fn
+
+    async def throttle(self, prompt: str):
+        now = time.time()
+        tokens = self.token_counter_fn(prompt)
+        # Clean up old entries
+        while self.request_times and now - self.request_times[0] > 60:
+            self.request_times.popleft()
+        while self.token_times and now - self.token_times[0][0] > 60:
+            self.token_times.popleft()
+        # Check per-request token limit
+        if tokens > self.max_tokens_per_request:
+            raise ValueError(
+                f"Request exceeds max_tokens_per_request ({self.max_tokens_per_request})."
+            )
+        # Wait if over per-minute request or token limit
+        while (
+            len(self.request_times) >= self.max_requests_per_minute
+            or sum(t for _, t in self.token_times) + tokens > self.max_tokens_per_minute
+        ):
+            await asyncio.sleep(self.throttle_sleep)
+            now = time.time()
+            while self.request_times and now - self.request_times[0] > 60:
+                self.request_times.popleft()
+            while self.token_times and now - self.token_times[0][0] > 60:
+                self.token_times.popleft()
+        # Record this request
+        self.request_times.append(now)
+        self.token_times.append((now, tokens))
+
+
+llm_rate_limitter = LLMRateLimiter()

{zrb-1.8.3 → zrb-1.8.5}/src/zrb/session/session.py

@@ -186,17 +186,16 @@ class Session(AnySession):
 
     def defer_monitoring(self, task: AnyTask, coro: Coroutine):
         self._register_single_task(task)
-        self._monitoring_coros[task] = asyncio.create_task(coro)
+        self._monitoring_coros[task] = coro
 
     def defer_action(self, task: AnyTask, coro: Coroutine):
         self._register_single_task(task)
-        self._action_coros[task] = asyncio.create_task(coro)
+        self._action_coros[task] = coro
 
     def defer_coro(self, coro: Coroutine):
-        task = asyncio.create_task(coro)
-        self._coros.append(task)
+        self._coros.append(coro)
         self._coros = [
-            existing_task for existing_task in self._coros if not existing_task.done()
+            existing_coro for existing_coro in self._coros if not existing_coro.done()
         ]
 
     async def wait_deferred(self):