zrb 1.21.9__py3-none-any.whl → 1.21.28__py3-none-any.whl

zrb/config/default_prompt/summarization_prompt.md

@@ -1,11 +1,57 @@
-You are a memory management AI. Your only task is to process the provided conversation history and call the `final_result` tool **once**.
+You are a smart memory management AI. Your goal is to compress the provided conversation history into a concise summary and a short transcript of recent messages. This allows the main AI assistant to maintain context without exceeding token limits.
 
-Follow these instructions carefully:
+You will receive a JSON string representing the full conversation history. This JSON contains a list of message objects.
 
-1. **Summarize:** Create a concise narrative summary that integrates the `Past Conversation Summary` with the `Recent Conversation`. **This summary must not be more than two paragraphs.**
-2. **Transcript:** Extract ONLY the last 4 (four) turns of the `Recent Conversation` to serve as the new transcript.
-  * **Do not change or shorten the content of these turns, with one exception:** If a tool call returns a very long output, do not include the full output. Instead, briefly summarize the result of the tool call.
-  * Ensure the timestamp format is `[YYYY-MM-DD HH:MM:SS UTC+Z] Role: Message/Tool name being called`.
-3. **Update Memory:** Call the `final_result` tool with all the information you consolidated.
+Your task is to call the `save_conversation_summary` tool **once** with the following data. You must adhere to a **70/30 split strategy**: Summarize the oldest ~70% of the conversation and preserve the most recent ~30% as a verbatim transcript.
 
-After you have called the tool, your task is complete.
+1. **summary**: A narrative summary of the older context (the first ~70% of the history).
+  * **Length:** Comprehensive but concise.
+  * **Content - YOU MUST USE THESE SECTIONS:**
+    * **[Completed Actions]:** detailed list of files created, modified, or bugs fixed. **Do not omit file paths.**
+    * **[Active Context]:** What is the current high-level goal?
+    * **[Pending Steps]:** What specifically remains to be done?
+    * **[Constraints]:** Key user preferences or technical constraints.
+  * **Critical Logic:**
+    * **Anti-Looping:** If a task is listed in **[Completed Actions]**, do NOT list it in **[Pending Steps]**.
+    * **Context Merging:** If the input history already contains a summary, merge it intelligently. Updates to files supersede older descriptions.
+
+2. **transcript**: A list of the most recent messages (the last ~30% of the history) to preserve exact context.
+  * **Format:** A list of objects with `role`, `time`, and `content`.
+  * **Time Format:** Use "yyyy-mm-ddTHH:MM:SSZ" (e.g., "2023-10-27T10:00:00Z").
+  * **Content Rules:**
+    * **Preserve Verbatim:** Do not summarize user instructions or code in this section. The main AI needs the exact recent commands to function correctly.
+    * **Tool Outputs:** If a tool output in this recent section is huge (e.g., > 100 lines of file content), you may summarize it (e.g., "File content of X read successfully... "), but preserve any error messages or short confirmations exactly.
+
+**Input Structure Hint:**
+The input JSON is a list of Pydantic AI messages.
+- `kind="request"` -> usually User.
+- `kind="response"` -> usually Model.
+- Tool Results -> `part_kind="tool-return"`.
+
+**Example:**
+
+**Input (Abstract Representation of ~6 turns):**
+```json
+[
+  { "role": "user", "content": "Previous Summary: \n[Completed Actions]: Created `src/app.py`.\n[Active Context]: Fixing login bug.\n[Pending Steps]: Verify fix." },
+  { "role": "model", "content": "I see the bug. I will fix `src/app.py` now." },
+  { "role": "tool_call", "content": "write_file('src/app.py', '...fixed code...')" },
+  { "role": "tool_result", "content": "Success" },
+  { "role": "user", "content": "Great. Now add a test for it." },
+  { "role": "model", "content": "Okay, I will create `tests/test_login.py`." }
+]
+```
+
+**Output (Tool Call `save_conversation_summary`):**
+```json
+{
+  "summary": "[Completed Actions]: Created `src/app.py` and fixed login bug in `src/app.py`.\n[Active Context]: Adding tests for login functionality.\n[Pending Steps]: Create `tests/test_login.py`.\n[Constraints]: None.",
+  "transcript": [
+    { "role": "user", "time": "2023-10-27T10:05:00Z", "content": "Great. Now add a test for it." },
+    { "role": "model", "time": "2023-10-27T10:05:05Z", "content": "Okay, I will create `tests/test_login.py`." }
+  ]
+}
+```
+
+**Final Note:**
+The `summary` + `transcript` is the ONLY memory the main AI will have. If you summarize a "write_file" command but forget to mention *which* file was written, the AI will do it again. **Be specific.**

