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

zrb/task/base/execution.py

@@ -53,7 +53,9 @@ def check_execute_condition(task: "BaseTask", session: AnySession) -> bool:
     Evaluates the task's execute_condition attribute.
     """
     ctx = task.get_ctx(session)
-    execute_condition_attr = getattr(task, "_execute_condition", True)
+    execute_condition_attr = (
+        task._execute_condition if task._execute_condition is not None else True
+    )
     return get_bool_attr(ctx, execute_condition_attr, True, auto_render=True)
 
 
@@ -63,8 +65,12 @@ async def execute_action_until_ready(task: "BaseTask", session: AnySession):
     """
     ctx = task.get_ctx(session)
     readiness_checks = task.readiness_checks
-    readiness_check_delay = getattr(task, "_readiness_check_delay", 0.5)
-    monitor_readiness = getattr(task, "_monitor_readiness", False)
+    readiness_check_delay = (
+        task._readiness_check_delay if task._readiness_check_delay is not None else 0.5
+    )
+    monitor_readiness = (
+        task._monitor_readiness if task._monitor_readiness is not None else False
+    )
 
     if not readiness_checks:  # Simplified check for empty list
         ctx.log_info("No readiness checks")
@@ -140,8 +146,8 @@ async def execute_action_with_retry(task: "BaseTask", session: AnySession) -> An
     handling success (triggering successors) and failure (triggering fallbacks).
     """
     ctx = task.get_ctx(session)
-    retries = getattr(task, "_retries", 2)
-    retry_period = getattr(task, "_retry_period", 0)
+    retries = task._retries if task._retries is not None else 2
+    retry_period = task._retry_period if task._retry_period is not None else 0
     max_attempt = retries + 1
     ctx.set_max_attempt(max_attempt)
 
@@ -163,8 +169,9 @@ async def execute_action_with_retry(task: "BaseTask", session: AnySession) -> An
             session.get_task_status(task).mark_as_completed()
 
             # Store result in XCom
-            task_xcom: Xcom = ctx.xcom.get(task.name)
-            task_xcom.push(result)
+            task_xcom: Xcom | None = ctx.xcom.get(task.name)
+            if task_xcom is not None:
+                task_xcom.push(result)
 
             # Skip fallbacks and execute successors on success
             skip_fallbacks(task, session)
@@ -201,7 +208,7 @@ async def run_default_action(task: "BaseTask", ctx: AnyContext) -> Any:
     This is the default implementation called by BaseTask._exec_action.
     Subclasses like LLMTask override _exec_action with their own logic.
     """
-    action = getattr(task, "_action", None)
+    action = task._action
     if action is None:
         ctx.log_debug("No action defined for this task.")
         return None

zrb/task/base/lifecycle.py

@@ -12,7 +12,8 @@ from zrb.util.run import run_async
 async def run_and_cleanup(
     task: AnyTask,
     session: AnySession | None = None,
-    str_kwargs: dict[str, str] = {},
+    str_kwargs: dict[str, str] | None = None,
+    kwargs: dict[str, Any] | None = None,
 ) -> Any:
     """
     Wrapper for async_run that ensures session termination and cleanup of
@@ -23,7 +24,9 @@ async def run_and_cleanup(
         session = Session(shared_ctx=SharedContext())
 
     # Create the main task execution coroutine
-    main_task_coro = asyncio.create_task(run_task_async(task, session, str_kwargs))
+    main_task_coro = asyncio.create_task(
+        run_task_async(task, session, str_kwargs, kwargs)
+    )
 
     try:
         result = await main_task_coro
@@ -67,7 +70,8 @@ async def run_and_cleanup(
 async def run_task_async(
     task: AnyTask,
     session: AnySession | None = None,
-    str_kwargs: dict[str, str] = {},
+    str_kwargs: dict[str, str] | None = None,
+    kwargs: dict[str, Any] | None = None,
 ) -> Any:
     """
     Asynchronous entry point for running a task (`task.async_run()`).
@@ -77,7 +81,7 @@ async def run_task_async(
         session = Session(shared_ctx=SharedContext())
 
     # Populate shared context with inputs and environment variables
-    fill_shared_context_inputs(task, session.shared_ctx, str_kwargs)
+    fill_shared_context_inputs(session.shared_ctx, task, str_kwargs, kwargs)
     fill_shared_context_envs(session.shared_ctx)  # Inject OS env vars
 
     # Start the execution chain from the root tasks

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

@@ -1,20 +1,18 @@
-import json
+import inspect
 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
 
 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
     from pydantic_ai.toolsets import AbstractToolset
 
@@ -24,19 +22,59 @@ 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]" = [],
-    toolsets: list["AbstractToolset[Agent]"] = [],
+    tools: list["ToolOrCallable"] = [],
+    toolsets: list["AbstractToolset[None]"] = [],
     retries: int = 3,
-    is_yolo_mode: bool | None = None,
-) -> "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 is_yolo_mode is None:
-        is_yolo_mode = False
+    if yolo_mode is None:
+        yolo_mode = False
     # Normalize tools
     tool_list = []
     for tool_or_callable in tools:
@@ -46,7 +84,7 @@ def create_agent_instance(
             tool = tool_or_callable
             tool_list.append(
                 Tool(
-                    function=wrap_func(tool.function, ctx, is_yolo_mode),
+                    function=wrap_func(tool.function, ctx, yolo_mode),
                     takes_ctx=tool.takes_ctx,
                     max_retries=tool.max_retries,
                     name=tool.name,
@@ -60,184 +98,107 @@ def create_agent_instance(
             )
         else:
             # Turn function into tool
-            tool_list.append(wrap_tool(tool_or_callable, ctx, is_yolo_mode))
+            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,
-        toolsets=toolsets,
+        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",
-    system_prompt: str,
-    model_settings: "ModelSettings | None",
+    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]",
-    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,
-    is_yolo_mode: bool | None = None,
+    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 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,
+        rate_limitter=rate_limitter,
+        output_type=output_type,
         system_prompt=system_prompt,
         tools=tools,
-        toolsets=tool_sets,
+        toolsets=toolsets,
         model_settings=model_settings,
         retries=retries,
-        is_yolo_mode=is_yolo_mode,
-    )
-
-
-async def run_agent_iteration(
-    ctx: AnyContext,
-    agent: "Agent",
-    user_prompt: str,
-    attachments: "list[UserContent] | None" = None,
-    history_list: ListOfDict | None = None,
-    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,
-                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
-                ),
-            )
-        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,
-) -> "AgentRun":
-    from openai import APIError
-    from pydantic_ai.messages import ModelMessagesTypeAdapter
-
-    agent_payload = _estimate_request_payload(
-        agent, user_prompt, attachments, history_list
+        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,
     )
-    if rate_limitter:
-        await rate_limitter.throttle(agent_payload)
-    else:
-        await llm_rate_limitter.throttle(agent_payload)
-
-    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
-                # 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,
-    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
+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