zrb 1.16.0__tar.gz → 1.16.2__tar.gz

{zrb-1.16.0 → zrb-1.16.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: zrb
-Version: 1.16.0
+Version: 1.16.2
 Summary: Your Automation Powerhouse
 License: AGPL-3.0-or-later
 Keywords: Automation,Task Runner,Code Generator,Monorepo,Low Code
@@ -24,12 +24,12 @@ Requires-Dist: isort (>=6.0.1,<7.0.0)
 Requires-Dist: libcst (>=1.8.2,<2.0.0)
 Requires-Dist: markdownify (>=1.2.0,<2.0.0)
 Requires-Dist: mcp (>1.12.3)
-Requires-Dist: openai (>=1.99.9)
+Requires-Dist: openai (>=1.107.2)
 Requires-Dist: pdfplumber (>=0.11.7,<0.12.0)
 Requires-Dist: playwright (>=1.54.0,<2.0.0) ; extra == "playwright" or extra == "all"
 Requires-Dist: prompt-toolkit (>=3)
 Requires-Dist: psutil (>=7.0.0,<8.0.0)
-Requires-Dist: pydantic-ai-slim[anthropic,bedrock,cohere,google,groq,huggingface,mistral,openai,vertexai] (>=1.0.1,<1.1.0)
+Requires-Dist: pydantic-ai-slim[anthropic,bedrock,cohere,google,groq,huggingface,mistral,openai,vertexai] (>=1.0.10,<1.1.0)
 Requires-Dist: pyjwt (>=2.10.1,<3.0.0)
 Requires-Dist: python-dotenv (>=1.1.1,<2.0.0)
 Requires-Dist: python-jose[cryptography] (>=3.5.0,<4.0.0)
@@ -37,6 +37,7 @@ Requires-Dist: requests (>=2.32.2,<3.0.0)
 Requires-Dist: rich (>=13)
 Requires-Dist: tiktoken (>=0.8.0,<0.9.0)
 Requires-Dist: ulid-py (>=1.1.0,<2.0.0)
+Requires-Dist: xmltodict (>=1.0.2,<2.0.0)
 Project-URL: Documentation, https://github.com/state-alchemists/zrb
 Project-URL: Homepage, https://github.com/state-alchemists/zrb
 Project-URL: Repository, https://github.com/state-alchemists/zrb

{zrb-1.16.0 → zrb-1.16.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.16.0"
+version = "1.16.2"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"
@@ -50,10 +50,10 @@ ulid-py = "^1.1.0"
 psutil = "^7.0.0"
 
 # NOTE: Look at https://github.com/pydantic/pydantic-ai/blob/main/pydantic_ai_slim/pyproject.toml
-pydantic-ai-slim = {extras = ["openai", "cohere", "vertexai", "google", "anthropic", "groq", "mistral", "bedrock", "huggingface"], version = "~1.0.1"} 
+pydantic-ai-slim = {extras = ["openai", "cohere", "vertexai", "google", "anthropic", "groq", "mistral", "bedrock", "huggingface"], version = "~1.0.10"} 
 
 # NOTE: Make sure the following dependencies are not conflicting with pydantic-ai-slim
-openai = ">=1.99.9"
+openai = ">=1.107.2"
 prompt-toolkit = ">=3"
 rich = ">=13"
 mcp = ">1.12.3"
@@ -64,6 +64,7 @@ playwright = {version = "^1.54.0", optional = true}
 chromadb = {version = "^0.6.3", optional = true}
 markdownify = "^1.2.0"
 pyjwt = "^2.10.1"
+xmltodict = "^1.0.2"
 
 [tool.poetry.extras]
 # poetry install -E rag

{zrb-1.16.0 → zrb-1.16.2}/src/zrb/builtin/llm/chat_session.py

@@ -31,15 +31,20 @@ async def read_user_prompt(ctx: AnyContext) -> str:
     final_result: str = ""
     while True:
         await asyncio.sleep(0.01)
-        previous_session_name = ctx.input.previous_session if is_first_time else None
-        if is_first_time:
-            is_first_time = False
-        # Get user input based on mode
-        if not multiline_mode:
-            ctx.print("💬 >>", plain=True)
-        user_input = await _read_next_line(reader, ctx)
-        if not multiline_mode:
-            ctx.print("", plain=True)
+        previous_session_name: str | None = (
+            ctx.input.previous_session if is_first_time else None
+        )
+        start_new: bool = ctx.input.start_new if is_first_time else False
+        if is_first_time and ctx.input.message.strip() != "":
+            user_input = ctx.input.message
+        else:
+            # Get user input based on mode
+            if not multiline_mode:
+                ctx.print("💬 >>", plain=True)
+            user_input = await _read_next_line(reader, ctx)
+            if not multiline_mode:
+                ctx.print("", plain=True)
+        is_first_time = False
         # Handle user input
         if user_input.strip().lower() in ("/bye", "/quit", "/q", "/exit"):
             user_prompt = "\n".join(user_inputs)
@@ -50,6 +55,7 @@ async def read_user_prompt(ctx: AnyContext) -> str:
                 modes=current_modes,
                 yolo_mode=current_yolo_mode,
                 previous_session_name=previous_session_name,
+                start_new=start_new,
             )
             if result is not None:
                 final_result = result
@@ -67,6 +73,7 @@ async def read_user_prompt(ctx: AnyContext) -> str:
                 modes=current_modes,
                 yolo_mode=current_yolo_mode,
                 previous_session_name=previous_session_name,
+                start_new=start_new,
             )
             if result is not None:
                 final_result = result