zrb/config/default_prompt/system_prompt.md

@@ -1,4 +1,4 @@
-You are an expert AI agent designed for completing a single request. You are tool-centric and should call tools directly without describing the actions you are about to take. Only communicate to report the final result.
+This is a single request session. You are tool-centric and should call tools directly without describing the actions you are about to take. Only communicate to report the final result.
 
 # Core Principles
 

zrb/config/llm_rate_limitter.py

@@ -7,7 +7,7 @@ from typing import Any, Callable
 from zrb.config.config import CFG
 
 
-class LLMRateLimiter:
+class LLMRateLimitter:
     """
     Helper class to enforce LLM API rate limits and throttling.
     Tracks requests and tokens in a rolling 60-second window.
@@ -129,7 +129,7 @@ class LLMRateLimiter:
     async def throttle(
         self,
         prompt: Any,
-        throttle_notif_callback: Callable | None = None,
+        throttle_notif_callback: Callable[[str], Any] | None = None,
     ):
         now = time.time()
         str_prompt = self._prompt_to_str(prompt)
@@ -142,7 +142,17 @@ class LLMRateLimiter:
         # 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})."
+                (
+                    "Request exceeds max_tokens_per_request "
+                    f"({tokens} > {self.max_tokens_per_request})."
+                )
+            )
+        if tokens > self.max_tokens_per_minute:
+            raise ValueError(
+                (
+                    "Request exceeds max_tokens_per_minute "
+                    f"({tokens} > {self.max_tokens_per_minute})."
+                )
             )
         # Wait if over per-minute request or token limit
         while (
@@ -150,7 +160,16 @@ class LLMRateLimiter:
             or sum(t for _, t in self.token_times) + tokens > self.max_tokens_per_minute
         ):
             if throttle_notif_callback is not None:
-                throttle_notif_callback()
+                if len(self.request_times) >= self.max_requests_per_minute:
+                    rpm = len(self.request_times)
+                    throttle_notif_callback(
+                        f"Max request per minute exceeded: {rpm} of {self.max_requests_per_minute}"
+                    )
+                else:
+                    tpm = sum(t for _, t in self.token_times) + tokens
+                    throttle_notif_callback(
+                        f"Max token per minute exceeded: {tpm} of {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:
@@ -168,4 +187,4 @@ class LLMRateLimiter:
             return f"{prompt}"
 
 
-llm_rate_limitter = LLMRateLimiter()
+llm_rate_limitter = LLMRateLimitter()

zrb/input/option_input.py

@@ -47,9 +47,21 @@ class OptionInput(BaseInput):
         option_str = ", ".join(options)
         if default_value != "":
             prompt_message = f"{prompt_message} ({option_str}) [{default_value}]"
-        value = input(f"{prompt_message}: ")
+        value = self._get_value_from_user_input(shared_ctx, prompt_message, options)
         if value.strip() != "" and value.strip() not in options:
             value = self._prompt_cli_str(shared_ctx)
         if value.strip() == "":
             value = default_value
         return value
+
+    def _get_value_from_user_input(
+        self, shared_ctx: AnySharedContext, prompt_message: str, options: list[str]
+    ) -> str:
+        from prompt_toolkit import PromptSession
+        from prompt_toolkit.completion import WordCompleter
+
+        if shared_ctx.is_tty:
+            reader = PromptSession()
+            option_completer = WordCompleter(options, ignore_case=True)
+            return reader.prompt(f"{prompt_message}: ", completer=option_completer)
+        return input(f"{prompt_message}: ")

zrb/task/llm/agent.py

@@ -1,22 +1,16 @@
 import inspect
-import json
 from collections.abc import Callable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
-from zrb.config.llm_rate_limitter import LLMRateLimiter, llm_rate_limitter
+from zrb.config.llm_rate_limitter import LLMRateLimitter
 from zrb.context.any_context import AnyContext
-from zrb.context.any_shared_context import AnySharedContext
-from zrb.task.llm.error import extract_api_error_details
-from zrb.task.llm.print_node import print_node
+from zrb.task.llm.history_processor import create_summarize_history_processor
 from zrb.task.llm.tool_wrapper import wrap_func, wrap_tool
-from zrb.task.llm.typing import ListOfDict
-from zrb.util.cli.style import stylize_faint
 
 if TYPE_CHECKING:
     from pydantic_ai import Agent, Tool
-    from pydantic_ai.agent import AgentRun
-    from pydantic_ai.messages import UserContent
+    from pydantic_ai._agent_graph import HistoryProcessor
     from pydantic_ai.models import Model
     from pydantic_ai.output import OutputDataT, OutputSpec
     from pydantic_ai.settings import ModelSettings
@@ -28,13 +22,21 @@ if TYPE_CHECKING:
 def create_agent_instance(
     ctx: AnyContext,
     model: "str | Model",
+    rate_limitter: LLMRateLimitter | None = None,
     output_type: "OutputSpec[OutputDataT]" = str,
     system_prompt: str = "",
     model_settings: "ModelSettings | None" = None,
-    tools: "list[ToolOrCallable]" = [],
+    tools: list["ToolOrCallable"] = [],
     toolsets: list["AbstractToolset[None]"] = [],
     retries: int = 3,
     yolo_mode: bool | list[str] | None = None,
+    summarization_model: "Model | str | None" = None,
+    summarization_model_settings: "ModelSettings | None" = None,
+    summarization_system_prompt: str | None = None,
+    summarization_retries: int = 2,
+    summarization_token_threshold: int | None = None,
+    history_processors: list["HistoryProcessor"] | None = None,
+    auto_summarize: bool = True,
 ) -> "Agent[None, Any]":
     """Creates a new Agent instance with configured tools and servers."""
     from pydantic_ai import Agent, RunContext, Tool
@@ -102,6 +104,21 @@ def create_agent_instance(
         ConfirmationWrapperToolset(wrapped=toolset, ctx=ctx, yolo_mode=yolo_mode)
         for toolset in toolsets
     ]
+    # Create History processor with summarizer
+    history_processors = [] if history_processors is None else history_processors
+    if auto_summarize:
+        history_processors += [
+            create_summarize_history_processor(
+                ctx=ctx,
+                system_prompt=system_prompt,
+                rate_limitter=rate_limitter,
+                summarization_model=summarization_model,
+                summarization_model_settings=summarization_model_settings,
+                summarization_system_prompt=summarization_system_prompt,
+                summarization_token_threshold=summarization_token_threshold,
+                summarization_retries=summarization_retries,
+            )
+        ]
     # Return Agent
     return Agent[None, Any](
         model=model,
@@ -111,12 +128,14 @@ def create_agent_instance(
         toolsets=wrapped_toolsets,
         model_settings=model_settings,
         retries=retries,
+        history_processors=history_processors,
     )
 
 
 def get_agent(
     ctx: AnyContext,
     model: "str | Model",
+    rate_limitter: LLMRateLimitter | None = None,
     output_type: "OutputSpec[OutputDataT]" = str,
     system_prompt: str = "",
     model_settings: "ModelSettings | None" = None,
@@ -128,6 +147,12 @@ def get_agent(
     additional_toolsets: "list[AbstractToolset[None] | str]" = [],
     retries: int = 3,
     yolo_mode: bool | list[str] | None = None,
+    summarization_model: "Model | str | None" = None,
+    summarization_model_settings: "ModelSettings | None" = None,
+    summarization_system_prompt: str | None = None,
+    summarization_retries: int = 2,
+    summarization_token_threshold: int | None = None,
+    history_processors: list["HistoryProcessor"] | None = None,
 ) -> "Agent":
     """Retrieves the configured Agent instance or creates one if necessary."""
     # Get tools for agent
@@ -143,6 +168,7 @@ def get_agent(
     return create_agent_instance(
         ctx=ctx,
         model=model,
+        rate_limitter=rate_limitter,
         output_type=output_type,
         system_prompt=system_prompt,
         tools=tools,
@@ -150,6 +176,12 @@ def get_agent(
         model_settings=model_settings,
         retries=retries,
         yolo_mode=yolo_mode,
+        summarization_model=summarization_model,
+        summarization_model_settings=summarization_model_settings,
+        summarization_system_prompt=summarization_system_prompt,
+        summarization_retries=summarization_retries,
+        summarization_token_threshold=summarization_token_threshold,
+        history_processors=history_processors,
     )
 
 
@@ -170,137 +202,3 @@ def _render_toolset_or_str_list(
             continue
         toolsets.append(toolset_or_str)
     return toolsets
-
-
-async def run_agent_iteration(
-    ctx: AnyContext,
-    agent: "Agent[None, Any]",
-    user_prompt: str,
-    attachments: "list[UserContent] | None" = None,
-    history_list: ListOfDict | None = None,
-    rate_limitter: LLMRateLimiter | None = None,
-    max_retry: int = 2,
-    log_indent_level: int = 0,
-) -> "AgentRun":
-    """
-    Runs a single iteration of the agent execution loop.
-
-    Args:
-        ctx: The task context.
-        agent: The Pydantic AI agent instance.
-        user_prompt: The user's input prompt.
-        history_list: The current conversation history.
-
-    Returns:
-        The agent run result object.
-
-    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,
-                attachments=[] if attachments is None else attachments,
-                history_list=[] if history_list is None else history_list,
-                rate_limitter=(
-                    llm_rate_limitter if rate_limitter is None else rate_limitter
-                ),
-                log_indent_level=log_indent_level,
-            )
-        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,
-    attachments: "list[UserContent]",
-    history_list: ListOfDict,
-    rate_limitter: LLMRateLimiter,
-    log_indent_level: int,
-) -> "AgentRun":
-    from openai import APIError
-    from pydantic_ai.messages import ModelMessagesTypeAdapter
-
-    agent_payload = _estimate_request_payload(
-        agent, user_prompt, attachments, history_list
-    )
-    callback = _create_print_throttle_notif(ctx)
-    if rate_limitter:
-        await rate_limitter.throttle(agent_payload, callback)
-    else:
-        await llm_rate_limitter.throttle(agent_payload, callback)
-
-    user_prompt_with_attachments = [user_prompt] + attachments
-    async with agent:
-        async with agent.iter(
-            user_prompt=user_prompt_with_attachments,
-            message_history=ModelMessagesTypeAdapter.validate_python(history_list),
-        ) as agent_run:
-            async for node in agent_run:
-                # Each node represents a step in the agent's execution
-                try:
-                    await print_node(
-                        _get_plain_printer(ctx), agent_run, node, log_indent_level
-                    )
-                except APIError as e:
-                    # Extract detailed error information from the response
-                    error_details = extract_api_error_details(e)
-                    ctx.log_error(f"API Error: {error_details}")
-                    raise
-                except Exception as e:
-                    ctx.log_error(f"Error processing node: {str(e)}")
-                    ctx.log_error(f"Error type: {type(e).__name__}")
-                    raise
-            return agent_run
-
-
-def _create_print_throttle_notif(ctx: AnyContext) -> Callable[[], None]:
-    def _print_throttle_notif():
-        ctx.print(stylize_faint("  ⌛>> Request Throttled"), plain=True)
-
-    return _print_throttle_notif
-
-
-def _estimate_request_payload(
-    agent: "Agent",
-    user_prompt: str,
-    attachments: "list[UserContent]",
-    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},
-            *[_estimate_attachment_payload(attachment) for attachment in attachments],
-        ]
-    )
-
-
-def _estimate_attachment_payload(attachment: "UserContent") -> Any:
-    if hasattr(attachment, "url"):
-        return {"role": "user", "content": attachment.url}
-    if hasattr(attachment, "data"):
-        return {"role": "user", "content": "x" * len(attachment.data)}
-    return ""
-
-
-def _get_plain_printer(ctx: AnyContext):
-    def printer(*args, **kwargs):
-        if "plain" not in kwargs:
-            kwargs["plain"] = True
-        return ctx.print(*args, **kwargs)
-
-    return printer

