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

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,16 +4,17 @@ 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.context.any_shared_context import AnySharedContext
 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:
@@ -69,7 +70,7 @@ def get_model_api_key(
 ) -> str | None:
     """Gets the model API key, rendering if configured."""
     api_key = get_attr(ctx, model_api_key_attr, None, auto_render=render_model_api_key)
-    if api_key is None and llm_config.default_api_key is not None:
+    if api_key is None and llm_config.default_model_api_key is not None:
         return llm_config.default_model_api_key
     if isinstance(api_key, str) or api_key is None:
         return api_key
@@ -78,11 +79,11 @@ def get_model_api_key(
 
 def get_model(
     ctx: AnyContext,
-    model_attr: "Callable[[AnyContext], Model | str | fstring] | Model | 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":

zrb/task/llm/conversation_history.py

@@ -1,7 +1,6 @@
 import json
 import os
 from collections.abc import Callable
-from copy import deepcopy
 from typing import Any
 
 from zrb.attr.type import StrAttr
@@ -13,6 +12,37 @@ from zrb.util.attr import get_str_attr
 from zrb.util.file import read_file, write_file
 from zrb.util.markdown import make_markdown_section
 from zrb.util.run import run_async
+from zrb.xcom.xcom import Xcom
+
+
+def _get_global_subagent_messages_xcom(ctx: AnyContext) -> Xcom:
+    if "_global_subagents" not in ctx.xcom:
+        ctx.xcom["_global_subagents"] = Xcom([{}])
+    if not isinstance(ctx.xcom["_global_subagents"], Xcom):
+        raise ValueError("ctx.xcom._global_subagents must be an Xcom")
+    return ctx.xcom["_global_subagents"]
+
+
+def inject_subagent_history_into_ctx(
+    ctx: AnyContext, conversation_history: ConversationHistory
+):
+    subagent_messages_xcom = _get_global_subagent_messages_xcom(ctx)
+    existing_subagent_history = subagent_messages_xcom.get({})
+    subagent_messages_xcom.set(
+        {**existing_subagent_history, **conversation_history.subagent_history}
+    )
+
+
+def set_ctx_subagent_history(ctx: AnyContext, subagent_name: str, messages: ListOfDict):
+    subagent_messages_xcom = _get_global_subagent_messages_xcom(ctx)
+    subagent_history = subagent_messages_xcom.get({})
+    subagent_history[subagent_name] = messages
+    subagent_messages_xcom.set(subagent_history)
+
+
+def get_subagent_histories_from_ctx(ctx: AnyContext) -> dict[str, ListOfDict]:
+    subagent_messsages_xcom = _get_global_subagent_messages_xcom(ctx)
+    return subagent_messsages_xcom.get({})
 
 
 def inject_conversation_history_notes(conversation_history: ConversationHistory):
@@ -56,7 +86,10 @@ def get_history_file(
 
 async def _read_from_source(
     ctx: AnyContext,
-    reader: Callable[[AnyContext], dict[str, Any] | list | None] | None,
+    reader: (
+        Callable[[AnyContext], ConversationHistory | dict[str, Any] | list | None]
+        | None
+    ),
     file_path: str | None,
 ) -> "ConversationHistory | None":
     # Priority 1: Reader function
@@ -165,28 +198,6 @@ async def write_conversation_history(
         write_file(history_file, json.dumps(history_data.to_dict(), indent=2))
 
 
-def replace_system_prompt_in_history(
-    history_list: ListOfDict, replacement: str = "<main LLM system prompt>"
-) -> ListOfDict:
-    """
-    Returns a new history list where any part with part_kind 'system-prompt'
-    has its 'content' replaced with the given replacement string.
-    Args:
-        history: List of history items (each item is a dict with a 'parts' list).
-        replacement: The string to use in place of system-prompt content.
-
-    Returns:
-        A deep-copied list of history items with system-prompt content replaced.
-    """
-    new_history = deepcopy(history_list)
-    for item in new_history:
-        parts = item.get("parts", [])
-        for part in parts:
-            if part.get("part_kind") == "system-prompt":
-                part["content"] = replacement
-    return new_history
-
-
 def count_part_in_history_list(history_list: ListOfDict) -> int:
     """Calculates the total number of 'parts' in a history list."""
     history_part_len = 0

zrb/task/llm/conversation_history_model.py

@@ -10,27 +10,24 @@ class ConversationHistory:
 
     def __init__(
         self,
-        past_conversation_summary: str = "",
-        past_conversation_transcript: str = "",
         history: ListOfDict | None = None,
         contextual_note: str | None = None,
         long_term_note: str | None = None,
         project_path: str | None = None,
+        subagent_history: dict[str, ListOfDict] | None = None,
     ):
-        self.past_conversation_transcript = past_conversation_transcript
-        self.past_conversation_summary = past_conversation_summary
         self.history = history if history is not None else []
         self.contextual_note = contextual_note if contextual_note is not None else ""
         self.long_term_note = long_term_note if long_term_note is not None else ""
         self.project_path = project_path if project_path is not None else os.getcwd()
+        self.subagent_history = subagent_history if subagent_history is not None else {}
 
     def to_dict(self) -> dict[str, Any]:
         return {
-            "past_conversation_summary": self.past_conversation_summary,
-            "past_conversation_transcript": self.past_conversation_transcript,
             "history": self.history,
             "contextual_note": self.contextual_note,
             "long_term_note": self.long_term_note,
+            "subagent_history": self.subagent_history,
         }
 
     def model_dump_json(self, indent: int = 2) -> str:
@@ -44,15 +41,11 @@ class ConversationHistory:
             if isinstance(data, cls):
                 return data  # Already a valid instance
             if isinstance(data, dict):
-                # This handles both the new format and the old {'context': ..., 'history': ...}
                 return cls(
-                    past_conversation_summary=data.get("past_conversation_summary", ""),
-                    past_conversation_transcript=data.get(
-                        "past_conversation_transcript", ""
-                    ),
                     history=data.get("history", data.get("messages", [])),
                     contextual_note=data.get("contextual_note", ""),
                     long_term_note=data.get("long_term_note", ""),
+                    subagent_history=data.get("subagent_history", {}),
                 )
             elif isinstance(data, list):
                 # Handle very old format (just a list) - wrap it

zrb/task/llm/default_workflow/coding/workflow.md

@@ -1,5 +1,5 @@
 ---
-description: "A comprehensive workflow for software engineering tasks, including writing, modifying, and debugging code, as well as creating new applications."
+description: "A comprehensive workflow for software engineering tasks, including writing, modifying, and debugging code, as well as creating new applications. ALWAYS activate this workflow whenever you need to deal with software engineering tasks."
 ---
 
 This workflow provides a structured approach to software engineering tasks. Adhere to these guidelines to deliver high-quality, idiomatic code that respects the project's existing patterns and conventions.
@@ -35,8 +35,7 @@ Always consider if a more specific workflow is available and appropriate for the
 When requested to perform tasks like fixing bugs, adding features, refactoring, or explaining code, follow this sequence:
 1. **Understand & Strategize:** Think about the user's request and the relevant codebase context. When the task involves **complex refactoring, codebase exploration or system-wide analysis**, your **first and primary tool** must be 'codebase_investigator'. Use it to build a comprehensive understanding of the code, its structure, and dependencies. For **simple, targeted searches** (like finding a specific function name, file path, or variable declaration), you should use 'search_file_content' or 'glob' directly.
 2. **Plan:** Build a coherent and grounded (based on the understanding in step 1) plan for how you intend to resolve the user's task. Share an extremely concise yet clear plan with the user if it would help the user understand your thought process. As part of the plan, you should use an iterative development process that includes writing unit tests to verify your changes. Use output logs or debug statements as part of this process to arrive at a solution.
-3. **Implement:** Use the available tools (e.g., 'replace', 'write_file' 'run_shell_command' ...) to act on the plan, strictly adhering to the project's established conventions (detailed under 'Core 
-Mandates').
+3. **Implement:** Use the available tools (e.g., 'replace_in_file', 'write_to_file' 'run_shell_command' ...) to act on the plan, strictly adhering to the project's established conventions (detailed under 'Core Mandates').
 4. **Verify (Tests):** If applicable and feasible, verify the changes using the project's testing procedures. Identify the correct test commands and frameworks by examining 'README' files, build/package configuration (e.g., 'package.json'), or existing test execution patterns. NEVER assume standard test commands.
 5. **Verify (Standards):** VERY IMPORTANT: After making code changes, execute the project-specific build, linting and type-checking commands (e.g., 'tsc', 'npm run lint', 'ruff check .') that you have identified for this project (or obtained from the user). This ensures code quality and adherence to standards. If unsure about these commands, you can ask the user if they'd like you to run them and if so how to.
 6. **Finalize:** After all verification passes, consider the task complete. Do not remove or revert any changes or created files (like tests). Await the user's next instruction.