zrb 1.21.37__py3-none-any.whl → 1.21.43__py3-none-any.whl

zrb/builtin/llm/history.py

@@ -17,13 +17,10 @@ def read_chat_conversation(ctx: AnyContext) -> dict[str, Any] | list | None:
         return None  # Indicate no history to load
     previous_session_name = ctx.input.previous_session
     if not previous_session_name:  # Check for empty string or None
-        last_session_file_path = os.path.join(CFG.LLM_HISTORY_DIR, "last-session")
-        if os.path.isfile(last_session_file_path):
-            previous_session_name = read_file(last_session_file_path).strip()
-            if not previous_session_name:  # Handle empty last-session file
-                return None
-        else:
-            return None  # No previous session specified and no last session found
+        last_session_name = get_last_session_name()
+        if last_session_name is None:
+            return None
+        previous_session_name = last_session_name
     conversation_file_path = os.path.join(
         CFG.LLM_HISTORY_DIR, f"{previous_session_name}.json"
     )
@@ -51,6 +48,16 @@ def read_chat_conversation(ctx: AnyContext) -> dict[str, Any] | list | None:
         return None
 
 
+def get_last_session_name() -> str | None:
+    last_session_file_path = os.path.join(CFG.LLM_HISTORY_DIR, "last-session")
+    if not os.path.isfile(last_session_file_path):
+        return None
+    last_session_name = read_file(last_session_file_path).strip()
+    if not last_session_name:  # Handle empty last-session file
+        return None
+    return last_session_name
+
+
 def write_chat_conversation(ctx: AnyContext, history_data: ConversationHistory):
     """Writes the conversation history data (including context) to a session file."""
     os.makedirs(CFG.LLM_HISTORY_DIR, exist_ok=True)

zrb/builtin/llm/llm_ask.py

@@ -31,6 +31,11 @@ from zrb.builtin.llm.tool.web import (
     create_search_internet_tool,
     open_web_page,
 )
+from zrb.builtin.llm.xcom_names import (
+    LLM_ASK_ERROR_XCOM_NAME,
+    LLM_ASK_RESULT_XCOM_NAME,
+    LLM_ASK_SESSION_XCOM_NAME,
+)
 from zrb.callback.callback import Callback
 from zrb.config.config import CFG
 from zrb.config.llm_config import llm_config
@@ -40,6 +45,7 @@ from zrb.input.bool_input import BoolInput
 from zrb.input.str_input import StrInput
 from zrb.input.text_input import TextInput
 from zrb.task.base_trigger import BaseTrigger
+from zrb.task.llm.workflow import LLM_LOADED_WORKFLOW_XCOM_NAME
 from zrb.task.llm_task import LLMTask
 from zrb.util.string.conversion import to_boolean
 
@@ -99,6 +105,8 @@ def _get_default_yolo_mode(ctx: AnyContext) -> str:
 
 
 def _render_yolo_mode_input(ctx: AnyContext) -> list[str] | bool:
+    if isinstance(ctx.input.yolo, bool):
+        return ctx.input.yolo
     if ctx.input.yolo.strip() == "":
         return []
     elements = [element.strip() for element in ctx.input.yolo.split(",")]