zrb/task/llm/agent_runner.py

@@ -0,0 +1,152 @@
+import json
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from zrb.config.llm_rate_limitter import LLMRateLimitter, llm_rate_limitter
+from zrb.context.any_context import AnyContext
+from zrb.task.llm.error import extract_api_error_details
+from zrb.task.llm.print_node import print_node
+from zrb.task.llm.typing import ListOfDict
+from zrb.util.cli.style import stylize_faint
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent, Tool
+    from pydantic_ai.agent import AgentRun
+    from pydantic_ai.messages import UserContent
+
+    ToolOrCallable = Tool | Callable
+
+
+async def run_agent_iteration(
+    ctx: AnyContext,
+    agent: "Agent[None, Any]",
+    user_prompt: str,
+    attachments: "list[UserContent] | None" = None,
+    history_list: ListOfDict | None = None,
+    rate_limitter: LLMRateLimitter | None = None,
+    max_retry: int = 2,
+    log_indent_level: int = 0,
+) -> "AgentRun":
+    """
+    Runs a single iteration of the agent execution loop.
+
+    Args:
+        ctx: The task context.
+        agent: The Pydantic AI agent instance.
+        user_prompt: The user's input prompt.
+        history_list: The current conversation history.
+
+    Returns:
+        The agent run result object.
+
+    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,
+                attachments=[] if attachments is None else attachments,
+                history_list=[] if history_list is None else history_list,
+                rate_limitter=(
+                    llm_rate_limitter if rate_limitter is None else rate_limitter
+                ),
+                log_indent_level=log_indent_level,
+            )
+        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,
+    attachments: "list[UserContent]",
+    history_list: ListOfDict,
+    rate_limitter: LLMRateLimitter,
+    log_indent_level: int,
+) -> "AgentRun":
+    from openai import APIError
+    from pydantic_ai import UsageLimits
+    from pydantic_ai.messages import ModelMessagesTypeAdapter
+
+    agent_payload = _estimate_request_payload(
+        agent, user_prompt, attachments, history_list
+    )
+    callback = _create_print_throttle_notif(ctx)
+    if rate_limitter:
+        await rate_limitter.throttle(agent_payload, callback)
+    else:
+        await llm_rate_limitter.throttle(agent_payload, callback)
+    user_prompt_with_attachments = [user_prompt] + attachments
+    async with agent:
+        async with agent.iter(
+            user_prompt=user_prompt_with_attachments,
+            message_history=ModelMessagesTypeAdapter.validate_python(history_list),
+            usage_limits=UsageLimits(request_limit=None),  # We don't want limit
+        ) as agent_run:
+            async for node in agent_run:
+                # Each node represents a step in the agent's execution
+                try:
+                    await print_node(
+                        _get_plain_printer(ctx), agent_run, node, log_indent_level
+                    )
+                except APIError as e:
+                    # Extract detailed error information from the response
+                    error_details = extract_api_error_details(e)
+                    ctx.log_error(f"API Error: {error_details}")
+                    raise
+                except Exception as e:
+                    ctx.log_error(f"Error processing node: {str(e)}")
+                    ctx.log_error(f"Error type: {type(e).__name__}")
+                    raise
+            return agent_run
+
+
+def _create_print_throttle_notif(ctx: AnyContext) -> Callable[[str], None]:
+    def _print_throttle_notif(reason: str):
+        ctx.print(stylize_faint(f"  ⌛>> Request Throttled: {reason}"), plain=True)
+
+    return _print_throttle_notif
+
+
+def _estimate_request_payload(
+    agent: "Agent",
+    user_prompt: str,
+    attachments: "list[UserContent]",
+    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},
+            *[_estimate_attachment_payload(attachment) for attachment in attachments],
+        ]
+    )
+
+
+def _estimate_attachment_payload(attachment: "UserContent") -> Any:
+    if hasattr(attachment, "url"):
+        return {"role": "user", "content": attachment.url}
+    if hasattr(attachment, "data"):
+        return {"role": "user", "content": "x" * len(attachment.data)}
+    return ""
+
+
+def _get_plain_printer(ctx: AnyContext):
+    def printer(*args, **kwargs):
+        if "plain" not in kwargs:
+            kwargs["plain"] = True
+        return ctx.print(*args, **kwargs)
+
+    return printer

zrb/task/llm/config.py

@@ -4,7 +4,7 @@ if TYPE_CHECKING:
     from pydantic_ai.models import Model
     from pydantic_ai.settings import ModelSettings
 
-from zrb.attr.type import BoolAttr, StrAttr, StrListAttr, fstring
+from zrb.attr.type import BoolAttr, StrAttr, StrListAttr
 from zrb.config.llm_config import LLMConfig, llm_config
 from zrb.context.any_context import AnyContext
 from zrb.util.attr import get_attr, get_bool_attr, get_str_list_attr
@@ -12,7 +12,9 @@ from zrb.util.attr import get_attr, get_bool_attr, get_str_list_attr
 
 def get_yolo_mode(
     ctx: AnyContext,
-    yolo_mode_attr: BoolAttr | StrListAttr | None = None,
+    yolo_mode_attr: (
+        Callable[[AnyContext], list[str] | bool | None] | StrListAttr | BoolAttr | None
+    ) = None,
     render_yolo_mode: bool = True,
 ) -> bool | list[str]:
     if yolo_mode_attr is None:
@@ -77,11 +79,11 @@ def get_model_api_key(
 
 def get_model(
     ctx: AnyContext,
-    model_attr: "Callable[[AnyContext], Model | str | fstring] | Model | str | None",
+    model_attr: "Callable[[AnyContext], Model | str | None] | Model | str | None",
     render_model: bool,
-    model_base_url_attr: StrAttr | None = None,
+    model_base_url_attr: "Callable[[AnyContext], Model | str | None] | Model | str | None",
     render_model_base_url: bool = True,
-    model_api_key_attr: StrAttr | None = None,
+    model_api_key_attr: "Callable[[AnyContext], Model | str | None] | Model | str | None" = None,
     render_model_api_key: bool = True,
     is_small_model: bool = False,
 ) -> "str | Model":