zrb 1.15.3__py3-none-any.whl → 1.21.29__py3-none-any.whl

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,30 +4,40 @@ if TYPE_CHECKING:
     from pydantic_ai.models import Model
     from pydantic_ai.settings import ModelSettings
 
-from zrb.attr.type import BoolAttr, StrAttr, 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
+from zrb.util.attr import get_attr, get_bool_attr, get_str_list_attr
 
 
-def get_is_yolo_mode(
+def get_yolo_mode(
     ctx: AnyContext,
-    is_yolo_mode_attr: BoolAttr | None = None,
+    yolo_mode_attr: (
+        Callable[[AnyContext], list[str] | bool | None] | StrListAttr | BoolAttr | None
+    ) = None,
     render_yolo_mode: bool = True,
-):
-    return get_bool_attr(
-        ctx,
-        is_yolo_mode_attr,
-        llm_config.default_yolo_mode,
-        auto_render=render_yolo_mode,
-    )
+) -> bool | list[str]:
+    if yolo_mode_attr is None:
+        return llm_config.default_yolo_mode
+    try:
+        return get_bool_attr(
+            ctx,
+            yolo_mode_attr,
+            False,
+            auto_render=render_yolo_mode,
+        )
+    except Exception:
+        return get_str_list_attr(
+            ctx,
+            yolo_mode_attr,
+            auto_render=render_yolo_mode,
+        )
 
 
 def get_model_settings(
     ctx: AnyContext,
     model_settings_attr: (
-        "ModelSettings | Callable[[AnySharedContext], ModelSettings] | None"
+        "ModelSettings | Callable[[AnyContext], ModelSettings] | None"
     ) = None,
 ) -> "ModelSettings | None":
     """Gets the model settings, resolving callables if necessary."""
