zrb 1.13.1__py3-none-any.whl → 1.21.17__py3-none-any.whl

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/llm/agent.py

@@ -1,5 +1,7 @@
+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
@@ -9,32 +11,68 @@ 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_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.models import Model
+    from pydantic_ai.output import OutputDataT, OutputSpec
     from pydantic_ai.settings import ModelSettings
     from pydantic_ai.toolsets import AbstractToolset
 
     ToolOrCallable = Tool | Callable
-else:
-    ToolOrCallable = Any
 
 
 def create_agent_instance(
     ctx: AnyContext,
-    model: "str | Model | None" = None,
+    model: "str | Model",
+    output_type: "OutputSpec[OutputDataT]" = str,
     system_prompt: str = "",
     model_settings: "ModelSettings | None" = None,
-    tools: list[ToolOrCallable] = [],
-    toolsets: list["AbstractToolset[Agent]"] = [],
+    tools: "list[ToolOrCallable]" = [],
+    toolsets: list["AbstractToolset[None]"] = [],
     retries: int = 3,
-) -> "Agent":
+    yolo_mode: bool | list[str] | None = None,
+) -> "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:
@@ -44,7 +82,7 @@ def create_agent_instance(
             tool = tool_or_callable
             tool_list.append(
                 Tool(
-                    function=wrap_func(tool.function),
+                    function=wrap_func(tool.function, ctx, yolo_mode),
                     takes_ctx=tool.takes_ctx,
                     max_retries=tool.max_retries,
                     name=tool.name,
@@ -58,13 +96,19 @@ def create_agent_instance(
             )
         else:
             # Turn function into tool
-            tool_list.append(wrap_tool(tool_or_callable, ctx))
+            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
+    ]
     # Return Agent
-    return Agent(
+    return Agent[None, Any](
         model=model,
-        system_prompt=system_prompt,
+        output_type=output_type,
+        instructions=system_prompt,
         tools=tool_list,
-        toolsets=toolsets,
+        toolsets=wrapped_toolsets,
         model_settings=model_settings,
         retries=retries,
     )
@@ -72,58 +116,71 @@ def create_agent_instance(
 
 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",
+    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],
-    toolsets_attr: "list[AbstractToolset[Agent]] | Callable[[AnySharedContext], list[AbstractToolset[Agent]]]",  # noqa
-    additional_toolsets: "list[AbstractToolset[Agent]]",
+        "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,
+    yolo_mode: bool | list[str] | 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 Toolsets for agent
-    tool_sets = list(toolsets_attr(ctx) if callable(toolsets_attr) else toolsets_attr)
-    tool_sets.extend(additional_toolsets)
+    toolset_or_str_list = list(
+        toolsets_attr(ctx) if callable(toolsets_attr) else toolsets_attr
+    )
+    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,
+        output_type=output_type,
         system_prompt=system_prompt,
         tools=tools,
-        toolsets=tool_sets,
+        toolsets=toolsets,
         model_settings=model_settings,
         retries=retries,
+        yolo_mode=yolo_mode,
     )
 
 
+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
+
+
 async def run_agent_iteration(
     ctx: AnyContext,
-    agent: "Agent",
+    agent: "Agent[None, Any]",
     user_prompt: str,
-    history_list: ListOfDict,
+    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.
@@ -149,8 +206,12 @@ async def run_agent_iteration(
                 ctx=ctx,
                 agent=agent,
                 user_prompt=user_prompt,
-                history_list=history_list,
-                rate_limitter=rate_limitter,
+                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
@@ -163,28 +224,34 @@ async def _run_single_agent_iteration(
     ctx: AnyContext,
     agent: "Agent",
     user_prompt: str,
+    attachments: "list[UserContent]",
     history_list: ListOfDict,
-    rate_limitter: LLMRateLimiter | None = None,
+    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, history_list)
+    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)
+        await rate_limitter.throttle(agent_payload, callback)
     else:
-        await llm_rate_limitter.throttle(agent_payload)
-
+        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,
+            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
-                # Reference: https://ai.pydantic.dev/agents/#streaming
                 try:
-                    await print_node(_get_plain_printer(ctx), agent_run, node)
+                    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)
@@ -197,8 +264,18 @@ async def _run_single_agent_iteration(
             return agent_run
 
 
-def estimate_request_payload(
-    agent: "Agent", user_prompt: str, history_list: ListOfDict
+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(
@@ -206,10 +283,19 @@ def estimate_request_payload(
             {"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:

zrb/task/llm/config.py

@@ -1,20 +1,43 @@
-from typing import TYPE_CHECKING, Any, Callable
+from typing import TYPE_CHECKING, Callable
 
 if TYPE_CHECKING:
     from pydantic_ai.models import Model
     from pydantic_ai.settings import ModelSettings
 
-from zrb.attr.type import 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
+from zrb.util.attr import get_attr, get_bool_attr, get_str_list_attr
+
+
+def get_yolo_mode(
+    ctx: AnyContext,
+    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:
+        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."""
@@ -47,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
@@ -56,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,
-) -> "str | Model | None":
+    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(
@@ -76,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
@@ -88,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