@@ -101,6 +108,7 @@ async def read_user_prompt(ctx: AnyContext) -> str:
                 modes=current_modes,
                 yolo_mode=current_yolo_mode,
                 previous_session_name=previous_session_name,
+                start_new=start_new,
             )
             if result is not None:
                 final_result = result

{zrb-1.16.0 → zrb-1.16.2}/src/zrb/builtin/llm/llm_ask.py

@@ -1,3 +1,4 @@
+import os
 from collections.abc import Callable
 from typing import TYPE_CHECKING
 
@@ -28,6 +29,7 @@ from zrb.callback.callback import Callback
 from zrb.config.config import CFG
 from zrb.config.llm_config import llm_config
 from zrb.context.any_context import AnyContext
+from zrb.input.any_input import AnyInput
 from zrb.input.bool_input import BoolInput
 from zrb.input.str_input import StrInput
 from zrb.input.text_input import TextInput
@@ -37,10 +39,23 @@ from zrb.util.string.conversion import to_boolean
 
 if TYPE_CHECKING:
     from pydantic_ai import Tool
+    from pydantic_ai.toolsets import AbstractToolset
 
     ToolOrCallable = Tool | Callable
 
 
+def _get_toolset(ctx: AnyContext) -> list["AbstractToolset[None] | str"]:
+    cwd = os.getcwd()
+    toolsets = []
+    for config_path in [
+        os.path.join(cwd, "mcp_config.json"),
+        os.path.join(cwd, "mcp-config.json"),
+    ]:
+        if os.path.isfile(config_path):
+            toolsets.append(config_path)
+    return toolsets
+
+
 def _get_tool(ctx: AnyContext) -> list["ToolOrCallable"]:
     tools = []
     if CFG.LLM_ALLOW_ANALYZE_REPO:
@@ -91,81 +106,89 @@ def _render_yolo_mode_input(ctx: AnyContext) -> list[str] | bool | None:
     return elements
 
 
