zrb 1.8.9__tar.gz → 1.8.10__tar.gz

{zrb-1.8.9 → zrb-1.8.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.8.9
+Version: 1.8.10
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -19,18 +19,19 @@ Provides-Extra: rag
 Requires-Dist: beautifulsoup4 (>=4.13.3,<5.0.0)
 Requires-Dist: black (>=25.1.0,<25.2.0)
 Requires-Dist: chromadb (>=0.6.3,<0.7.0) ; extra == "rag" or extra == "all"
-Requires-Dist: fastapi[standard] (>=0.115.12,<0.116.0)
+Requires-Dist: fastapi[standard] (>=0.115.14,<0.116.0)
 Requires-Dist: isort (>=6.0.1,<6.1.0)
 Requires-Dist: libcst (>=1.7.0,<2.0.0)
-Requires-Dist: openai (>=1.76.0,<2.0.0) ; extra == "rag" or extra == "all"
+Requires-Dist: openai (>=1.86.0,<2.0.0) ; extra == "rag" or extra == "all"
 Requires-Dist: pdfplumber (>=0.11.6,<0.12.0) ; extra == "rag" or extra == "all"
-Requires-Dist: playwright (>=1.51.0,<2.0.0) ; extra == "playwright" or extra == "all"
+Requires-Dist: playwright (>=1.53.0,<2.0.0) ; extra == "playwright" or extra == "all"
 Requires-Dist: psutil (>=7.0.0,<8.0.0)
-Requires-Dist: pydantic-ai (>=0.2.18,<0.3.0)
+Requires-Dist: pydantic-ai (>=0.3.4,<0.4.0)
 Requires-Dist: pyjwt (>=2.10.1,<3.0.0)
-Requires-Dist: python-dotenv (>=1.1.0,<2.0.0)
+Requires-Dist: python-dotenv (>=1.1.1,<2.0.0)
 Requires-Dist: python-jose[cryptography] (>=3.4.0,<4.0.0)
 Requires-Dist: requests (>=2.32.4,<3.0.0)
+Requires-Dist: tiktoken (>=0.8.0,<0.9.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
@@ -81,9 +82,11 @@ Or run our installation script to set up Zrb along with all prerequisites:
 bash -c "$(curl -fsSL https://raw.githubusercontent.com/state-alchemists/zrb/main/install.sh)"
 ```
 
+You can also [run Zrb as container](https://github.com/state-alchemists/zrb?tab=readme-ov-file#-run-zrb-as-a-container)
+
 # 🍲 Quick Start: Build Your First Automation Workflow
 
-Zrb empowers you to create custom automation tasks using Python. This guide shows you how to define two simple tasks: one to generate a PlantUML script from your source code and another to convert that script into a PNG image.
+Zrb empowers you to create custom automation tasks using Python. This guide shows you how to define two simple tasks: one to generate a Mermaid script from your source code and another to convert that script into a PNG image.
 
 ## 1. Create Your Task Definition File
 
@@ -92,59 +95,61 @@ Place a file named `zrb_init.py` in a directory that's accessible from your proj
 Add the following content to your zrb_init.py:
 
 ```python
-import os
 from zrb import cli, LLMTask, CmdTask, StrInput, Group
-from zrb.builtin.llm.tool.file import (
-    list_files, read_from_file, search_files, write_to_file
-)
-
+from zrb.builtin.llm.tool.code import analyze_repo
+from zrb.builtin.llm.tool.file import write_to_file
 
-CURRENT_DIR = os.getcwd()
 
-# Create a group for UML-related tasks
-uml_group = cli.add_group(Group(name="uml", description="UML related tasks"))
+# Create a group for Mermaid-related tasks
+mermaid_group = cli.add_group(Group(name="mermaid", description="🧜 Mermaid diagram related tasks"))
 
-# Task 1: Generate a PlantUML script from your source code
-make_uml_script = uml_group.add_task(
+# Task 1: Generate a Mermaid script from your source code
+make_mermaid_script = mermaid_group.add_task(
     LLMTask(
         name="make-script",
-        description="Creating plantuml diagram based on source code in current directory",
-        input=StrInput(name="diagram", default="state diagram"),
+        description="Creating mermaid diagram based on source code in current directory",
+        input=[
+            StrInput(name="dir", default="./"),
+            StrInput(name="diagram", default="state-diagram"),
+        ],
         message=(
-            f"Read all necessary files in {CURRENT_DIR}, "
-            "make a {ctx.input.diagram} in plantuml format. "
-            f"Write the script into {CURRENT_DIR}/{{ctx.input.diagram}}.uml"
+            "Read all necessary files in {ctx.input.dir}, "
+            "make a {ctx.input.diagram} in mermaid format. "
+            "Write the script into `{ctx.input.dir}/{ctx.input.diagram}.mmd`"
         ),
         tools=[
-            list_files, read_from_file, search_files, write_to_file
+            analyze_repo, write_to_file
         ],
     )
 )
 
-# Task 2: Convert the PlantUML script into a PNG image
-make_uml_image = uml_group.add_task(
+# Task 2: Convert the Mermaid script into a PNG image
+make_mermaid_image = mermaid_group.add_task(
     CmdTask(
         name="make-image",
         description="Creating png based on source code in current directory",
-        input=StrInput(name="diagram", default="state diagram"),
-        cmd="plantuml -tpng '{ctx.input.diagram}.uml'",
-        cwd=CURRENT_DIR,
+        input=[
+            StrInput(name="dir", default="./"),
+            StrInput(name="diagram", default="state-diagram"),
+        ],
+        cmd="mmdc -i '{ctx.input.diagram}.mmd' -o '{ctx.input.diagram}.png'",
+        cwd="{ctx.input.dir}",
     )
 )
 
 # Set up the dependency: the image task runs after the script is created
-make_uml_script >> make_uml_image
+make_mermaid_script >> make_mermaid_image
 ```
 
 **What This Does**
 
 - **Task 1 – make-script**:
 
-    Uses an LLM to read all files in your current directory and generate a PlantUML script (e.g., `state diagram.uml`).
+    Uses an LLM to read all files in your current directory and generate a Mermaid script (e.g., `state diagram.mmd`).
 
 - **Task 2 – make-image**:
 
-    Executes a command that converts the PlantUML script into a PNG image (e.g., `state diagram.png`). This task will run only after the script has been generated.
+    Executes a command that converts the Mermaid script into a PNG image (e.g., `state diagram.png`). This task will run only after the script has been generated.
 
 
 ## 2. Run Your Tasks
@@ -161,18 +166,19 @@ After setting up your tasks, you can execute them from any project. For example:
 - Create a state diagram:
 
     ```bash
-    zrb uml make-image --diagram "state diagram"
+    zrb mermaid make-image --diagram "state diagram" --dir ./
     ```
 
 - Or use the interactive mode:
 
     ```bash
-    zrb uml make-image
+    zrb mermaid make-image
     ```
 
     Zrb will prompt:
 
     ```bash
+    dir [./]:
     diagram [state diagram]:
     ```
 
@@ -194,7 +200,7 @@ Then open your browser and visit `http://localhost:21213`
 ![Zrb Web UI](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/zrb-web-ui.png)
 
 
-# 🐋 Running Zrb as a Container
+# 🐋 Run Zrb as a Container
 
 Zrb can be run in a containerized environment, offering two distinct versions to suit different needs:
 

{zrb-1.8.9 → zrb-1.8.10}/README.md

@@ -43,9 +43,11 @@ Or run our installation script to set up Zrb along with all prerequisites:
 bash -c "$(curl -fsSL https://raw.githubusercontent.com/state-alchemists/zrb/main/install.sh)"
 ```
 
+You can also [run Zrb as container](https://github.com/state-alchemists/zrb?tab=readme-ov-file#-run-zrb-as-a-container)
+
 # 🍲 Quick Start: Build Your First Automation Workflow
 
-Zrb empowers you to create custom automation tasks using Python. This guide shows you how to define two simple tasks: one to generate a PlantUML script from your source code and another to convert that script into a PNG image.
+Zrb empowers you to create custom automation tasks using Python. This guide shows you how to define two simple tasks: one to generate a Mermaid script from your source code and another to convert that script into a PNG image.
 
 ## 1. Create Your Task Definition File
 
@@ -54,59 +56,61 @@ Place a file named `zrb_init.py` in a directory that's accessible from your proj
 Add the following content to your zrb_init.py:
 
 ```python
-import os
 from zrb import cli, LLMTask, CmdTask, StrInput, Group
-from zrb.builtin.llm.tool.file import (
-    list_files, read_from_file, search_files, write_to_file
-)
-
+from zrb.builtin.llm.tool.code import analyze_repo
+from zrb.builtin.llm.tool.file import write_to_file
 
-CURRENT_DIR = os.getcwd()
 
-# Create a group for UML-related tasks
-uml_group = cli.add_group(Group(name="uml", description="UML related tasks"))
+# Create a group for Mermaid-related tasks
+mermaid_group = cli.add_group(Group(name="mermaid", description="🧜 Mermaid diagram related tasks"))
 
-# Task 1: Generate a PlantUML script from your source code
-make_uml_script = uml_group.add_task(
+# Task 1: Generate a Mermaid script from your source code
+make_mermaid_script = mermaid_group.add_task(
     LLMTask(
         name="make-script",
-        description="Creating plantuml diagram based on source code in current directory",
-        input=StrInput(name="diagram", default="state diagram"),
+        description="Creating mermaid diagram based on source code in current directory",
+        input=[
+            StrInput(name="dir", default="./"),
+            StrInput(name="diagram", default="state-diagram"),
+        ],
         message=(
-            f"Read all necessary files in {CURRENT_DIR}, "
-            "make a {ctx.input.diagram} in plantuml format. "
-            f"Write the script into {CURRENT_DIR}/{{ctx.input.diagram}}.uml"
+            "Read all necessary files in {ctx.input.dir}, "
+            "make a {ctx.input.diagram} in mermaid format. "
+            "Write the script into `{ctx.input.dir}/{ctx.input.diagram}.mmd`"
         ),
         tools=[
-            list_files, read_from_file, search_files, write_to_file
+            analyze_repo, write_to_file
         ],
     )
 )
 
-# Task 2: Convert the PlantUML script into a PNG image
-make_uml_image = uml_group.add_task(
+# Task 2: Convert the Mermaid script into a PNG image
+make_mermaid_image = mermaid_group.add_task(
     CmdTask(
         name="make-image",
         description="Creating png based on source code in current directory",
-        input=StrInput(name="diagram", default="state diagram"),
-        cmd="plantuml -tpng '{ctx.input.diagram}.uml'",
-        cwd=CURRENT_DIR,
+        input=[
+            StrInput(name="dir", default="./"),
+            StrInput(name="diagram", default="state-diagram"),
+        ],
+        cmd="mmdc -i '{ctx.input.diagram}.mmd' -o '{ctx.input.diagram}.png'",
+        cwd="{ctx.input.dir}",
     )
 )
 
 # Set up the dependency: the image task runs after the script is created
-make_uml_script >> make_uml_image
+make_mermaid_script >> make_mermaid_image
 ```
 
 **What This Does**
 
 - **Task 1 – make-script**:
 
-    Uses an LLM to read all files in your current directory and generate a PlantUML script (e.g., `state diagram.uml`).
+    Uses an LLM to read all files in your current directory and generate a Mermaid script (e.g., `state diagram.mmd`).
 
 - **Task 2 – make-image**:
 
-    Executes a command that converts the PlantUML script into a PNG image (e.g., `state diagram.png`). This task will run only after the script has been generated.
+    Executes a command that converts the Mermaid script into a PNG image (e.g., `state diagram.png`). This task will run only after the script has been generated.
 
 
 ## 2. Run Your Tasks
@@ -123,18 +127,19 @@ After setting up your tasks, you can execute them from any project. For example:
 - Create a state diagram:
 
     ```bash
-    zrb uml make-image --diagram "state diagram"
+    zrb mermaid make-image --diagram "state diagram" --dir ./
     ```
 
 - Or use the interactive mode:
 
     ```bash
-    zrb uml make-image
+    zrb mermaid make-image
     ```
 
     Zrb will prompt:
 
     ```bash
+    dir [./]:
     diagram [state diagram]:
     ```
 
@@ -156,7 +161,7 @@ Then open your browser and visit `http://localhost:21213`
 ![Zrb Web UI](https://raw.githubusercontent.com/state-alchemists/zrb/main/_images/zrb-web-ui.png)
 
 
-# 🐋 Running Zrb as a Container
+# 🐋 Run Zrb as a Container
 
 Zrb can be run in a containerized environment, offering two distinct versions to suit different needs:
 

{zrb-1.8.9 → zrb-1.8.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.8.9"
+version = "1.8.10"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"
@@ -38,7 +38,7 @@ exclude = [
 
 [tool.poetry.dependencies]
 python = ">=3.10.0,<4.0.0"
-python-dotenv = "^1.1.0"
+python-dotenv = "^1.1.1"
 black = "~25.1.0"
 isort = "~6.0.1"
 requests = "^2.32.4"
@@ -46,12 +46,13 @@ libcst = "^1.7.0"
 chromadb = {version = "^0.6.3", optional = true}
 pdfplumber = {version = "^0.11.6", optional = true}
 beautifulsoup4 = "^4.13.3"
-fastapi = {extras = ["standard"], version = "^0.115.12"}
+fastapi = {extras = ["standard"], version = "^0.115.14"}
 python-jose = {extras = ["cryptography"], version = "^3.4.0"}
 ulid-py = "^1.1.0"
-pydantic-ai = "^0.2.18"
-openai = {version = "^1.76.0", optional = true}
-playwright = {version = "^1.51.0", optional = true}
+pydantic-ai = "^0.3.4"
+tiktoken = "^0.8.0"
+openai = {version = "^1.86.0", optional = true}
+playwright = {version = "^1.53.0", optional = true}
 psutil = "^7.0.0"
 pyjwt = "^2.10.1"
 
@@ -62,7 +63,7 @@ playwright = ["playwright"]
 all = ["chromadb", "pdfplumber", "openai", "playwright"]
 
 [tool.poetry.group.dev.dependencies]
-alembic = "^1.15.2" # FastApp dependencies
+alembic = "^1.16.2" # FastApp dependencies
 autopep8 = "^2.0.4"
 flake8 = "~7.2.0"
 pytest = "~8.3.5"

{zrb-1.8.9 → zrb-1.8.10}/src/zrb/builtin/llm/tool/code.py

@@ -4,6 +4,7 @@ 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
+from zrb.llm_rate_limitter import llm_rate_limitter
 
 _EXTRACT_INFO_FROM_REPO_SYSTEM_PROMPT = """
 You are an extraction info agent.
@@ -81,8 +82,8 @@ async def analyze_repo(
     goal: str,
     extensions: list[str] = _DEFAULT_EXTENSIONS,
     exclude_patterns: list[str] = DEFAULT_EXCLUDED_PATTERNS,
-    extraction_char_limit: int = 150000,
-    summarization_char_limit: int = 150000,
+    extraction_token_limit: int = 30000,
+    summarization_token_limit: int = 30000,
 ) -> str:
     """
     Extract and summarize information from a directory that probably
@@ -101,9 +102,9 @@ async def analyze_repo(
             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
+        extraction_token_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
+        summarization_token_limit(Optional[int]): Max resource content char length
             the summarization assistant able to handle. Defaults to 150000
     Returns:
         str: The analysis result
@@ -117,22 +118,19 @@ async def analyze_repo(
         ctx,
         file_metadatas=file_metadatas,
         goal=goal,
-        char_limit=extraction_char_limit,
+        token_limit=extraction_token_limit,
     )
+    if len(extracted_infos) == 1:
+        return extracted_infos[0]
     ctx.print("Summarization")
-    summarized_infos = await _summarize_info(
-        ctx,
-        extracted_infos=extracted_infos,
-        goal=goal,
-        char_limit=summarization_char_limit,
-    )
+    summarized_infos = extracted_infos
     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,
+            token_limit=summarization_token_limit,
         )
     return summarized_infos[0]
 
@@ -165,7 +163,7 @@ async def _extract_info(
     ctx: AnyContext,
     file_metadatas: list[dict[str, str]],
     goal: str,
-    char_limit: int,
+    token_limit: int,
 ) -> list[str]:
     extract = create_sub_agent_tool(
         tool_name="extract",
@@ -174,27 +172,31 @@ async def _extract_info(
     )
     extracted_infos = []
     content_buffer = []
-    current_char_count = 0
+    current_token_count = 0
     for metadata in file_metadatas:
         path = metadata.get("path", "")
         content = metadata.get("content", "")
         file_obj = {"path": path, "content": content}
         file_str = json.dumps(file_obj)
-        if current_char_count + len(file_str) > char_limit:
+        if current_token_count + llm_rate_limitter.count_token(file_str) > token_limit:
             if content_buffer:
                 prompt = _create_extract_info_prompt(goal, content_buffer)
-                extracted_info = await extract(ctx, prompt)
+                extracted_info = await extract(
+                    ctx, llm_rate_limitter.clip_prompt(prompt, token_limit)
+                )
                 extracted_infos.append(extracted_info)
             content_buffer = [file_obj]
-            current_char_count = len(file_str)
+            current_token_count = llm_rate_limitter.count_token(file_str)
         else:
             content_buffer.append(file_obj)
-            current_char_count += len(file_str)
+            current_token_count += llm_rate_limitter.count_token(file_str)
 
     # 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_info = await extract(
+            ctx, llm_rate_limitter.clip_prompt(prompt, token_limit)
+        )
         extracted_infos.append(extracted_info)
     return extracted_infos
 
@@ -212,7 +214,7 @@ async def _summarize_info(
     ctx: AnyContext,
     extracted_infos: list[str],
     goal: str,
-    char_limit: int,
+    token_limit: int,
 ) -> list[str]:
     summarize = create_sub_agent_tool(
         tool_name="extract",
@@ -222,10 +224,13 @@ async def _summarize_info(
     summarized_infos = []
     content_buffer = ""
     for extracted_info in extracted_infos:
-        if len(content_buffer) + len(extracted_info) > char_limit:
+        new_prompt = content_buffer + extracted_info
+        if llm_rate_limitter.count_token(new_prompt) > token_limit:
             if content_buffer:
                 prompt = _create_summarize_info_prompt(goal, content_buffer)
-                summarized_info = await summarize(ctx, prompt)
+                summarized_info = await summarize(
+                    ctx, llm_rate_limitter.clip_prompt(prompt, token_limit)
+                )
                 summarized_infos.append(summarized_info)
             content_buffer = extracted_info
         else:
@@ -234,7 +239,9 @@ async def _summarize_info(
     # 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_info = await summarize(
+            ctx, llm_rate_limitter.clip_prompt(prompt, token_limit)
+        )
         summarized_infos.append(summarized_info)
     return summarized_infos
 

{zrb-1.8.9 → zrb-1.8.10}/src/zrb/builtin/llm/tool/file.py

@@ -6,6 +6,7 @@ from typing import Any, Optional
 
 from zrb.builtin.llm.tool.sub_agent import create_sub_agent_tool
 from zrb.context.any_context import AnyContext
+from zrb.llm_rate_limitter import llm_rate_limitter
 from zrb.util.file import read_file, read_file_with_line_numbers, write_file
 
 _EXTRACT_INFO_FROM_FILE_SYSTEM_PROMPT = """
@@ -464,7 +465,9 @@ def apply_diff(
         raise RuntimeError(f"Unexpected error applying diff to {path}: {e}")
 
 
-async def analyze_file(ctx: AnyContext, path: str, query: str) -> str:
+async def analyze_file(
+    ctx: AnyContext, path: str, query: str, token_limit: int = 30000
+) -> str:
     """Analyze file using LLM capability to reduce context usage.
     Use this tool for:
     - summarization
@@ -474,6 +477,7 @@ async def analyze_file(ctx: AnyContext, path: str, query: str) -> str:
     Args:
         path (str): File path to be analyze. Pass exactly as provided, including '~'.
         query(str): Instruction to analyze the file
+        token_limit(Optional[int]): Max token length to be taken from file
     Returns:
         str: The analysis result
     Raises:
@@ -489,9 +493,8 @@ async def analyze_file(ctx: AnyContext, path: str, query: str) -> str:
         system_prompt=_EXTRACT_INFO_FROM_FILE_SYSTEM_PROMPT,
         tools=[read_from_file, search_files],
     )
-    return await _analyze_file(
-        ctx,
-        json.dumps(
-            {"instruction": query, "file_path": abs_path, "file_content": file_content}
-        ),
+    payload = json.dumps(
+        {"instruction": query, "file_path": abs_path, "file_content": file_content}
     )
+    clipped_payload = llm_rate_limitter.clip_prompt(payload, token_limit)
+    return await _analyze_file(ctx, clipped_payload)

{zrb-1.8.9 → zrb-1.8.10}/src/zrb/builtin/llm/tool/sub_agent.py

@@ -1,7 +1,7 @@
 import json
 from collections.abc import Callable
 from textwrap import dedent
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Coroutine
 
 if TYPE_CHECKING:
     from pydantic_ai import Tool
@@ -33,7 +33,7 @@ def create_sub_agent_tool(
     model_settings: ModelSettings | None = None,
     tools: list[ToolOrCallable] = [],
     mcp_servers: list[MCPServer] = [],
-) -> Callable[[AnyContext, str], str]:
+) -> Callable[[AnyContext, str], Coroutine[Any, Any, str]]:
     """
     Create an LLM "sub-agent" tool function for use by a main LLM agent.
 
@@ -97,6 +97,8 @@ def create_sub_agent_tool(
             tools=tools,
             mcp_servers=mcp_servers,
         )
+
+        sub_agent_run = None
         # Run the sub-agent iteration
         # Start with an empty history for the sub-agent
         sub_agent_run = await run_agent_iteration(

{zrb-1.8.9 → zrb-1.8.10}/src/zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/requirements.txt

@@ -1,5 +1,5 @@
-fastapi[standard]~=0.115.12
-alembic~=1.15.2
+fastapi[standard]~=0.115.14
+alembic~=1.16.2
 sqlmodel~=0.0.24
 ulid-py~=1.1.0
 passlib~=1.7.4

{zrb-1.8.9 → zrb-1.8.10}/src/zrb/config.py

@@ -245,12 +245,12 @@ class Config:
     @property
     def LLM_MAX_TOKENS_PER_MINUTE(self) -> int:
         """Maximum number of LLM tokens allowed per minute."""
-        return int(os.getenv("ZRB_LLM_MAX_TOKENS_PER_MINUTE", "120000"))
+        return int(os.getenv("ZRB_LLM_MAX_TOKENS_PER_MINUTE", "200000"))
 
     @property
     def LLM_MAX_TOKENS_PER_REQUEST(self) -> int:
         """Maximum number of tokens allowed per individual LLM request."""
-        return int(os.getenv("ZRB_LLM_MAX_TOKENS_PER_REQUEST", "30000"))
+        return int(os.getenv("ZRB_LLM_MAX_TOKENS_PER_REQUEST", "50000"))
 
     @property
     def LLM_THROTTLE_SLEEP(self) -> float:

{zrb-1.8.9 → zrb-1.8.10}/src/zrb/llm_rate_limitter.py

@@ -3,9 +3,16 @@ import time
 from collections import deque
 from typing import Callable
 
+import tiktoken
+
 from zrb.config import CFG
 
 
+def _estimate_token(text: str) -> int:
+    enc = tiktoken.encoding_for_model("gpt-4o")
+    return len(enc.encode(text))
+
+
 class LLMRateLimiter:
     """
     Helper class to enforce LLM API rate limits and throttling.
@@ -53,10 +60,10 @@ class LLMRateLimiter:
         return CFG.LLM_THROTTLE_SLEEP
 
     @property
-    def token_counter_fn(self) -> Callable[[str], int]:
+    def count_token(self) -> Callable[[str], int]:
         if self._token_counter_fn is not None:
             return self._token_counter_fn
-        return lambda x: len(x.split())
+        return _estimate_token
 
     def set_max_requests_per_minute(self, value: int):
         self._max_requests_per_minute = value
@@ -73,9 +80,22 @@ class LLMRateLimiter:
     def set_token_counter_fn(self, fn: Callable[[str], int]):
         self._token_counter_fn = fn
 
+    def clip_prompt(self, prompt: str, limit: int) -> str:
+        token_count = self.count_token(prompt)
+        if token_count <= limit:
+            return prompt
+        while token_count > limit:
+            prompt_parts = prompt.split(" ")
+            last_part_index = len(prompt_parts) - 2
+            clipped_prompt = " ".join(prompt_parts[:last_part_index])
+            token_count = self.count_token(clipped_prompt)
+            if token_count < limit:
+                return clipped_prompt
+        return prompt[:limit]
+
     async def throttle(self, prompt: str):
         now = time.time()
-        tokens = self.token_counter_fn(prompt)
+        tokens = self.count_token(prompt)
         # Clean up old entries
         while self.request_times and now - self.request_times[0] > 60:
             self.request_times.popleft()