zrb 1.21.17__py3-none-any.whl → 1.21.33__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,136 +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/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