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

zrb/task/base/monitoring.py

@@ -17,9 +17,13 @@ async def monitor_task_readiness(
     """
     ctx = task.get_ctx(session)
     readiness_checks = task.readiness_checks
-    readiness_check_period = getattr(task, "_readiness_check_period", 5.0)
-    readiness_failure_threshold = getattr(task, "_readiness_failure_threshold", 1)
-    readiness_timeout = getattr(task, "_readiness_timeout", 60)
+    readiness_check_period = (
+        task._readiness_check_period if task._readiness_check_period else 5.0
+    )
+    readiness_failure_threshold = (
+        task._readiness_failure_threshold if task._readiness_failure_threshold else 1
+    )
+    readiness_timeout = task._readiness_timeout if task._readiness_timeout else 60
 
     if not readiness_checks:
         ctx.log_debug("No readiness checks defined, monitoring is not applicable.")
@@ -41,8 +45,9 @@ async def monitor_task_readiness(
                 session.get_task_status(check).reset_history()
                 session.get_task_status(check).reset()
                 # Clear previous XCom data for the check task if needed
-                check_xcom: Xcom = ctx.xcom.get(check.name)
-                check_xcom.clear()
+                check_xcom: Xcom | None = ctx.xcom.get(check.name)
+                if check_xcom is not None:
+                    check_xcom.clear()
 
             readiness_check_coros = [
                 run_async(check.exec_chain(session)) for check in readiness_checks
@@ -77,7 +82,7 @@ async def monitor_task_readiness(
                 )
                 # Ensure check tasks are marked as failed on timeout
                 for check in readiness_checks:
-                    if not session.get_task_status(check).is_finished:
+                    if not session.get_task_status(check).is_ready:
                         session.get_task_status(check).mark_as_failed()
 
             except (asyncio.CancelledError, KeyboardInterrupt):
@@ -92,7 +97,7 @@ async def monitor_task_readiness(
                 )
                 # Mark checks as failed
                 for check in readiness_checks:
-                    if not session.get_task_status(check).is_finished:
+                    if not session.get_task_status(check).is_ready:
                         session.get_task_status(check).mark_as_failed()
 
         # If failure threshold is reached

zrb/task/base_task.py

@@ -21,6 +21,7 @@ from zrb.task.base.execution import (
 )
 from zrb.task.base.lifecycle import execute_root_tasks, run_and_cleanup, run_task_async
 from zrb.task.base.operators import handle_lshift, handle_rshift
+from zrb.util.string.conversion import to_snake_case
 
 
 class BaseTask(AnyTask):
@@ -216,7 +217,10 @@ class BaseTask(AnyTask):
         return build_task_context(self, session)
 
     def run(
-        self, session: AnySession | None = None, str_kwargs: dict[str, str] = {}
+        self,
+        session: AnySession | None = None,
+        str_kwargs: dict[str, str] | None = None,
+        kwargs: dict[str, Any] | None = None,
     ) -> Any:
         """
         Synchronously runs the task and its dependencies, handling async setup and cleanup.
@@ -235,12 +239,19 @@ class BaseTask(AnyTask):
             Any: The final result of the main task execution.
         """
         # Use asyncio.run() to execute the async cleanup wrapper
-        return asyncio.run(run_and_cleanup(self, session, str_kwargs))
+        return asyncio.run(
+            run_and_cleanup(self, session=session, str_kwargs=str_kwargs, kwargs=kwargs)
+        )
 
     async def async_run(
-        self, session: AnySession | None = None, str_kwargs: dict[str, str] = {}
+        self,
+        session: AnySession | None = None,
+        str_kwargs: dict[str, str] | None = None,
+        kwargs: dict[str, Any] | None = None,
     ) -> Any:
-        return await run_task_async(self, session, str_kwargs)
+        return await run_task_async(
+            self, session=session, str_kwargs=str_kwargs, kwargs=kwargs
+        )
 
     async def exec_root_tasks(self, session: AnySession):
         return await execute_root_tasks(self, session)
@@ -276,7 +287,60 @@ class BaseTask(AnyTask):
             # Add definition location to the error
             if hasattr(e, "add_note"):
                 e.add_note(additional_error_note)
-            else:
+            elif hasattr(e, "__notes__"):
                 # fallback: use the __notes__ attribute directly
                 e.__notes__ = getattr(e, "__notes__", []) + [additional_error_note]
             raise e
+
+    def to_function(self) -> Callable[..., Any]:
+        from zrb.context.shared_context import SharedContext
+        from zrb.session.session import Session
+
+        def task_runner_fn(**kwargs) -> Any:
+            task_kwargs = self._get_func_kwargs(kwargs)
+            shared_ctx = SharedContext()
+            session = Session(shared_ctx=shared_ctx)
+            return self.run(session=session, kwargs=task_kwargs)
+
+        task_runner_fn.__doc__ = self._create_fn_docstring()
+        task_runner_fn.__signature__ = self._create_fn_signature()
+        task_runner_fn.__name__ = self.name
+        return task_runner_fn
+
+    def _get_func_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
+        fn_kwargs = {}
+        for inp in self.inputs:
+            snake_input_name = to_snake_case(inp.name)
+            if snake_input_name in kwargs:
+                fn_kwargs[inp.name] = kwargs[snake_input_name]
+        return fn_kwargs
+
+    def _create_fn_docstring(self) -> str:
+        from zrb.context.shared_context import SharedContext
+
+        stub_shared_ctx = SharedContext()
+        str_input_default_values = {}
+        for inp in self.inputs:
+            str_input_default_values[inp.name] = inp.get_default_str(stub_shared_ctx)
+        # Create docstring
+        doc = f"{self.description}\n\n"
+        if len(self.inputs) > 0:
+            doc += "Args:\n"
+            for inp in self.inputs:
+                str_input_default = str_input_default_values.get(inp.name, "")
+                doc += (
+                    f"    {inp.name}: {inp.description} (default: {str_input_default})"
+                )
+                doc += "\n"
+        return doc
+
+    def _create_fn_signature(self) -> inspect.Signature:
+        params = []
+        for inp in self.inputs:
+            params.append(
+                inspect.Parameter(
+                    name=to_snake_case(inp.name),
+                    kind=inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                )
+            )
+        return inspect.Signature(params)

zrb/task/base_trigger.py

@@ -36,7 +36,7 @@ class BaseTrigger(BaseTask):
         input: list[AnyInput | None] | AnyInput | None = None,
         env: list[AnyEnv | None] | AnyEnv | None = None,
         action: fstring | Callable[[AnyContext], Any] | None = None,
-        execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
+        execute_condition: bool | str | Callable[[AnyContext], bool] = True,
         queue_name: fstring | None = None,
         callback: list[AnyCallback] | AnyCallback = [],
         retries: int = 2,
@@ -127,7 +127,7 @@ class BaseTrigger(BaseTask):
         return self._callbacks
 
     async def exec_root_tasks(self, session: AnySession):
-        exchange_xcom = self.get_exchange_xcom(session)
+        exchange_xcom = self._get_exchange_xcom(session)
         exchange_xcom.add_push_callback(lambda: self._exchange_push_callback(session))
         return await super().exec_root_tasks(session)
 
@@ -136,8 +136,7 @@ class BaseTrigger(BaseTask):
         session.defer_coro(coro)
 
     async def _fanout_and_trigger_callback(self, session: AnySession):
-        exchange_xcom = self.get_exchange_xcom(session)
-        data = exchange_xcom.pop()
+        data = self.pop_exchange_xcom(session)
         coros = []
         for callback in self.callbacks:
             xcom_dict = DotDict({self.queue_name: Xcom([data])})
@@ -156,8 +155,16 @@ class BaseTrigger(BaseTask):
             )
         await asyncio.gather(*coros)
 
-    def get_exchange_xcom(self, session: AnySession) -> Xcom:
+    def _get_exchange_xcom(self, session: AnySession) -> Xcom:
         shared_ctx = session.shared_ctx
         if self.queue_name not in shared_ctx.xcom:
             shared_ctx.xcom[self.queue_name] = Xcom()
         return shared_ctx.xcom[self.queue_name]
+
+    def push_exchange_xcom(self, session: AnySession, data: Any):
+        exchange_xcom = self._get_exchange_xcom(session)
+        exchange_xcom.push(data)
+
+    def pop_exchange_xcom(self, session: AnySession) -> Any:
+        exchange_xcom = self._get_exchange_xcom(session)
+        return exchange_xcom.pop()

zrb/task/cmd_task.py

@@ -4,7 +4,7 @@ from functools import partial
 from zrb.attr.type import BoolAttr, IntAttr, StrAttr
 from zrb.cmd.cmd_result import CmdResult
 from zrb.cmd.cmd_val import AnyCmdVal, CmdVal, SingleCmdVal
-from zrb.config import CFG
+from zrb.config.config import CFG
 from zrb.context.any_context import AnyContext
 from zrb.env.any_env import AnyEnv
 from zrb.input.any_input import AnyInput

zrb/task/llm/agent.py

@@ -1,211 +1,204 @@
+import inspect
 from collections.abc import Callable
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
+from zrb.config.llm_rate_limitter import LLMRateLimitter
+from zrb.context.any_context import AnyContext
+from zrb.task.llm.history_processor import create_summarize_history_processor
+from zrb.task.llm.tool_wrapper import wrap_func, wrap_tool
+
 if TYPE_CHECKING:
     from pydantic_ai import Agent, Tool
-    from pydantic_ai.agent import AgentRun
-    from pydantic_ai.mcp import MCPServer
+    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
-else:
-    Agent = Any
-    Tool = Any
-    AgentRun = Any
-    MCPServer = Any
-    ModelMessagesTypeAdapter = Any
-    Model = Any
-    ModelSettings = Any
-
-import json
-
-from zrb.context.any_context import AnyContext
-from zrb.context.any_shared_context import AnySharedContext
-from zrb.llm_rate_limitter import LLMRateLimiter, llm_rate_limitter
-from zrb.task.llm.error import extract_api_error_details
-from zrb.task.llm.print_node import print_node
-from zrb.task.llm.tool_wrapper import wrap_tool
-from zrb.task.llm.typing import ListOfDict
+    from pydantic_ai.toolsets import AbstractToolset
 
-ToolOrCallable = Tool | Callable
+    ToolOrCallable = Tool | Callable
 
 
 def create_agent_instance(
     ctx: AnyContext,
-    model: str | Model | None = None,
+    model: "str | Model",
+    rate_limitter: LLMRateLimitter | None = None,
+    output_type: "OutputSpec[OutputDataT]" = str,
     system_prompt: str = "",
-    model_settings: ModelSettings | None = None,
-    tools: list[ToolOrCallable] = [],
-    mcp_servers: list[MCPServer] = [],
+    model_settings: "ModelSettings | None" = None,
+    tools: list["ToolOrCallable"] = [],
+    toolsets: list["AbstractToolset[None]"] = [],
     retries: int = 3,
-) -> Agent:
+    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, Tool
+    from pydantic_ai import Agent, RunContext, Tool
+    from pydantic_ai.tools import GenerateToolJsonSchema
+    from pydantic_ai.toolsets import ToolsetTool, WrapperToolset
+
+    @dataclass
+    class ConfirmationWrapperToolset(WrapperToolset):
+        ctx: AnyContext
+        yolo_mode: bool | list[str]
+
+        async def call_tool(
+            self, name: str, tool_args: dict, ctx: RunContext, tool: ToolsetTool[None]
+        ) -> Any:
+            # The `tool` object is passed in. Use it for inspection.
+            # Define a temporary function that performs the actual tool call.
+            async def execute_delegated_tool_call(**params):
+                # Pass all arguments down the chain.
+                return await self.wrapped.call_tool(name, tool_args, ctx, tool)
+
+            # For the confirmation UI, make our temporary function look like the real one.
+            try:
+                execute_delegated_tool_call.__name__ = name
+                execute_delegated_tool_call.__doc__ = tool.function.__doc__
+                execute_delegated_tool_call.__signature__ = inspect.signature(
+                    tool.function
+                )
+            except (AttributeError, TypeError):
+                pass  # Ignore if we can't inspect the original function
+            # Use the existing wrap_func to get the confirmation logic
+            wrapped_executor = wrap_func(
+                execute_delegated_tool_call, self.ctx, self.yolo_mode
+            )
+            # Call the wrapped executor. This will trigger the confirmation prompt.
+            return await wrapped_executor(**tool_args)
 
+    if yolo_mode is None:
+        yolo_mode = False
     # Normalize tools
     tool_list = []
     for tool_or_callable in tools:
         if isinstance(tool_or_callable, Tool):
             tool_list.append(tool_or_callable)
+            # Update tool's function
+            tool = tool_or_callable
+            tool_list.append(
+                Tool(
+                    function=wrap_func(tool.function, ctx, yolo_mode),
+                    takes_ctx=tool.takes_ctx,
+                    max_retries=tool.max_retries,
+                    name=tool.name,
+                    description=tool.description,
+                    prepare=tool.prepare,
+                    docstring_format=tool.docstring_format,
+                    require_parameter_descriptions=tool.require_parameter_descriptions,
+                    schema_generator=GenerateToolJsonSchema,
+                    strict=tool.strict,
+                )
+            )
         else:
-            # Pass ctx to wrap_tool
-            tool_list.append(wrap_tool(tool_or_callable, ctx))
+            # Turn function into tool
+            tool_list.append(wrap_tool(tool_or_callable, ctx, yolo_mode))
+    # Wrap toolsets
+    wrapped_toolsets = [
+        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(
+    return Agent[None, Any](
         model=model,
-        system_prompt=system_prompt,
+        output_type=output_type,
+        instructions=system_prompt,
         tools=tool_list,
-        mcp_servers=mcp_servers,
+        toolsets=wrapped_toolsets,
         model_settings=model_settings,
         retries=retries,
+        history_processors=history_processors,
     )
 
 
 def get_agent(
     ctx: AnyContext,
-    agent_attr: Agent | Callable[[AnySharedContext], Agent] | None,
-    model: str | Model | None,
-    system_prompt: str,
-    model_settings: ModelSettings | None,
+    model: "str | Model",
+    rate_limitter: LLMRateLimitter | None = None,
+    output_type: "OutputSpec[OutputDataT]" = str,
+    system_prompt: str = "",
+    model_settings: "ModelSettings | None" = None,
     tools_attr: (
-        list[ToolOrCallable] | Callable[[AnySharedContext], list[ToolOrCallable]]
-    ),
-    additional_tools: list[ToolOrCallable],
-    mcp_servers_attr: list[MCPServer] | Callable[[AnySharedContext], list[MCPServer]],
-    additional_mcp_servers: list[MCPServer],
+        "list[ToolOrCallable] | Callable[[AnyContext], list[ToolOrCallable]]"
+    ) = [],
+    additional_tools: "list[ToolOrCallable]" = [],
+    toolsets_attr: "list[AbstractToolset[None] | str] | Callable[[AnyContext], list[AbstractToolset[None] | str]]" = [],  # noqa
+    additional_toolsets: "list[AbstractToolset[None] | str]" = [],
     retries: int = 3,
-) -> Agent:
+    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."""
-    from pydantic_ai import Agent
-
-    # Render agent instance and return if agent_attr is already an agent
-    if isinstance(agent_attr, Agent):
-        return agent_attr
-    if callable(agent_attr):
-        agent_instance = agent_attr(ctx)
-        if not isinstance(agent_instance, Agent):
-            err_msg = (
-                "Callable agent factory did not return an Agent instance, "
-                f"got: {type(agent_instance)}"
-            )
-            raise TypeError(err_msg)
-        return agent_instance
     # Get tools for agent
     tools = list(tools_attr(ctx) if callable(tools_attr) else tools_attr)
     tools.extend(additional_tools)
-    # Get MCP Servers for agent
-    mcp_servers = list(
-        mcp_servers_attr(ctx) if callable(mcp_servers_attr) else mcp_servers_attr
+    # Get Toolsets for agent
+    toolset_or_str_list = list(
+        toolsets_attr(ctx) if callable(toolsets_attr) else toolsets_attr
     )
-    mcp_servers.extend(additional_mcp_servers)
+    toolset_or_str_list.extend(additional_toolsets)
+    toolsets = _render_toolset_or_str_list(ctx, toolset_or_str_list)
     # If no agent provided, create one using the configuration
     return create_agent_instance(
         ctx=ctx,
         model=model,
+        rate_limitter=rate_limitter,
+        output_type=output_type,
         system_prompt=system_prompt,
         tools=tools,
-        mcp_servers=mcp_servers,
+        toolsets=toolsets,
         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,
     )
 
 
-async def run_agent_iteration(
-    ctx: AnyContext,
-    agent: Agent,
-    user_prompt: str,
-    history_list: ListOfDict,
-    rate_limitter: LLMRateLimiter | None = None,
-    max_retry: int = 2,
-) -> AgentRun:
-    """
-    Runs a single iteration of the agent execution loop.
-
-    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,
-                history_list=history_list,
-                rate_limitter=rate_limitter,
-            )
-        except BaseException:
-            attempt += 1
-            if attempt == max_retry:
-                raise
-    raise Exception("Max retry exceeded")
-
-
-async def _run_single_agent_iteration(
-    ctx: AnyContext,
-    agent: Agent,
-    user_prompt: str,
-    history_list: ListOfDict,
-    rate_limitter: LLMRateLimiter | None = None,
-) -> AgentRun:
-    from openai import APIError
-    from pydantic_ai.messages import ModelMessagesTypeAdapter
-
-    agent_payload = estimate_request_payload(agent, user_prompt, history_list)
-    if rate_limitter:
-        await rate_limitter.throttle(agent_payload)
-    else:
-        await llm_rate_limitter.throttle(agent_payload)
-
-    async with agent.run_mcp_servers():
-        async with agent.iter(
-            user_prompt=user_prompt,
-            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
-                # Reference: https://ai.pydantic.dev/agents/#streaming
-                try:
-                    await print_node(_get_plain_printer(ctx), agent_run, node)
-                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 estimate_request_payload(
-    agent: Agent, user_prompt: str, history_list: ListOfDict
-) -> str:
-    system_prompts = agent._system_prompts if hasattr(agent, "_system_prompts") else ()
-    return json.dumps(
-        [
-            {"role": "system", "content": "\n".join(system_prompts)},
-            *history_list,
-            {"role": "user", "content": user_prompt},
-        ]
-    )
-
-
-def _get_plain_printer(ctx: AnyContext):
-    def printer(*args, **kwargs):
-        if "plain" not in kwargs:
-            kwargs["plain"] = True
-        return ctx.print(*args, **kwargs)
-
-    return printer
+def _render_toolset_or_str_list(
+    ctx: AnyContext, toolset_or_str_list: list["AbstractToolset[None] | str"]
+) -> list["AbstractToolset[None]"]:
+    from pydantic_ai.mcp import load_mcp_servers
+
+    toolsets = []
+    for toolset_or_str in toolset_or_str_list:
+        if isinstance(toolset_or_str, str):
+            try:
+                servers = load_mcp_servers(toolset_or_str)
+                for server in servers:
+                    toolsets.append(server)
+            except Exception as e:
+                ctx.log_error(f"Invalid MCP Config {toolset_or_str}: {e}")
+            continue
+        toolsets.append(toolset_or_str)
+    return toolsets