@@ -60,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
@@ -69,18 +79,21 @@ def get_model_api_key(
 
 def get_model(
     ctx: AnyContext,
-    model_attr: "Callable[[AnySharedContext], 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":
     """Gets the model instance or name, handling defaults and configuration."""
     from pydantic_ai.models import Model
 
     model = get_attr(ctx, model_attr, None, auto_render=render_model)
     if model is None:
+        if is_small_model:
+            return llm_config.default_small_model
         return llm_config.default_model
     if isinstance(model, str):
         model_base_url = get_model_base_url(
@@ -89,11 +102,11 @@ def get_model(
         model_api_key = get_model_api_key(ctx, model_api_key_attr, render_model_api_key)
         new_llm_config = LLMConfig(
             default_model_name=model,
-            default_base_url=model_base_url,
-            default_api_key=model_api_key,
+            default_model_base_url=model_base_url,
+            default_model_api_key=model_api_key,
         )
         if model_base_url is None and model_api_key is None:
-            default_model_provider = llm_config.default_model_provider
+            default_model_provider = _get_default_model_provider(is_small_model)
             if default_model_provider is not None:
                 new_llm_config.set_default_model_provider(default_model_provider)
         return new_llm_config.default_model
@@ -101,3 +114,9 @@ def get_model(
     if isinstance(model, Model):
         return model
     raise ValueError(f"Invalid model type resolved: {type(model)}, value: {model}")
+
+
+def _get_default_model_provider(is_small_model: bool = False):
+    if is_small_model:
+        return llm_config.default_small_model_provider
+    return llm_config.default_model_provider

zrb/task/llm/conversation_history.py

@@ -1,16 +1,73 @@
 import json
+import os
 from collections.abc import Callable
-from copy import deepcopy
 from typing import Any
 
 from zrb.attr.type import StrAttr
+from zrb.config.llm_context.config import llm_context_config
 from zrb.context.any_context import AnyContext
-from zrb.context.any_shared_context import AnySharedContext
 from zrb.task.llm.conversation_history_model import ConversationHistory
 from zrb.task.llm.typing import ListOfDict
 from zrb.util.attr import get_str_attr
-from zrb.util.file import write_file
+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):
+    conversation_history.long_term_note = _fetch_long_term_note(
+        conversation_history.project_path
+    )
+    conversation_history.contextual_note = _fetch_contextual_note(
+        conversation_history.project_path
+    )
+
+
+def _fetch_long_term_note(project_path: str) -> str:
+    contexts = llm_context_config.get_notes(cwd=project_path)
+    return contexts.get("/", "")
+
+
+def _fetch_contextual_note(project_path: str) -> str:
+    contexts = llm_context_config.get_notes(cwd=project_path)
+    return "\n".join(
+        [
+            make_markdown_section(header, content)
+            for header, content in contexts.items()
+            if header != "/"
+        ]
+    )
 
 
 def get_history_file(
@@ -27,16 +84,62 @@ def get_history_file(
     )
 
 
+async def _read_from_source(
+    ctx: AnyContext,
+    reader: (
+        Callable[[AnyContext], ConversationHistory | dict[str, Any] | list | None]
+        | None
+    ),
+    file_path: str | None,
+) -> "ConversationHistory | None":
+    # Priority 1: Reader function
+    if reader:
+        try:
+            raw_data = await run_async(reader(ctx))
+            if raw_data:
+                instance = ConversationHistory.parse_and_validate(
+                    ctx, raw_data, "reader"
+                )
+                if instance:
+                    return instance
+        except Exception as e:
+            ctx.log_warning(
+                f"Error executing conversation history reader: {e}. Ignoring."
+            )
+    # Priority 2: History file
+    if file_path and os.path.isfile(file_path):
+        try:
+            content = read_file(file_path)
+            raw_data = json.loads(content)
+            instance = ConversationHistory.parse_and_validate(
+                ctx, raw_data, f"file '{file_path}'"
+            )
+            if instance:
+                return instance
+        except json.JSONDecodeError:
+            ctx.log_warning(
+                f"Could not decode JSON from history file '{file_path}'. "
+                "Ignoring file content."
+            )
+        except Exception as e:
+            ctx.log_warning(
+                f"Error reading history file '{file_path}': {e}. "
+                "Ignoring file content."
+            )
+    # Fallback: Return default value
+    return None
+
+
 async def read_conversation_history(
     ctx: AnyContext,
     conversation_history_reader: (
-        Callable[[AnySharedContext], ConversationHistory | dict | list | None] | None
+        Callable[[AnyContext], ConversationHistory | dict | list | None] | None
     ),
     conversation_history_file_attr: StrAttr | None,
     render_history_file: bool,
     conversation_history_attr: (
         ConversationHistory
-        | Callable[[AnySharedContext], ConversationHistory | dict | list]
+        | Callable[[AnyContext], ConversationHistory | dict | list]
         | dict
         | list
     ),
@@ -46,7 +149,7 @@ async def read_conversation_history(
         ctx, conversation_history_file_attr, render_history_file
     )
     # Use the class method defined above
-    history_data = await ConversationHistory.read_from_source(
+    history_data = await _read_from_source(
         ctx=ctx,
         reader=conversation_history_reader,
         file_path=history_file,
@@ -80,7 +183,7 @@ async def write_conversation_history(
     ctx: AnyContext,
     history_data: ConversationHistory,
     conversation_history_writer: (
-        Callable[[AnySharedContext, ConversationHistory], None] | None
+        Callable[[AnyContext, ConversationHistory], None] | None
     ),
     conversation_history_file_attr: StrAttr | None,
     render_history_file: bool,
@@ -95,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

@@ -1,90 +1,38 @@
 import json
 import os
-from collections.abc import Callable
 from typing import Any
 
-from zrb.config.llm_context.config import llm_context_config
 from zrb.context.any_context import AnyContext
 from zrb.task.llm.typing import ListOfDict
-from zrb.util.file import read_file
-from zrb.util.llm.prompt import make_prompt_section
-from zrb.util.run import run_async
 
 
 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:
         return json.dumps(self.to_dict(), indent=indent)
 
-    @classmethod
-    async def read_from_source(
-        cls,
-        ctx: AnyContext,
-        reader: Callable[[AnyContext], dict[str, Any] | list | None] | None,
-        file_path: str | None,
-    ) -> "ConversationHistory | None":
-        # Priority 1: Reader function
-        if reader:
-            try:
-                raw_data = await run_async(reader(ctx))
-                if raw_data:
-                    instance = cls.parse_and_validate(ctx, raw_data, "reader")
-                    if instance:
-                        return instance
-            except Exception as e:
-                ctx.log_warning(
-                    f"Error executing conversation history reader: {e}. Ignoring."
-                )
-        # Priority 2: History file
-        if file_path and os.path.isfile(file_path):
-            try:
-                content = read_file(file_path)
-                raw_data = json.loads(content)
-                instance = cls.parse_and_validate(ctx, raw_data, f"file '{file_path}'")
-                if instance:
-                    return instance
-            except json.JSONDecodeError:
-                ctx.log_warning(
-                    f"Could not decode JSON from history file '{file_path}'. "
-                    "Ignoring file content."
-                )
-            except Exception as e:
-                ctx.log_warning(
-                    f"Error reading history file '{file_path}': {e}. "
-                    "Ignoring file content."
-                )
-        # Fallback: Return default value
-        return None
-
-    def fetch_newest_notes(self):
-        self._fetch_long_term_note()
-        self._fetch_contextual_note()
-
     @classmethod
     def parse_and_validate(
         cls, ctx: AnyContext, data: Any, source: str
@@ -93,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
@@ -121,122 +65,3 @@ class ConversationHistory:
                 f"Error validating/parsing history data from {source}: {e}. Ignoring."
             )
         return cls()
-
-    def write_past_conversation_summary(self, past_conversation_summary: str):
-        """
-        Write or update the past conversation summary.
-
-        Use this tool to store or update a summary of previous conversations for
-        future reference. This is useful for providing context to LLMs or other tools
-        that need a concise history.
-
-        Args:
-            past_conversation_summary (str): The summary text to store.
-
-        Returns:
-            str: A JSON object indicating the success or failure of the operation.
-
-        Raises:
-            Exception: If the summary cannot be written.
-        """
-        self.past_conversation_summary = past_conversation_summary
-        return json.dumps({"success": True})
-
-    def write_past_conversation_transcript(self, past_conversation_transcript: str):
-        """
-        Write or update the past conversation transcript.
-
-        Use this tool to store or update the full transcript of previous conversations.
-        This is useful for providing detailed context to LLMs or for record-keeping.
-
-        Args:
-            past_conversation_transcript (str): The transcript text to store.
-
-        Returns:
-            str: A JSON object indicating the success or failure of the operation.
-
-        Raises:
-            Exception: If the transcript cannot be written.
-        """
-        self.past_conversation_transcript = past_conversation_transcript
-        return json.dumps({"success": True})
-
-    def read_long_term_note(self) -> str:
-        """
-        Read the content of the long-term references.
-
-        This tool helps you retrieve knowledge or notes stored for long-term reference.
-        If the note does not exist, you may want to create it using the write tool.
-
-        Returns:
-            str: JSON with content of the notes.
-
-        Raises:
-            Exception: If the note cannot be read.
-        """
-        return json.dumps({"content": self._fetch_long_term_note()})
-
-    def write_long_term_note(self, content: str) -> str:
-        """
-        Write the entire content of the long-term references.
-        This will overwrite any existing long-term notes.
-
-        Args:
-            content (str): The full content of the long-term notes.
-
-        Returns:
-            str: JSON indicating success.
-        """
-        llm_context_config.write_context(content, context_path="/")
-        return json.dumps({"success": True})
-
-    def read_contextual_note(self) -> str:
-        """
-        Read the content of the contextual references for the current project.
-
-        This tool helps you retrieve knowledge or notes stored for contextual reference.
-        If the note does not exist, you may want to create it using the write tool.
-
-        Returns:
-            str: JSON with content of the notes.
-
-        Raises:
-            Exception: If the note cannot be read.
-        """
-        return json.dumps({"content": self._fetch_contextual_note()})
-
-    def write_contextual_note(
-        self, content: str, context_path: str | None = None
-    ) -> str:
-        """
-        Write the entire content of the contextual references for a specific path.
-        This will overwrite any existing contextual notes for that path.
-
-        Args:
-            content (str): The full content of the contextual notes.
-            context_path (str, optional): The directory path for the context.
-                                        Defaults to the current project path.
-
-        Returns:
-            str: JSON indicating success.
-        """
-        if context_path is None:
-            context_path = self.project_path
-        llm_context_config.write_context(content, context_path=context_path)
-        return json.dumps({"success": True})
-
-    def _fetch_long_term_note(self):
-        contexts = llm_context_config.get_contexts(cwd=self.project_path)
-        self.long_term_note = contexts.get("/", "")
-        return self.long_term_note
-
-    def _fetch_contextual_note(self):
-        contexts = llm_context_config.get_contexts(cwd=self.project_path)
-        self.contextual_note = "\n".join(
-            [
-                make_prompt_section(header, content)
-                for header, content in contexts.items()
-                if header != "/"
-            ]
-        )
-        return self.contextual_note