@@ -172,9 +180,9 @@ def _get_inputs(require_message: bool = True) -> list[AnyInput | None]:
             always_prompt=False,
         ),
         TextInput(
-            "workflows",
-            description="Workflows",
-            prompt="Workflows",
+            "workflow",
+            description="Workflows (comma separated)",
+            prompt="Workflows (comma separated)",
             default=lambda ctx: ",".join(llm_config.default_workflows),
             allow_positional_parsing=False,
             always_prompt=False,
@@ -237,7 +245,7 @@ llm_ask = LLMTask(
         None if ctx.input.system_prompt.strip() == "" else ctx.input.system_prompt
     ),
     workflows=lambda ctx: (
-        None if ctx.input.workflows.strip() == "" else ctx.input.workflows.split(",")
+        None if ctx.input.workflow.strip() == "" else ctx.input.workflow.split(",")
     ),
     attachment=_render_attach_input,
     message="{ctx.input.message}",
@@ -258,9 +266,10 @@ llm_group.add_task(
         callback=Callback(
             task=llm_ask,
             input_mapping=get_llm_ask_input_mapping,
-            result_queue="ask_result",
-            error_queue="ask_error",
-            session_name_queue="ask_session_name",
+            xcom_mapping={LLM_LOADED_WORKFLOW_XCOM_NAME: LLM_LOADED_WORKFLOW_XCOM_NAME},
+            result_queue=LLM_ASK_RESULT_XCOM_NAME,
+            error_queue=LLM_ASK_ERROR_XCOM_NAME,
+            session_name_queue=LLM_ASK_SESSION_XCOM_NAME,
         ),
         retries=0,
         cli_only=True,

zrb/builtin/llm/tool/file.py

@@ -274,8 +274,9 @@ def write_to_file(
        - CORRECT: "content": "He said \"Hello\""
        - WRONG:   "content": "He said \\"Hello\\""  <-- This breaks JSON parsing!
     2. **SIZE LIMIT:** Content MUST NOT exceed 4000 characters.
-       - Exceeding this causes truncation and EOF errors.
-       - Split larger content into multiple sequential calls (first 'w', then 'a').
+       - **STRICT PROHIBITION:** You are FORBIDDEN from writing more than 4000 characters in a single call.
+       - This is due to LLM output token limits, which will cause truncation and failure.
+       - To write larger files, you MUST split the content into multiple sequential calls (e.g., first 'w', then 'a').
 
     Examples:
     ```

zrb/builtin/llm/tool/search/brave.py

@@ -1,7 +1,5 @@
 from typing import Any
 
-import requests
-
 from zrb.config.config import CFG
 
 
@@ -36,6 +34,8 @@ def search_internet(
     Returns:
         dict: Summary of search results (titles, links, snippets).
     """
+    import requests
+
     if safe_search is None:
         safe_search = CFG.BRAVE_API_SAFE
     if language is None:

zrb/builtin/llm/tool/search/searxng.py

@@ -1,7 +1,5 @@
 from typing import Any
 
-import requests
-
 from zrb.config.config import CFG
 
 
@@ -36,6 +34,8 @@ def search_internet(
     Returns:
         dict: Summary of search results (titles, links, snippets).
     """
+    import requests
+
     if safe_search is None:
         safe_search = CFG.SEARXNG_SAFE
     if language is None:

zrb/builtin/llm/tool/search/serpapi.py

@@ -1,7 +1,5 @@
 from typing import Any
 
-import requests
-
 from zrb.config.config import CFG
 
 
@@ -36,6 +34,8 @@ def search_internet(
     Returns:
         dict: Summary of search results (titles, links, snippets).
     """
+    import requests
+
     if safe_search is None:
         safe_search = CFG.SERPAPI_SAFE
     if language is None:

zrb/builtin/llm/xcom_names.py

@@ -0,0 +1,3 @@
+LLM_ASK_RESULT_XCOM_NAME = "ask_result"
+LLM_ASK_ERROR_XCOM_NAME = "ask_error"
+LLM_ASK_SESSION_XCOM_NAME = "ask_session_name"

zrb/callback/callback.py

@@ -6,7 +6,6 @@ from zrb.callback.any_callback import AnyCallback
 from zrb.session.any_session import AnySession
 from zrb.task.any_task import AnyTask
 from zrb.util.attr import get_str_dict_attr
-from zrb.util.cli.style import stylize_faint
 from zrb.util.string.conversion import to_snake_case
 from zrb.xcom.xcom import Xcom
 
@@ -24,6 +23,7 @@ class Callback(AnyCallback):
         task: AnyTask,
         input_mapping: StrDictAttr,
         render_input_mapping: bool = True,
+        xcom_mapping: dict[str, str] | None = None,
         result_queue: str | None = None,
         error_queue: str | None = None,
         session_name_queue: str | None = None,
@@ -36,6 +36,7 @@ class Callback(AnyCallback):
             input_mapping: A dictionary or attribute mapping to prepare inputs for the task.
             render_input_mapping: Whether to render the input mapping using
                 f-string like syntax.
+            xcom_mapping: Map of parent session's xcom names to current session's xcom names
             result_queue: The name of the XCom queue in the parent session
                 to publish the task result.
             result_queue: The name of the Xcom queue in the parent session
@@ -46,6 +47,7 @@ class Callback(AnyCallback):
         self._task = task
         self._input_mapping = input_mapping
         self._render_input_mapping = render_input_mapping
+        self._xcom_mapping = xcom_mapping
         self._result_queue = result_queue
         self._error_queue = error_queue
         self._session_name_queue = session_name_queue
@@ -63,6 +65,11 @@ class Callback(AnyCallback):
         for name, value in inputs.items():
             session.shared_ctx.input[name] = value
             session.shared_ctx.input[to_snake_case(name)] = value
+        # map xcom
+        if self._xcom_mapping is not None:
+            for parent_xcom_name, current_xcom_name in self._xcom_mapping.items():
+                parent_xcom = parent_session.shared_ctx.xcom[parent_xcom_name]
+                session.shared_ctx.xcom[current_xcom_name] = parent_xcom
         # run task and get result
         try:
             result = await self._task.async_run(session)

zrb/config/config.py

@@ -373,7 +373,7 @@ class Config:
         """
         return int(
             self._getenv(
-                ["LLM_MAX_TOKEN_PER_MINUTE", "LLM_MAX_TOKENS_PER_MINUTE"], "100000"
+                ["LLM_MAX_TOKEN_PER_MINUTE", "LLM_MAX_TOKENS_PER_MINUTE"], "120000"
             )
         )
 

zrb/context/context.py

@@ -139,6 +139,17 @@ class Context(AnyContext):
         stylized_prefix = stylize(prefix, color=color)
         print(f"{stylized_prefix} {message}", sep=sep, end=end, file=file, flush=flush)
 
+    def print_err(
+        self,
+        *values: object,
+        sep: str | None = " ",
+        end: str | None = "\n",
+        file: TextIO | None = sys.stderr,
+        flush: bool = True,
+        plain: bool = False,
+    ):
+        self.print(*values, sep=sep, end=end, file=file, flush=flush, plain=plain)
+
     def log_debug(
         self,
         *values: object,

zrb/task/base/context.py

@@ -79,24 +79,36 @@ def combine_inputs(
             input_names.append(task_input.name)  # Update names list
 
 
+def combine_envs(
+    existing_envs: list[AnyEnv],
+    new_envs: list[AnyEnv | None] | AnyEnv | None,
+):
+    """
+    Combines new envs into an existing list.
+    Modifies the existing_envs list in place.
+    """
+    if isinstance(new_envs, AnyEnv):
+        existing_envs.append(new_envs)
+    elif new_envs is None:
+        pass
+    else:
+        # new_envs is a list
+        for env in new_envs:
+            if env is not None:
+                existing_envs.append(env)
+
+
 def get_combined_envs(task: "BaseTask") -> list[AnyEnv]:
     """
     Aggregates environment variables from the task and its upstreams.
     """
-    envs = []
+    envs: list[AnyEnv] = []
     for upstream in task.upstreams:
-        envs.extend(upstream.envs)  # Use extend for list concatenation
-
-    # Access _envs directly as task is BaseTask
-    task_envs: list[AnyEnv | None] | AnyEnv | None = task._envs
-    if isinstance(task_envs, AnyEnv):
-        envs.append(task_envs)
-    elif isinstance(task_envs, list):
-        # Filter out None while extending
-        envs.extend(env for env in task_envs if env is not None)
-
-    # Filter out None values efficiently from the combined list
-    return [env for env in envs if env is not None]
+        combine_envs(envs, upstream.envs)
+
+    combine_envs(envs, task._envs)
+
+    return envs
 
 
 def get_combined_inputs(task: "BaseTask") -> list[AnyInput]:

zrb/task/base/execution.py

@@ -88,56 +88,61 @@ async def execute_action_until_ready(task: "BaseTask", session: AnySession):
         run_async(execute_action_with_retry(task, session))
     )
 
-    await asyncio.sleep(readiness_check_delay)
-
-    readiness_check_coros = [
-        run_async(check.exec_chain(session)) for check in readiness_checks
-    ]
-
-    # Wait primarily for readiness checks to complete
-    ctx.log_info("Waiting for readiness checks")
-    readiness_passed = False
     try:
-        # Gather results, but primarily interested in completion/errors
-        await asyncio.gather(*readiness_check_coros)
-        # Check if all readiness tasks actually completed successfully
-        all_readiness_completed = all(
-            session.get_task_status(check).is_completed for check in readiness_checks
-        )
-        if all_readiness_completed:
-            ctx.log_info("Readiness checks completed successfully")
-            readiness_passed = True
-            # Mark task as ready only if checks passed and action didn't fail during checks
-            if not session.get_task_status(task).is_failed:
-                ctx.log_info("Marked as ready")
-                session.get_task_status(task).mark_as_ready()
-        else:
-            ctx.log_warning(
-                "One or more readiness checks did not complete successfully."
-            )
+        await asyncio.sleep(readiness_check_delay)
 
-    except Exception as e:
-        ctx.log_error(f"Readiness check failed with exception: {e}")
-        # If readiness checks fail with an exception, the task is not ready.
-        # The action_coro might still be running or have failed.
-        # execute_action_with_retry handles marking the main task status.
-
-    # Defer the main action coroutine; it will be awaited later if needed
-    session.defer_action(task, action_coro)
-
-    # Start monitoring only if readiness passed and monitoring is enabled
-    if readiness_passed and monitor_readiness:
-        # Import dynamically to avoid circular dependency if monitoring imports execution
-        from zrb.task.base.monitoring import monitor_task_readiness
+        readiness_check_coros = [
+            run_async(check.exec_chain(session)) for check in readiness_checks
+        ]
 
-        monitor_coro = asyncio.create_task(
-            run_async(monitor_task_readiness(task, session, action_coro))
-        )
-        session.defer_monitoring(task, monitor_coro)
+        # Wait primarily for readiness checks to complete
+        ctx.log_info("Waiting for readiness checks")
+        readiness_passed = False
+        try:
+            # Gather results, but primarily interested in completion/errors
+            await asyncio.gather(*readiness_check_coros)
+            # Check if all readiness tasks actually completed successfully
+            all_readiness_completed = all(
+                session.get_task_status(check).is_completed
+                for check in readiness_checks
+            )
+            if all_readiness_completed:
+                ctx.log_info("Readiness checks completed successfully")
+                readiness_passed = True
+                # Mark task as ready only if checks passed and action didn't fail during checks
+                if not session.get_task_status(task).is_failed:
+                    ctx.log_info("Marked as ready")
+                    session.get_task_status(task).mark_as_ready()
+            else:
+                ctx.log_warning(
+                    "One or more readiness checks did not complete successfully."
+                )
+
+        except Exception as e:
+            ctx.log_error(f"Readiness check failed with exception: {e}")
+            # If readiness checks fail with an exception, the task is not ready.
+            # The action_coro might still be running or have failed.
+            # execute_action_with_retry handles marking the main task status.
+
+        # Defer the main action coroutine; it will be awaited later if needed
+        session.defer_action(task, action_coro)
+
+        # Start monitoring only if readiness passed and monitoring is enabled
+        if readiness_passed and monitor_readiness:
+            # Import dynamically to avoid circular dependency if monitoring imports execution
+            from zrb.task.base.monitoring import monitor_task_readiness
+
+            monitor_coro = asyncio.create_task(
+                run_async(monitor_task_readiness(task, session, action_coro))
+            )
+            session.defer_monitoring(task, monitor_coro)
 
-    # The result here is primarily about readiness check completion.
-    # The actual task result is handled by the deferred action_coro.
-    return None
+        # The result here is primarily about readiness check completion.
+        # The actual task result is handled by the deferred action_coro.
+        return None
+    except (asyncio.CancelledError, KeyboardInterrupt, GeneratorExit):
+        action_coro.cancel()
+        raise
 
 
 async def execute_action_with_retry(task: "BaseTask", session: AnySession) -> Any:
@@ -178,7 +183,7 @@ async def execute_action_with_retry(task: "BaseTask", session: AnySession) -> An
             await run_async(execute_successors(task, session))
             return result
 
-        except (asyncio.CancelledError, KeyboardInterrupt):
+        except (asyncio.CancelledError, KeyboardInterrupt, GeneratorExit):
             ctx.log_warning("Task cancelled or interrupted")
             session.get_task_status(task).mark_as_failed()  # Mark as failed on cancel
             # Do not trigger fallbacks/successors on cancellation

zrb/task/base/lifecycle.py

@@ -176,7 +176,7 @@ async def log_session_state(task: AnyTask, session: AnySession):
     try:
         while not session.is_terminated:
             session.state_logger.write(session.as_state_log())
-            await asyncio.sleep(0.1)  # Log interval
+            await asyncio.sleep(0)  # Log interval
         # Log one final time after termination signal
         session.state_logger.write(session.as_state_log())
     except asyncio.CancelledError:

zrb/task/base_task.py

@@ -3,7 +3,7 @@ import inspect
 from collections.abc import Callable
 from typing import Any
 
-from zrb.attr.type import BoolAttr, fstring
+from zrb.attr.type import fstring
 from zrb.context.any_context import AnyContext
 from zrb.env.any_env import AnyEnv
 from zrb.input.any_input import AnyInput
@@ -55,7 +55,7 @@ class BaseTask(AnyTask):
         input: list[AnyInput | None] | AnyInput | None = None,
         env: list[AnyEnv | None] | AnyEnv | None = None,
         action: fstring | Callable[[AnyContext], Any] | None = None,
-        execute_condition: BoolAttr = True,
+        execute_condition: bool | str | Callable[[AnyContext], bool] = True,
         retries: int = 2,
         retry_period: float = 0,
         readiness_check: list[AnyTask] | AnyTask | None = None,
@@ -68,9 +68,18 @@ class BaseTask(AnyTask):
         fallback: list[AnyTask] | AnyTask | None = None,
         successor: list[AnyTask] | AnyTask | None = None,
     ):
-        caller_frame = inspect.stack()[1]
-        self.__decl_file = caller_frame.filename
-        self.__decl_line = caller_frame.lineno
+        # Optimized stack retrieval
+        frame = inspect.currentframe()
+        if frame is not None:
+            caller_frame = frame.f_back
+            self.__decl_file = (
+                caller_frame.f_code.co_filename if caller_frame else "unknown"
+            )
+            self.__decl_line = caller_frame.f_lineno if caller_frame else 0
+        else:
+            self.__decl_file = "unknown"
+            self.__decl_line = 0
+
         self._name = name
         self._color = color
         self._icon = icon
@@ -80,10 +89,10 @@ class BaseTask(AnyTask):
         self._envs = env
         self._retries = retries
         self._retry_period = retry_period
-        self._upstreams = upstream
-        self._fallbacks = fallback
-        self._successors = successor
-        self._readiness_checks = readiness_check
+        self._upstreams = self._ensure_task_list(upstream)
+        self._fallbacks = self._ensure_task_list(fallback)
+        self._successors = self._ensure_task_list(successor)
+        self._readiness_checks = self._ensure_task_list(readiness_check)
         self._readiness_check_delay = readiness_check_delay
         self._readiness_check_period = readiness_check_period
         self._readiness_failure_threshold = readiness_failure_threshold
@@ -92,6 +101,13 @@ class BaseTask(AnyTask):
         self._execute_condition = execute_condition
         self._action = action
 
+    def _ensure_task_list(self, tasks: AnyTask | list[AnyTask] | None) -> list[AnyTask]:
+        if tasks is None:
+            return []
+        if isinstance(tasks, list):
+            return tasks
+        return [tasks]
+
     def __repr__(self):
         return f"<{self.__class__.__name__} name={self.name}>"
 
@@ -132,18 +148,10 @@ class BaseTask(AnyTask):
     @property
     def fallbacks(self) -> list[AnyTask]:
         """Returns the list of fallback tasks."""
-        if self._fallbacks is None:
-            return []
-        elif isinstance(self._fallbacks, list):
-            return self._fallbacks
-        return [self._fallbacks]  # Assume single task
+        return self._fallbacks
 
     def append_fallback(self, fallbacks: AnyTask | list[AnyTask]):
         """Appends fallback tasks, ensuring no duplicates."""
-        if self._fallbacks is None:
-            self._fallbacks = []
-        elif not isinstance(self._fallbacks, list):
-            self._fallbacks = [self._fallbacks]
         to_add = fallbacks if isinstance(fallbacks, list) else [fallbacks]
         for fb in to_add:
             if fb not in self._fallbacks:
@@ -152,18 +160,10 @@ class BaseTask(AnyTask):
     @property
     def successors(self) -> list[AnyTask]:
         """Returns the list of successor tasks."""
-        if self._successors is None:
-            return []
-        elif isinstance(self._successors, list):
-            return self._successors
-        return [self._successors]  # Assume single task
+        return self._successors
 
     def append_successor(self, successors: AnyTask | list[AnyTask]):
         """Appends successor tasks, ensuring no duplicates."""
-        if self._successors is None:
-            self._successors = []
-        elif not isinstance(self._successors, list):
-            self._successors = [self._successors]
         to_add = successors if isinstance(successors, list) else [successors]
         for succ in to_add:
             if succ not in self._successors:
@@ -172,18 +172,10 @@ class BaseTask(AnyTask):
     @property
     def readiness_checks(self) -> list[AnyTask]:
         """Returns the list of readiness check tasks."""
-        if self._readiness_checks is None:
-            return []
-        elif isinstance(self._readiness_checks, list):
-            return self._readiness_checks
-        return [self._readiness_checks]  # Assume single task
+        return self._readiness_checks
 
     def append_readiness_check(self, readiness_checks: AnyTask | list[AnyTask]):
         """Appends readiness check tasks, ensuring no duplicates."""
-        if self._readiness_checks is None:
-            self._readiness_checks = []
-        elif not isinstance(self._readiness_checks, list):
-            self._readiness_checks = [self._readiness_checks]
         to_add = (
             readiness_checks
             if isinstance(readiness_checks, list)
@@ -196,18 +188,10 @@ class BaseTask(AnyTask):
     @property
     def upstreams(self) -> list[AnyTask]:
         """Returns the list of upstream tasks."""
-        if self._upstreams is None:
-            return []
-        elif isinstance(self._upstreams, list):
-            return self._upstreams
-        return [self._upstreams]  # Assume single task
+        return self._upstreams
 
     def append_upstream(self, upstreams: AnyTask | list[AnyTask]):
         """Appends upstream tasks, ensuring no duplicates."""
-        if self._upstreams is None:
-            self._upstreams = []
-        elif not isinstance(self._upstreams, list):
-            self._upstreams = [self._upstreams]
         to_add = upstreams if isinstance(upstreams, list) else [upstreams]
         for up in to_add:
             if up not in self._upstreams:
@@ -277,6 +261,8 @@ class BaseTask(AnyTask):
         try:
             # Delegate to the helper function for the default behavior
             return await run_default_action(self, ctx)
+        except (KeyboardInterrupt, GeneratorExit):
+            raise
         except BaseException as e:
             additional_error_note = (
                 f"Task: {self.name} ({self.__decl_file}:{self.__decl_line})"

zrb/task/base_trigger.py

@@ -5,7 +5,6 @@ from typing import Any
 from zrb.attr.type import fstring
 from zrb.callback.any_callback import AnyCallback
 from zrb.context.any_context import AnyContext
-from zrb.context.any_shared_context import AnySharedContext
 from zrb.context.shared_context import SharedContext
 from zrb.dot_dict.dot_dict import DotDict
 from zrb.env.any_env import AnyEnv

zrb/task/llm/agent.py

@@ -39,39 +39,10 @@ def create_agent_instance(
     auto_summarize: bool = True,
 ) -> "Agent[None, Any]":
     """Creates a new Agent instance with configured tools and servers."""
-    from pydantic_ai import Agent, RunContext, Tool
+    from pydantic_ai import Agent, 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)
+    ConfirmationWrapperToolset = _get_confirmation_wrapper_toolset_class()
 
     if yolo_mode is None:
         yolo_mode = False
@@ -132,6 +103,43 @@ def create_agent_instance(
     )
 
 
+def _get_confirmation_wrapper_toolset_class():
+    from pydantic_ai import RunContext
+    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)
+
+    return ConfirmationWrapperToolset
+
+
 def get_agent(
     ctx: AnyContext,
     model: "str | Model",