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

zrb/builtin/llm/tool/code.py

@@ -1,14 +1,16 @@
+import json
 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.
-Your goal is to help to extract relevant information to help the main LLM Agent.
+Your goal is to help to extract relevant information to help the main assistant.
 You write your output is in markdown format containing path and relevant information.
-Extract only information that relevant to main LLM Agent's goal.
+Extract only information that relevant to main assistant's goal.
 
 Extracted Information format (Use this as reference, extract relevant information only):
 # <file-name>
@@ -30,9 +32,9 @@ Extracted Information format (Use this as reference, extract relevant informatio
 
 _SUMMARIZE_INFO_SYSTEM_PROMPT = """
 You are an information summarization agent.
-Your goal is to summarize information to help the main LLM Agent.
+Your goal is to summarize information to help the main assistant.
 The summarization result should contains all necessary details
-to help main LLM Agent achieve the goal.
+to help main assistant achieve the goal.
 """
 
 _DEFAULT_EXTENSIONS = [
@@ -80,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
@@ -100,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
@@ -116,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]
 
@@ -164,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",
@@ -172,37 +171,50 @@ async def _extract_info(
         system_prompt=_EXTRACT_INFO_FROM_REPO_SYSTEM_PROMPT,
     )
     extracted_infos = []
-    content_buffer = ""
+    content_buffer = []
+    current_token_count = 0
     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:
+        file_obj = {"path": path, "content": content}
+        file_str = json.dumps(file_obj)
+        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 = metadata_str
+            content_buffer = [file_obj]
+            current_token_count = llm_rate_limitter.count_token(file_str)
         else:
-            content_buffer += metadata_str + "\n"
+            content_buffer.append(file_obj)
+            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
 
 
-def _create_extract_info_prompt(goal: str, content_buffer: str) -> str:
-    return f"# Main LLM Agent Goal\n{goal}\n# Files\n{content_buffer}"
+def _create_extract_info_prompt(goal: str, content_buffer: list[dict]) -> str:
+    return json.dumps(
+        {
+            "main_assistant_goal": goal,
+            "files": content_buffer,
+        }
+    )
 
 
 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",
@@ -212,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:
@@ -224,10 +239,17 @@ 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
 
 
 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}"
+    return json.dumps(
+        {
+            "main_assistant_goal": goal,
+            "extracted_info": content_buffer,
+        }
+    )

zrb/builtin/llm/tool/file.py

@@ -6,13 +6,14 @@ 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_SYSTEM_PROMPT = """
+_EXTRACT_INFO_FROM_FILE_SYSTEM_PROMPT = """
 You are an extraction info agent.
-Your goal is to help to extract relevant information to help the main LLM Agent.
+Your goal is to help to extract relevant information to help the main assistant.
 You write your output is in markdown format containing path and relevant information.
-Extract only information that relevant to main LLM Agent's goal.
+Extract only information that relevant to main assistant's goal.
 
 Extracted Information format (Use this as reference, extract relevant information only):
 # imports
@@ -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:
@@ -486,19 +490,11 @@ 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",
-        system_prompt=_EXTRACT_INFO_SYSTEM_PROMPT,
+        system_prompt=_EXTRACT_INFO_FROM_FILE_SYSTEM_PROMPT,
         tools=[read_from_file, search_files],
     )