-_llm_ask_inputs = [
-    StrInput(
-        "model",
-        description="LLM Model",
-        prompt="LLM Model",
-        default="",
-        allow_positional_parsing=False,
-        always_prompt=False,
-        allow_empty=True,
-    ),
-    StrInput(
-        "base-url",
-        description="LLM API Base URL",
-        prompt="LLM API Base URL",
-        default="",
-        allow_positional_parsing=False,
-        always_prompt=False,
-        allow_empty=True,
-    ),
-    StrInput(
-        "api-key",
-        description="LLM API Key",
-        prompt="LLM API Key",
-        default="",
-        allow_positional_parsing=False,
-        always_prompt=False,
-        allow_empty=True,
-    ),
-    TextInput(
-        "system-prompt",
-        description="System prompt",
-        prompt="System prompt",
-        default="",
-        allow_positional_parsing=False,
-        always_prompt=False,
-    ),
-    TextInput(
-        "modes",
-        description="Modes",
-        prompt="Modes",
-        default=lambda ctx: ",".join(llm_config.default_modes),
-        allow_positional_parsing=False,
-        always_prompt=False,
-    ),
-    BoolInput(
-        "start-new",
-        description="Start new session (LLM Agent will forget past conversation)",
-        prompt="Start new session (LLM Agent will forget past conversation)",
-        default=False,
-        allow_positional_parsing=False,
-        always_prompt=False,
-    ),
-    StrInput(
-        "yolo",
-        description="YOLO mode (LLM Agent will start in YOLO Mode)",
-        prompt="YOLO mode (LLM Agent will start in YOLO Mode)",
-        default=_get_default_yolo_mode,
-        allow_positional_parsing=False,
-        always_prompt=False,
-    ),
-    TextInput("message", description="User message", prompt="Your message"),
-    PreviousSessionInput(
-        "previous-session",
-        description="Previous conversation session",
-        prompt="Previous conversation session (can be empty)",
-        allow_positional_parsing=False,
-        allow_empty=True,
-        always_prompt=False,
-    ),
-]
+def _get_inputs(require_message: bool = True) -> list[AnyInput]:
+    return [
+        StrInput(
+            "model",
+            description="LLM Model",
+            prompt="LLM Model",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+            allow_empty=True,
+        ),
+        StrInput(
+            "base-url",
+            description="LLM API Base URL",
+            prompt="LLM API Base URL",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+            allow_empty=True,
+        ),
+        StrInput(
+            "api-key",
+            description="LLM API Key",
+            prompt="LLM API Key",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+            allow_empty=True,
+        ),
+        TextInput(
+            "system-prompt",
+            description="System prompt",
+            prompt="System prompt",
+            default="",
+            allow_positional_parsing=False,
+            always_prompt=False,
+        ),
+        TextInput(
+            "modes",
+            description="Modes",
+            prompt="Modes",
+            default=lambda ctx: ",".join(llm_config.default_modes),
+            allow_positional_parsing=False,
+            always_prompt=False,
+        ),
+        BoolInput(
+            "start-new",
+            description="Start new session (LLM Agent will forget past conversation)",
+            prompt="Start new session (LLM Agent will forget past conversation)",
+            default=False,
+            allow_positional_parsing=False,
+            always_prompt=False,
+        ),
+        StrInput(
+            "yolo",
+            description="YOLO mode (LLM Agent will start in YOLO Mode)",
+            prompt="YOLO mode (LLM Agent will start in YOLO Mode)",
+            default=_get_default_yolo_mode,
+            allow_positional_parsing=False,
+            always_prompt=False,
+        ),
+        TextInput(
+            "message",
+            description="User message",
+            prompt="Your message",
+            always_prompt=require_message,
+            allow_empty=not require_message,
+        ),
+        PreviousSessionInput(
+            "previous-session",
+            description="Previous conversation session",
+            prompt="Previous conversation session (can be empty)",
+            allow_positional_parsing=False,
+            allow_empty=True,
+            always_prompt=False,
+        ),
+    ]
+
 
 llm_ask: LLMTask = llm_group.add_task(
     LLMTask(
         name="llm-ask",
-        input=_llm_ask_inputs,
+        input=_get_inputs(True),
         description="❓ Ask LLM",
         model=lambda ctx: None if ctx.input.model.strip() == "" else ctx.input.model,
         model_base_url=lambda ctx: (
@@ -184,6 +207,7 @@ llm_ask: LLMTask = llm_group.add_task(
         ),
         message="{ctx.input.message}",
         tools=_get_tool,
+        toolsets=_get_toolset,
         yolo_mode=_render_yolo_mode_input,
         retries=0,
     ),
@@ -193,7 +217,7 @@ llm_ask: LLMTask = llm_group.add_task(
 llm_group.add_task(
     BaseTrigger(
         name="llm-chat",
-        input=[input for input in _llm_ask_inputs if input.name != "message"],
+        input=_get_inputs(False),
         description="💬 Chat with LLM",
         queue_name="ask_trigger",
         action=read_user_prompt,

{zrb-1.16.0 → zrb-1.16.2}/src/zrb/builtin/llm/tool/file.py

@@ -304,7 +304,7 @@ def write_to_file(
     Writes content to a file, completely overwriting it if it exists or
     creating it if it doesn't.
 
-    Use this tool to create new files or to replace the entire content of
+    Use this tool to create new files or to overwrite the entire content of
     existing files. This is a destructive operation, so be certain of your
     actions. Always read the file first to understand its contents before
     overwriting it, unless you are creating a new file.
@@ -497,7 +497,7 @@ async def analyze_file(
     ctx: AnyContext, path: str, query: str, token_limit: int | None = None
 ) -> dict[str, Any]:
     """
-    Performs a deep, goal-oriented analysis of a single file using a sub-agent.
+    Performs a high level, goal-oriented analysis of a single file using a sub-agent.
 
     This tool is ideal for complex questions about a single file that go beyond
     simple reading or searching. It uses a specialized sub-agent to analyze the

{zrb-1.16.0 → zrb-1.16.2}/src/zrb/builtin/llm/tool/sub_agent.py

@@ -24,7 +24,7 @@ def create_sub_agent_tool(
     model: "str | Model | None" = None,
     model_settings: "ModelSettings | None" = None,
     tools: "list[ToolOrCallable]" = [],
-    toolsets: list["AbstractToolset[Agent]"] = [],
+    toolsets: list["AbstractToolset[None]"] = [],
     yolo_mode: bool | list[str] | None = None,
     log_indent_level: int = 2,
 ) -> Callable[[AnyContext, str], Coroutine[Any, Any, dict[str, Any]]]:

{zrb-1.16.0 → zrb-1.16.2}/src/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-1.16.0 → zrb-1.16.2}/src/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-1.16.0 → zrb-1.16.2}/src/zrb/task/llm/agent.py

@@ -125,8 +125,8 @@ def get_agent(
         "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]]" = [],
+    toolsets_attr: "list[AbstractToolset[None] | str] | Callable[[AnySharedContext], list[AbstractToolset[None] | str]]" = [],  # noqa
+    additional_toolsets: "list[AbstractToolset[None] | str]" = [],
     retries: int = 3,
     yolo_mode: bool | list[str] | None = None,
 ) -> "Agent":
@@ -149,8 +149,11 @@ def get_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,
@@ -158,13 +161,32 @@ def get_agent(
         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[None, Any]",

{zrb-1.16.0 → zrb-1.16.2}/src/zrb/task/llm/config.py

@@ -4,7 +4,7 @@ if TYPE_CHECKING:
     from pydantic_ai.models import Model
     from pydantic_ai.settings import ModelSettings
 
-from zrb.attr.type import BoolAttr, StrAttr, fstring
+from zrb.attr.type import BoolAttr, StrAttr, StrListAttr, fstring
 from zrb.config.llm_config import LLMConfig, llm_config
 from zrb.context.any_context import AnyContext
 from zrb.context.any_shared_context import AnySharedContext
@@ -13,7 +13,7 @@ from zrb.util.attr import get_attr, get_bool_attr, get_str_list_attr
 
 def get_yolo_mode(
     ctx: AnyContext,
-    yolo_mode_attr: BoolAttr | StrAttr | None = None,
+    yolo_mode_attr: BoolAttr | StrListAttr | None = None,
     render_yolo_mode: bool = True,
 ) -> bool | list[str]:
     if yolo_mode_attr is None:

{zrb-1.16.0 → zrb-1.16.2}/src/zrb/task/llm/tool_wrapper.py

@@ -1,5 +1,6 @@
 import functools
 import inspect
+import json
 import traceback
 import typing
 from collections.abc import Callable
@@ -172,18 +173,17 @@ def _get_edited_kwargs(
         return kwargs, False
     while len(user_edit_responses) < 3:
         user_edit_responses.append("")
-    key, val = user_edit_responses[1:]
+    key, val_str = user_edit_responses[1:]
     if key not in kwargs:
         return kwargs, True
-    if val != "":
-        kwargs[key] = val
-        return kwargs, True
-    val = edit_text(
-        prompt_message=f"// {key}",
-        value=kwargs.get(key, ""),
-        editor=CFG.DEFAULT_EDITOR,
-    )
-    kwargs[key] = val
+    is_str_param = isinstance(kwargs[key], str)
+    if val_str == "":
+        val_str = edit_text(
+            prompt_message=f"// {key}",
+            value=_get_val_str(kwargs[key]),
+            editor=CFG.DEFAULT_EDITOR,
+        )
+    kwargs[key] = val_str if is_str_param else json.loads(val_str)
     return kwargs, True
 
 
@@ -233,7 +233,7 @@ def _get_detail_func_param(args: list[Any] | tuple[Any], kwargs: dict[str, Any])
 
 def _get_func_param_item(key: str, val: Any) -> str:
     upper_key = key.upper()
-    val_str = f"{val}"
+    val_str = _get_val_str(val)
     val_parts = val_str.split("\n")
     if len(val_parts) == 1:
         return f"- {upper_key} `{val}`"
@@ -244,6 +244,15 @@ def _get_func_param_item(key: str, val: Any) -> str:
     return "\n".join(lines)
 
 
+def _get_val_str(val: Any) -> str:
+    if isinstance(val, str):
+        return val
+    try:
+        return json.dumps(val, indent=4)
+    except Exception:
+        return f"{val}"
+
+
 def _get_func_call_str(
     func: Callable, args: list[Any] | tuple[Any], kwargs: dict[str, Any]
 ) -> str:

{zrb-1.16.0 → zrb-1.16.2}/src/zrb/task/llm_task.py

@@ -80,7 +80,8 @@ class LLMTask(BaseTask):
             | Callable[[AnySharedContext], list["ToolOrCallable"]]
         ) = [],
         toolsets: (
-            list["AbstractToolset[Agent]"] | Callable[[AnySharedContext], list["Tool"]]
+            list["AbstractToolset[None] | str"]
+            | Callable[[AnySharedContext], list["AbstractToolset[None] | str"]]
         ) = [],
         conversation_history: (
             ConversationHistory
@@ -169,7 +170,7 @@ class LLMTask(BaseTask):
         self._rate_limitter = rate_limitter
         self._additional_tools: list["ToolOrCallable"] = []
         self._toolsets = toolsets
-        self._additional_toolsets: list["AbstractToolset[Agent]"] = []
+        self._additional_toolsets: list["AbstractToolset[None] | str"] = []
         self._conversation_history = conversation_history
         self._conversation_history_reader = conversation_history_reader
         self._conversation_history_writer = conversation_history_writer
@@ -200,10 +201,12 @@ class LLMTask(BaseTask):
         for single_tool in tool:
             self._additional_tools.append(single_tool)
 
-    def add_toolset(self, *toolset: "AbstractToolset[Agent]"):
+    def add_toolset(self, *toolset: "AbstractToolset[None] | str"):
         self.append_toolset(*toolset)
 
-    def append_toolset(self, *toolset: "AbstractToolset[Agent]"):
+    def append_toolset(self, *toolset: "AbstractToolset[None] | str"):
+        from pydantic_ai.mcp import load_mcp_servers
+
         for single_toolset in toolset:
             self._additional_toolsets.append(single_toolset)