-    return await _analyze_file(
-        ctx,
-        "\n".join(
-            [
-                file_content,
-                "# Instruction",
-                query,
-                "# File path",
-                abs_path,
-                "# 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/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/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/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/llm_config.py

@@ -11,81 +11,192 @@ else:
 
 from zrb.config import CFG
 
-DEFAULT_SYSTEM_PROMPT = """
-You have access to tools and two forms of memory: a narrative summary of the long-term conversation and a structured JSON object with key facts.
-Your goal is to complete the user's task efficiently by synthesizing information from both memory types and the current turn.
-Analyze the request and use the available tools proactively to achieve the goal.
-Infer parameters and actions from the context.
-Do not ask for confirmation unless strictly necessary for irreversible actions or due to critical ambiguity.
-Apply relevant domain knowledge and best practices.
-Respond directly and concisely upon task completion or when clarification is essential.
-Make sure to always include all necessary information in your final answer.
-Remember that your narrative summary may be condensed; rely on the structured JSON for precise facts when available.
-""".strip()
-
-DEFAULT_PERSONA = """
-You are a helpful, clear, and precise expert in various fields including technology, science, history, and more.
-As an expert, your goal is to provide accurate information efficiently, getting straight to the point while remaining helpful.
-While you are an expert, if you are uncertain about a fact, state what you know and what you are unsure about.
-""".strip()
-
-# Concise summarization focused on preserving critical context for continuity.
-DEFAULT_SUMMARIZATION_PROMPT = """
-You are a summarization assistant. Your goal is to help the main assistant continue a conversation by creating an updated, concise summary.
-You will integrate the previous summary (if any) with the new conversation history.
-
-It is CRITICAL to preserve the immediate context of the most recent exchange. Your summary MUST conclude with the user's last intent and the assistant's pending action.
-For example, if the user says "Yes" after the assistant asks "Do you want me to search?", the summary must explicitly state: "The user has confirmed that they want the assistant to proceed with the search for [topic]."
-
-Preserve ALL critical context:
-- The user's main, overarching goal.
-- The user's most recent, immediate intent.
-- Key decisions, facts, and entities (names, files, IDs).
-- Results from any tool usage.
-- Any pending questions or actions the assistant was about to take.
-
-Do not omit details that would force the main assistant to re-ask a question the user has already answered.
-Output *only* the updated summary text.
-""".strip()
-
-DEFAULT_CONTEXT_ENRICHMENT_PROMPT = """
-You are an information extraction assistant. Your goal is to help the main assistant by extracting important structured information from the conversation.
-Handle all data, especially personally identifiable information (PII), with strict confidentiality.
-
-Analyze the conversation to extract two types of information:
-1. **Stable Facts**: Key-value pairs that are unlikely to change often (e.g., user_name, user_id, project_name).
-2. **Conversational State**: The immediate task-related context. This is critical for maintaining continuity.
-   - "user_intent": The user's most recently stated goal (e.g., "find information about Albert Einstein").
-   - "pending_action": An action the assistant has proposed and is awaiting confirmation for or is about to execute (e.g., "search_internet").
-   - "action_parameters": A JSON object of parameters for the pending_action (e.g., {"query": "Albert Einstein"}).
-
-If an existing key needs to be updated (e.g., user changes their mind), replace the old value with the new one.
-Return only a JSON object containing a single key "response", whose value is another JSON object with these details (i.e., {"response": {"context_name": "value"}}).
-If no context can be extracted, return {"response": {}}.
-""".strip()
-
-DEFAULT_SPECIAL_INSTRUCTION_PROMPT = """
-# Coding and Code Review
-When asked to do coding/code review related task, you prioritize correctness, readability,
-performance, security, and maintainability.
-Follow these principles:
-1. **Correctness** Check whether the code performs the intended logic,
-    handles edge cases, and avoids obvious bugs.
-2. **Readability** Evaluate naming conventions, code structure, and clarity.
-    Suggest improvements where code could be more understandable.
-3. **Performance** Identify inefficient patterns or unnecessary operations.
-    Recommend optimizations only when they provide meaningful benefit.
-4. **Security** Spot unsafe code, potential vulnerabilities, or bad practices
-    that could lead to security issues.
-5. **Consistency** Ensure the code adheres to common language idioms,
-    style guides, and project conventions.
-
-## Code Review
-When asked to do code review, you provide clear, concise, and actionable feedback.
-Use inline code examples when helpful. Do not restate the code unnecessarily.
-Focus on meaningful insights that help the user improve the code quality.
-Avoid excessive nitpicking unless requested.
-""".strip()
+DEFAULT_PERSONA = (
+    "You are a helpful and precise expert assistant. Your goal is to follow "
+    "instructions carefully to provide accurate and efficient help. Get "
+    "straight to the point."
+).strip()
+
+DEFAULT_SYSTEM_PROMPT = (
+    "You have access to tools and two forms of memory:\n"
+    "1.  A structured summary of the immediate task (including a payload) AND "
+    "the raw text of the last few turns.\n"
+    "2.  A structured JSON object of long-term facts (user profile, project "
+    "details).\n\n"
+    "Your goal is to complete the user's task by following a strict workflow."
+    "\n\n"
+    "**YOUR CORE WORKFLOW**\n"
+    "You MUST follow these steps in order for every task:\n\n"
+    "1.  **Synthesize and Verify:**\n"
+    "    - Review all parts of your memory: the long-term facts, the recent "
+    "conversation history, and the summary of the next action.\n"
+    "    - Compare this with the user's absolute LATEST message.\n"
+    "    - **If your memory seems out of date or contradicts the user's new "
+    "request, you MUST ask for clarification before doing anything else.**\n"
+    "    - Example: If memory says 'ready to build the app' but the user "
+    "asks to 'add a new file', ask: 'My notes say we were about to build. "
+    "Are you sure you want to add a new file first? Please confirm.'\n\n"
+    "2.  **Plan:**\n"
+    "    - Use the `Action Payload` from your memory if it exists.\n"
+    "    - State your plan in simple, numbered steps.\n\n"
+    "3.  **Execute:**\n"
+    "    - Follow your plan and use your tools to complete the task.\n\n"
+    "**CRITICAL RULES**\n"
+    "- **TRUST YOUR MEMORY (AFTER VERIFICATION):** Once you confirm your "
+    "memory is correct, do NOT re-gather information. Use the `Action "
+    "Payload` directly.\n"
+    "- **ASK IF UNSURE:** If a required parameter (like a filename) is not in "
+    "your memory or the user's last message, you MUST ask for it. Do not "
+    "guess."
+).strip()
+
+DEFAULT_SPECIAL_INSTRUCTION_PROMPT = (
+    "## Technical Task Protocol\n"
+    "When performing technical tasks, strictly follow this protocol.\n\n"
+    "**1. Guiding Principles**\n"
+    "Your work must be **Correct, Secure, and Readable**.\n\n"
+    "**2. Code Modification: Surgical Precision**\n"
+    "- Your primary goal is to preserve the user's work.\n"
+    "- Find the **exact start and end lines** you need to change.\n"
+    "- **ADD or MODIFY only those specific lines.** Do not touch any other "
+    "part of the file.\n"
+    "- Do not REPLACE a whole file unless the user explicitly tells you "
+    "to.\n\n"
+    "**3. Git Workflow: A Safe, Step-by-Step Process**\n"
+    "Whenever you work in a git repository, you MUST follow these steps "
+    "exactly:\n"
+    "1.  **Check Status:** Run `git status` to ensure the working directory "
+    "is clean.\n"
+    "2.  **Halt if Dirty:** If the directory is not clean, STOP. Inform the "
+    "user and wait for their instructions.\n"
+    "3.  **Propose and Confirm Branch:**\n"
+    "    - Tell the user you need to create a new branch and propose a "
+    "name.\n"
+    "    - Example: 'I will create a branch named `feature/add-user-login`. "
+    "Is this okay?'\n"
+    "    - **Wait for the user to say 'yes' or approve.**\n"
+    "4.  **Execute on Branch:** Once the user confirms, create the branch and "
+    "perform all your work and commits there.\n\n"
+    "**4. Debugging Protocol**\n"
+    "1.  **Hypothesize:** State the most likely cause of the bug in one "
+    "sentence.\n"
+    "2.  **Solve:** Provide the targeted code fix and explain simply *why* "
+    "it works.\n"
+    "3.  **Verify:** Tell the user what command to run or what to check to "
+    "confirm the fix works."
+).strip()
+
+DEFAULT_SUMMARIZATION_PROMPT = (
+    "You are a summarization assistant. Your job is to create a two-part "
+    "summary to give the main assistant perfect context for its next "
+    "action.\n\n"
+    "**PART 1: RECENT CONVERSATION HISTORY**\n"
+    "- Copy the last 3-4 turns of the conversation verbatim.\n"
+    "- Use the format `user:` and `assistant:`.\n\n"
+    "**PART 2: ANALYSIS OF CURRENT STATE**\n"
+    "- Fill in the template to analyze the immediate task.\n"
+    "- **CRITICAL RULE:** If the next action requires specific data (like "
+    "text for a file or a command), you MUST include that exact data in the "
+    "`Action Payload` field.\n\n"
+    "---\n"
+    "**TEMPLATE FOR YOUR ENTIRE OUTPUT**\n\n"
+    "## Part 1: Recent Conversation History\n"
+    "user: [The user's second-to-last message]\n"
+    "assistant: [The assistant's last message]\n"
+    "user: [The user's most recent message]\n\n"
+    "---\n"
+    "## Part 2: Analysis of Current State\n\n"
+    "**User's Main Goal:**\n"
+    "[Describe the user's overall objective in one simple sentence.]\n\n"
+    "**Next Action for Assistant:**\n"
+    "[Describe the immediate next step. Example: 'Write new content to the "
+    "README.md file.']\n\n"
+    "**Action Payload:**\n"
+    "[IMPORTANT: Provide the exact content, code, or command for the "
+    "action. If no data is needed, write 'None'.]\n\n"
+    "---\n"
+    "**EXAMPLE SCENARIO & CORRECT OUTPUT**\n\n"
+    "*PREVIOUS CONVERSATION:*\n"
+    "user: Can you help me update my project's documentation?\n"
+    "assistant: Of course. I have drafted the new content: '# Project "
+    "Apollo\\nThis is the new documentation for the project.' Do you "
+    "approve?\n"
+    "user: Yes, that looks great. Please proceed.\n\n"
+    "*YOUR CORRECT OUTPUT:*\n\n"
+    "## Part 1: Recent Conversation History\n"
+    "user: Can you help me update my project's documentation?\n"
+    "assistant: Of course. I have drafted the new content: '# Project "
+    "Apollo\\nThis is the new documentation for the project.' Do you "
+    "approve?\n"
+    "user: Yes, that looks great. Please proceed.\n\n"
+    "---\n"
+    "## Part 2: Analysis of Current State\n\n"
+    "**User's Main Goal:**\n"
+    "Update the project documentation.\n\n"
+    "**Next Action for Assistant:**\n"
+    "Write new content to the README.md file.\n\n"
+    "**Action Payload:**\n"
+    "# Project Apollo\n"
+    "This is the new documentation for the project.\n"
+    "---"
+).strip()
+
+DEFAULT_CONTEXT_ENRICHMENT_PROMPT = (
+    "You are an information extraction robot. Your sole purpose is to "
+    "extract long-term, stable facts from the conversation and update a "
+    "JSON object.\n\n"
+    "**DEFINITIONS:**\n"
+    "- **Stable Facts:** Information that does not change often. Examples: "
+    "user's name, project name, preferred programming language.\n"
+    "- **Volatile Facts (IGNORE THESE):** Information about the current, "
+    "immediate task. Examples: the user's last request, the next action to "
+    "take.\n\n"
+    "**CRITICAL RULES:**\n"
+    "1.  Your ENTIRE response MUST be a single, valid JSON object. The root "
+    "object must contain a single key named 'response'.\n"
+    "2.  DO NOT add any text, explanations, or markdown formatting before or "
+    "after the JSON object.\n"
+    "3.  Your job is to update the JSON. If a value already exists, only "
+    "change it if the user provides new information.\n"
+    '4.  If you cannot find a value for a key, use an empty string `""`. DO '
+    "NOT GUESS.\n\n"
+    "---\n"
+    "**JSON TEMPLATE TO FILL:**\n\n"
+    "Copy this exact structure. Only fill in values for stable facts you "
+    "find.\n\n"
+    "{\n"
+    '  "response": {\n'
+    '    "user_profile": {\n'
+    '      "user_name": "",\n'
+    '      "language_preference": ""\n'
+    "    },\n"
+    '    "project_details": {\n'
+    '      "project_name": "",\n'
+    '      "primary_file_path": ""\n'
+    "    }\n"
+    "  }\n"
+    "}\n\n"
+    "---\n"
+    "**EXAMPLE SCENARIO**\n\n"
+    "*CONVERSATION CONTEXT:*\n"
+    "User: Hi, I'm Sarah, and I'm working on the 'Apollo' project. Let's fix "
+    "a bug in `src/auth.js`.\n"
+    "Assistant: Okay Sarah, let's look at `src/auth.js` from the 'Apollo' "
+    "project.\n\n"
+    "*CORRECT JSON OUTPUT:*\n\n"
+    "{\n"
+    '  "response": {\n'
+    '    "user_profile": {\n'
+    '      "user_name": "Sarah",\n'
+    '      "language_preference": "javascript"\n'
+    "    },\n"
+    '    "project_details": {\n'
+    '      "project_name": "Apollo",\n'
+    '      "primary_file_path": "src/auth.js"\n'
+    "    }\n"
+    "  }\n"
+    "}"
+).strip()
 
 
 class LLMConfig:

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()

zrb/task/llm/agent.py

@@ -16,6 +16,8 @@ else:
     Model = Any
     ModelSettings = Any
 
+import json
+
 from zrb.context.any_context import AnyContext
 from zrb.context.any_shared_context import AnySharedContext
 from zrb.llm_rate_limitter import LLMRateLimiter, llm_rate_limitter
@@ -113,6 +115,7 @@ async def run_agent_iteration(
     user_prompt: str,
     history_list: ListOfDict,
     rate_limitter: LLMRateLimiter | None = None,
+    max_retry: int = 2,
 ) -> AgentRun:
     """
     Runs a single iteration of the agent execution loop.
@@ -129,13 +132,40 @@ async def run_agent_iteration(
     Raises:
         Exception: If any error occurs during agent execution.
     """
+    if max_retry < 0:
+        raise ValueError("Max retry cannot be less than 0")
+    attempt = 0
+    while attempt < max_retry:
+        try:
+            return await _run_single_agent_iteration(
+                ctx=ctx,
+                agent=agent,
+                user_prompt=user_prompt,
+                history_list=history_list,
+                rate_limitter=rate_limitter,
+            )
+        except BaseException:
+            attempt += 1
+            if attempt == max_retry:
+                raise
+    raise Exception("Max retry exceeded")
+
+
+async def _run_single_agent_iteration(
+    ctx: AnyContext,
+    agent: Agent,
+    user_prompt: str,
+    history_list: ListOfDict,
+    rate_limitter: LLMRateLimiter | None = None,
+) -> AgentRun:
     from openai import APIError
     from pydantic_ai.messages import ModelMessagesTypeAdapter
 
+    agent_payload = estimate_request_payload(agent, user_prompt, history_list)
     if rate_limitter:
-        await rate_limitter.throttle(user_prompt)
+        await rate_limitter.throttle(agent_payload)
     else:
-        await llm_rate_limitter.throttle(user_prompt)
+        await llm_rate_limitter.throttle(agent_payload)
 
     async with agent.run_mcp_servers():
         async with agent.iter(
@@ -159,6 +189,19 @@ async def run_agent_iteration(
             return agent_run
 
 
+def estimate_request_payload(
+    agent: Agent, user_prompt: str, history_list: ListOfDict
+) -> str:
+    system_prompts = agent._system_prompts if hasattr(agent, "_system_prompts") else ()
+    return json.dumps(
+        [
+            {"role": "system", "content": "\n".join(system_prompts)},
+            *history_list,
+            {"role": "user", "content": user_prompt},
+        ]
+    )
+
+
 def _get_plain_printer(ctx: AnyContext):
     def printer(*args, **kwargs):
         if "plain" not in kwargs:

{zrb-1.8.8.dist-info → zrb-1.8.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.8.8
+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.8.dist-info → zrb-1.8.10.dist-info}/RECORD

@@ -17,10 +17,10 @@ zrb/builtin/llm/previous-session.js,sha256=xMKZvJoAbrwiyHS0OoPrWuaKxWYLoyR5sgueP
 zrb/builtin/llm/tool/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/builtin/llm/tool/api.py,sha256=yR9I0ZsI96OeQl9pgwORMASVuXsAL0a89D_iPS4C8Dc,1699
 zrb/builtin/llm/tool/cli.py,sha256=_CNEmEc6K2Z0i9ppYeM7jGpqaEdT3uxaWQatmxP3jKE,858
-zrb/builtin/llm/tool/code.py,sha256=7fozry3yKJ2793W4BQkiGnH4OgU0FCe87kNRG05DJZ4,7428
-zrb/builtin/llm/tool/file.py,sha256=hgm2s5zjGu2lo05Y0XisBu9HeaPPs2FyzwaoHBSFAhI,18355
+zrb/builtin/llm/tool/code.py,sha256=gvRnimUh5kWqmpiYtJvEm6KDZhQArqhwAAkKI1_UapY,8133
+zrb/builtin/llm/tool/file.py,sha256=ufLCAaHB0JkEAqQS4fbM9OaTfLluqlCuSyMmnYhI0rY,18491
 zrb/builtin/llm/tool/rag.py,sha256=yqx7vXXyrOCJjhQJl4s0TnLL-2uQUTuKRnkWlSQBW0M,7883
-zrb/builtin/llm/tool/sub_agent.py,sha256=_ItDE5MV_RZtnY_-IUsSMmm6mYaDY3YRINT0hVNsGkA,4702
+zrb/builtin/llm/tool/sub_agent.py,sha256=GPHD8hLlIfme0h1Q0zzMUuAc2HiKl8CRqWGNcgE_H1Q,4764
 zrb/builtin/llm/tool/web.py,sha256=pXRLhcB_Y6z-2w4C4WezH8n-pg3PSMgt_bwn3aaqi6g,5479
 zrb/builtin/md5.py,sha256=690RV2LbW7wQeTFxY-lmmqTSVEEZv3XZbjEUW1Q3XpE,1480
 zrb/builtin/project/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -170,7 +170,7 @@ zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/module/gateway/view
 zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/module/gateway/view/static/pico-css/pico.yellow.min.css,sha256=_UXLKrhEsXonQ-VthBNB7zHUEcV67KDAE-SiDR1SrlU,83371
 zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/module/gateway/view/static/pico-css/pico.zinc.min.css,sha256=C9KHUa2PomYXZg2-rbpDPDYkuL_ZZTwFS-uj1Zo7azE,83337
 zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/module/gateway/view/template/default.html,sha256=Lg4vONCLOx8PSORFitg8JZa-4dp-pyUvKSVzCJEFsB8,4048
-zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/requirements.txt,sha256=jSLXX-ZOElSG69qjx5yVcpg6X_QJoisNDUYpCv2ldH8,195
+zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/requirements.txt,sha256=cV16jm4MpHK020eLm9ityrVBEZgJw0WvwrphOA-2lAE,195
 zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/schema/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/schema/permission.py,sha256=q66LXdZ-QTb30F1VTXNLnjyYBlK_ThVLiHgavtJy4LY,1424
 zrb/builtin/project/add/fastapp/fastapp_template/my_app_name/schema/role.py,sha256=7USbuhHhPc3xXkmwiqTVKsN8-eFWS8Q7emKxCGNGPw0,3244
@@ -217,7 +217,7 @@ zrb/callback/callback.py,sha256=PFhCqzfxdk6IAthmXcZ13DokT62xtBzJr_ciLw6I8Zg,4030
 zrb/cmd/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/cmd/cmd_result.py,sha256=L8bQJzWCpcYexIxHBNsXj2pT3BtLmWex0iJSMkvimOA,597
 zrb/cmd/cmd_val.py,sha256=7Doowyg6BK3ISSGBLt-PmlhzaEkBjWWm51cED6fAUOQ,1014
-zrb/config.py,sha256=XJCJl3u6fo9OlL8V1HNLn5iKcSJjIICQk87CkBrA7IU,10222
+zrb/config.py,sha256=qFtVVme30fMyi5x_mgvvULczNbORqK8ZEN8agXokXO4,10222
 zrb/content_transformer/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/content_transformer/any_content_transformer.py,sha256=v8ZUbcix1GGeDQwB6OKX_1TjpY__ksxWVeqibwa_iZA,850
 zrb/content_transformer/content_transformer.py,sha256=STl77wW-I69QaGzCXjvkppngYFLufow8ybPLSyAvlHs,2404
@@ -246,8 +246,8 @@ zrb/input/option_input.py,sha256=TQB82ko5odgzkULEizBZi0e9TIHEbIgvdP0AR3RhA74,213
 zrb/input/password_input.py,sha256=szBojWxSP9QJecgsgA87OIYwQrY2AQ3USIKdDZY6snU,1465
 zrb/input/str_input.py,sha256=NevZHX9rf1g8eMatPyy-kUX3DglrVAQpzvVpKAzf7bA,81
 zrb/input/text_input.py,sha256=6T3MngWdUs0u0ZVs5Dl11w5KS7nN1RkgrIR_zKumzPM,3695
-zrb/llm_config.py,sha256=9Pvm8SIW97C0C03tSZv-6aBsaFz4fiId3FG3znFteUk,13473
-zrb/llm_rate_limitter.py,sha256=RXdtPreMcmoYSE2Ab2StyHH95F0bD2pGmyySXs4gRio,3725
+zrb/llm_config.py,sha256=7xp4mhre3ULSzfyuqinXXMigNOYNcemVEgPIiWTNdMk,16875
+zrb/llm_rate_limitter.py,sha256=rFYJU2ngo1Hk3aSODsgFIpANI95qZydUvZ4WHOmdbHQ,4376
 zrb/runner/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/runner/cli.py,sha256=AbLTNqFy5FuyGQOWOjHZGaBC8e2yuE_Dx1sBdnisR18,6984
 zrb/runner/common_util.py,sha256=JDMcwvQ8cxnv9kQrAoKVLA40Q1omfv-u5_d5MvvwHeE,1373
@@ -337,7 +337,7 @@ zrb/task/base_trigger.py,sha256=WSGcmBcGAZw8EzUXfmCjqJQkz8GEmi1RzogpF6A1V4s,6902
 zrb/task/cmd_task.py,sha256=irGi0txTcsvGhxjfem4_radR4csNXhgtfcxruSF1LFI,10853
 zrb/task/http_check.py,sha256=Gf5rOB2Se2EdizuN9rp65HpGmfZkGc-clIAlHmPVehs,2565
 zrb/task/llm/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-zrb/task/llm/agent.py,sha256=pLldTsXzdaPs07WIw6DbpLpe-y6xIPU5MAPQ0fRA8Bk,5526
+zrb/task/llm/agent.py,sha256=pwGFkRSQ-maH92QaJr2dgfQhEiwFVK5XnyK4OYBJm4c,6782
 zrb/task/llm/config.py,sha256=Gb0lSHCgGXOAr7igkU7k_Ew5Yp_wOTpNQyZrLrtA7oc,3521
 zrb/task/llm/context.py,sha256=U9a8lxa2ikz6my0Sd5vpO763legHrMHyvBjbrqNmv0Y,3838
 zrb/task/llm/context_enrichment.py,sha256=BlW2CjSUsKJT8EZBXYxOE4MEBbRCoO34PlQQdzA-zBM,7201
@@ -390,7 +390,7 @@ zrb/util/string/name.py,sha256=SXEfxJ1-tDOzHqmSV8kvepRVyMqs2XdV_vyoh_9XUu0,1584
 zrb/util/todo.py,sha256=VGISej2KQZERpornK-8X7bysp4JydMrMUTnG8B0-liI,20708
 zrb/xcom/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 zrb/xcom/xcom.py,sha256=o79rxR9wphnShrcIushA0Qt71d_p3ZTxjNf7x9hJB78,1571
-zrb-1.8.8.dist-info/METADATA,sha256=x49O7t73u93ctMp-rTAz-MjelM8ZDyuVZ8vRFSEG54I,9761
-zrb-1.8.8.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-zrb-1.8.8.dist-info/entry_points.txt,sha256=-Pg3ElWPfnaSM-XvXqCxEAa-wfVI6BEgcs386s8C8v8,46
-zrb-1.8.8.dist-info/RECORD,,
+zrb-1.8.10.dist-info/METADATA,sha256=cFdiCmF4a4s7-jxcMZtXxhjG8L4hZO8-YBw_DecGSQQ,10108
+zrb-1.8.10.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+zrb-1.8.10.dist-info/entry_points.txt,sha256=-Pg3ElWPfnaSM-XvXqCxEAa-wfVI6BEgcs386s8C8v8,46
+zrb-1.8.10.dist-info/RECORD,,