zrb 1.5.11__py3-none-any.whl → 1.5.13__py3-none-any.whl

zrb/task/llm/tool_wrapper.py

@@ -1,58 +1,107 @@
 import functools
 import inspect
 import traceback
+import typing
 from collections.abc import Callable
 
+from pydantic_ai import RunContext, Tool
+
+from zrb.context.any_context import AnyContext
 from zrb.task.llm.error import ToolExecutionError
 from zrb.util.run import run_async
 
 
-def wrap_tool(func: Callable) -> Callable:
-    """Wraps a tool function to handle exceptions and context propagation.
+def wrap_tool(func: Callable, ctx: AnyContext) -> Tool:
+    """Wraps a tool function to handle exceptions and context propagation."""
+    original_sig = inspect.signature(func)
+    # Use helper function for clarity
+    needs_run_context_for_pydantic = _has_context_parameter(original_sig, RunContext)
+    needs_any_context_for_injection = _has_context_parameter(original_sig, AnyContext)
+    takes_no_args = len(original_sig.parameters) == 0
+    # Pass individual flags to the wrapper creator
+    wrapper = _create_wrapper(func, original_sig, ctx, needs_any_context_for_injection)
+    # Adjust signature - _adjust_signature determines exclusions based on type
+    _adjust_signature(wrapper, original_sig, takes_no_args)
+    # takes_ctx in pydantic-ai Tool is specifically for RunContext
+    return Tool(wrapper, takes_ctx=needs_run_context_for_pydantic)
+
 
-    - Catches exceptions during tool execution and returns a structured
-      JSON error message (`ToolExecutionError`).
-    - Inspects the original function signature for a 'ctx' parameter.
-      If found, the wrapper will accept 'ctx' and pass it to the function.
-    - If the original function has no parameters, injects a dummy '_dummy'
-      parameter into the wrapper's signature to ensure schema generation
-      for pydantic-ai.
+def _has_context_parameter(original_sig: inspect.Signature, context_type: type) -> bool:
+    """
+    Checks if the function signature includes a parameter with the specified
+    context type annotation.
+    """
+    return any(
+        _is_annotated_with_context(param.annotation, context_type)
+        for param in original_sig.parameters.values()
+    )
 
-    Args:
-        func: The tool function (sync or async) to wrap.
 
-    Returns:
-        An async wrapper function.
+def _is_annotated_with_context(param_annotation, context_type):
     """
-    original_sig = inspect.signature(func)
-    needs_ctx = "ctx" in original_sig.parameters
-    takes_no_args = len(original_sig.parameters) == 0
+    Checks if the parameter annotation is the specified context type
+    or a generic type containing it (e.g., Optional[ContextType]).
+    """
+    if param_annotation is inspect.Parameter.empty:
+        return False
+    if param_annotation is context_type:
+        return True
+    # Check for generic types like Optional[ContextType] or Union[ContextType, ...]
+    origin = typing.get_origin(param_annotation)
+    args = typing.get_args(param_annotation)
+    if origin is not None and args:
+        # Check if context_type is one of the arguments of the generic type
+        return any(arg is context_type for arg in args)
+    return False
+
+
+def _create_wrapper(
+    func: Callable,
+    original_sig: inspect.Signature,
+    ctx: AnyContext,  # Accept ctx
+    needs_any_context_for_injection: bool,
+) -> Callable:
+    """Creates the core wrapper function."""
 
     @functools.wraps(func)
     async def wrapper(*args, **kwargs):
-        ctx_arg = None
-        if needs_ctx:
-            if args:
-                ctx_arg = args[0]
-                args = args[1:]  # Remove ctx from args for the actual call if needed
-            elif "ctx" in kwargs:
-                ctx_arg = kwargs.pop("ctx")
+        # Identify AnyContext parameter name from the original signature if needed
+        any_context_param_name = None
+
+        if needs_any_context_for_injection:
+            for param in original_sig.parameters.values():
+                if _is_annotated_with_context(param.annotation, AnyContext):
+                    any_context_param_name = param.name
+                    break  # Found it, no need to continue
+
+            if any_context_param_name is None:
+                # This should not happen if needs_any_context_for_injection is True,
+                # but check for safety
+                raise ValueError(
+                    "AnyContext parameter name not found in function signature."
+                )
+            # Inject the captured ctx into kwargs. This will overwrite if the LLM
+            # somehow provided it.
+            kwargs[any_context_param_name] = ctx
+
+        # If the dummy argument was added for schema generation and is present in kwargs,
+        # remove it before calling the original function, unless the original function
+        # actually expects a parameter named '_dummy'.
+        if "_dummy" in kwargs and "_dummy" not in original_sig.parameters:
+            del kwargs["_dummy"]
 
         try:
-            # Remove dummy argument if it was added for no-arg functions
-            if takes_no_args and "_dummy" in kwargs:
-                del kwargs["_dummy"]
-
-            if needs_ctx:
-                # Ensure ctx is passed correctly, even if original func had only ctx
-                if ctx_arg is None:
-                    # This case should ideally not happen if takes_ctx is True in Tool
-                    raise ValueError("Context (ctx) was expected but not provided.")
-                # Call with context
-                return await run_async(func(ctx_arg, *args, **kwargs))
-            else:
-                # Call without context
-                return await run_async(func(*args, **kwargs))
+            # Call the original function.
+            # pydantic-ai is responsible for injecting RunContext if takes_ctx is True.
+            # Our wrapper injects AnyContext if needed.
+            # The arguments received by the wrapper (*args, **kwargs) are those
+            # provided by the LLM, potentially with RunContext already injected by
+            # pydantic-ai if takes_ctx is True. We just need to ensure AnyContext
+            # is injected if required by the original function.
+            # The dummy argument handling is moved to _adjust_signature's logic
+            # for schema generation, it's not needed here before calling the actual
+            # function.
+            return await run_async(func(*args, **kwargs))
         except Exception as e:
             error_model = ToolExecutionError(
                 tool_name=func.__name__,
@@ -62,8 +111,32 @@ def wrap_tool(func: Callable) -> Callable:
             )
             return error_model.model_dump_json()
 
-    if takes_no_args:
-        # Inject dummy parameter for schema generation if original func took no args
+    return wrapper
+
+
+def _adjust_signature(
+    wrapper: Callable, original_sig: inspect.Signature, takes_no_args: bool
+):
+    """Adjusts the wrapper function's signature for schema generation."""
+    # The wrapper's signature should represent the arguments the *LLM* needs to provide.
+    # The LLM does not provide RunContext (pydantic-ai injects it) or AnyContext
+    # (we inject it). So, the wrapper's signature should be the original signature,
+    # minus any parameters annotated with RunContext or AnyContext.
+
+    params_for_schema = [
+        param
+        for param in original_sig.parameters.values()
+        if not _is_annotated_with_context(param.annotation, RunContext)
+        and not _is_annotated_with_context(param.annotation, AnyContext)
+    ]
+
+    # If after removing context parameters, there are no parameters left,
+    # and the original function took no args, keep the dummy.
+    # If after removing context parameters, there are no parameters left,
+    # but the original function *did* take args (only context), then the schema
+    # should have no parameters.
+    if not params_for_schema and takes_no_args:
+        # Keep the dummy if the original function truly had no parameters
         new_sig = inspect.Signature(
             parameters=[
                 inspect.Parameter(
@@ -71,18 +144,7 @@ def wrap_tool(func: Callable) -> Callable:
                 )
             ]
         )
-        wrapper.__signature__ = new_sig
-    elif needs_ctx:
-        # Adjust signature if ctx was the *only* parameter originally
-        # This ensures the wrapper signature matches what pydantic-ai expects
-        params = list(original_sig.parameters.values())
-        if len(params) == 1 and params[0].name == "ctx":
-            new_sig = inspect.Signature(
-                parameters=[
-                    inspect.Parameter("ctx", inspect.Parameter.POSITIONAL_OR_KEYWORD)
-                ]
-            )
-            wrapper.__signature__ = new_sig
-        # Otherwise, the original signature (including ctx) is fine for the wrapper
+    else:
+        new_sig = inspect.Signature(parameters=params_for_schema)
 
-    return wrapper
+    wrapper.__signature__ = new_sig

zrb/task/llm_task.py

@@ -77,6 +77,8 @@ class LLMTask(BaseTask):
         render_enrich_context: bool = True,
         context_enrichment_prompt: StrAttr | None = None,
         render_context_enrichment_prompt: bool = True,
+        context_enrichment_threshold: IntAttr | None = None,
+        render_context_enrichment_threshold: bool = True,
         tools: (
             list[ToolOrCallable] | Callable[[AnySharedContext], list[ToolOrCallable]]
         ) = [],
@@ -161,6 +163,8 @@ class LLMTask(BaseTask):
         self._render_enrich_context = render_enrich_context
         self._context_enrichment_prompt = context_enrichment_prompt
         self._render_context_enrichment_prompt = render_context_enrichment_prompt
+        self._context_enrichment_threshold = context_enrichment_threshold
+        self._render_context_enrichment_threshold = render_context_enrichment_threshold
         self._tools = tools
         self._additional_tools: list[ToolOrCallable] = []
         self._mcp_servers = mcp_servers
@@ -180,14 +184,23 @@ class LLMTask(BaseTask):
         self._conversation_context = conversation_context
 
     def add_tool(self, tool: ToolOrCallable):
+        self.append_tool(tool)
+
+    def append_tool(self, tool: ToolOrCallable):
         self._additional_tools.append(tool)
 
     def add_mcp_server(self, mcp_server: MCPServer):
+        self.append_mcp_server(mcp_server)
+
+    def append_mcp_server(self, mcp_server: MCPServer):
         self._additional_mcp_servers.append(mcp_server)
 
     def set_should_enrich_context(self, enrich_context: bool):
         self._should_enrich_context = enrich_context
 
+    def set_context_enrichment_threshold(self, enrichment_threshold: int):
+        self._context_enrichment_threshold = enrichment_threshold
+
     def set_should_summarize_history(self, summarize_history: bool):
         self._should_summarize_history = summarize_history
 
@@ -244,6 +257,8 @@ class LLMTask(BaseTask):
             conversation_context=conversation_context,
             should_enrich_context_attr=self._should_enrich_context,
             render_enrich_context=self._render_enrich_context,
+            context_enrichment_threshold_attr=self._context_enrichment_threshold,
+            render_context_enrichment_threshold=self._render_context_enrichment_threshold,
             model=model,
             model_settings=model_settings,
             context_enrichment_prompt=context_enrichment_prompt,
@@ -321,7 +336,7 @@ class LLMTask(BaseTask):
                 usage = agent_run.result.usage()
                 ctx.xcom.get(xcom_usage_key).push(usage)
                 ctx.print(stylize_faint(f"[USAGE] {usage}"))
-                return agent_run.result.data
+                return agent_run.result.output
             else:
                 ctx.log_warning("Agent run did not produce a result.")
                 return None  # Or handle as appropriate

zrb/util/attr.py

@@ -16,6 +16,17 @@ from zrb.util.string.conversion import to_boolean
 def get_str_list_attr(
     shared_ctx: AnySharedContext, attr: StrListAttr | None, auto_render: bool = True
 ) -> list[str]:
+    """
+    Retrieve a list of strings from shared context attributes.
+
+    Args:
+        shared_ctx (AnySharedContext): The shared context object.
+        attr (StrListAttr | None): The string list attribute to retrieve.
+        auto_render (bool): Whether to auto-render the attribute values.
+
+    Returns:
+        list[str]: A list of string attributes.
+    """
     if callable(attr):
         return attr(shared_ctx)
     return {get_str_attr(shared_ctx, val, "", auto_render) for val in attr}
@@ -24,6 +35,17 @@ def get_str_list_attr(
 def get_str_dict_attr(
     shared_ctx: AnySharedContext, attr: StrDictAttr | None, auto_render: bool = True
 ) -> dict[str, Any]:
+    """
+    Retrieve a dictionary of strings from shared context attributes.
+
+    Args:
+        shared_ctx (AnySharedContext): The shared context object.
+        attr (StrDictAttr | None): The string dictionary attribute to retrieve.
+        auto_render (bool): Whether to auto-render the attribute values.
+
+    Returns:
+        dict[str, Any]: A dictionary of string attributes.
+    """
     if callable(attr):
         return attr(shared_ctx)
     return {
@@ -37,6 +59,18 @@ def get_str_attr(
     default: StrAttr = "",
     auto_render: bool = True,
 ) -> str:
+    """
+    Retrieve a string from shared context attributes.
+
+    Args:
+        shared_ctx (AnySharedContext): The shared context object.
+        attr (StrAttr | None): The string attribute to retrieve.
+        default (StrAttr): The default value if the attribute is None.
+        auto_render (bool): Whether to auto-render the attribute value.
+
+    Returns:
+        str: The string attribute value.
+    """
     val = get_attr(shared_ctx, attr, default, auto_render)
     if not isinstance(val, str):
         return str(val)
@@ -49,6 +83,18 @@ def get_bool_attr(
     default: BoolAttr = False,
     auto_render: bool = True,
 ) -> bool:
+    """
+    Retrieve a boolean from shared context attributes.
+
+    Args:
+        shared_ctx (AnySharedContext): The shared context object.
+        attr (BoolAttr | None): The boolean attribute to retrieve.
+        default (BoolAttr): The default value if the attribute is None.
+        auto_render (bool): Whether to auto-render the attribute value if it's a string.
+
+    Returns:
+        bool: The boolean attribute value.
+    """
     val = get_attr(shared_ctx, attr, default, auto_render)
     if isinstance(val, str):
         return to_boolean(val)
@@ -61,6 +107,18 @@ def get_int_attr(
     default: IntAttr = 0,
     auto_render: bool = True,
 ) -> int:
+    """
+    Retrieve an integer from shared context attributes.
+
+    Args:
+        shared_ctx (AnySharedContext): The shared context object.
+        attr (IntAttr | None): The integer attribute to retrieve.
+        default (IntAttr): The default value if the attribute is None.
+        auto_render (bool): Whether to auto-render the attribute value if it's a string.
+
+    Returns:
+        int: The integer attribute value.
+    """
     val = get_attr(shared_ctx, attr, default, auto_render)
     if isinstance(val, str):
         return int(val)
@@ -72,7 +130,19 @@ def get_float_attr(
     attr: FloatAttr | None,
     default: FloatAttr = 0.0,
     auto_render: bool = True,
-) -> str | None:
+) -> float | None:
+    """
+    Retrieve a float from shared context attributes.
+
+    Args:
+        shared_ctx (AnySharedContext): The shared context object.
+        attr (FloatAttr | None): The float attribute to retrieve.
+        default (FloatAttr): The default value if the attribute is None.
+        auto_render (bool): Whether to auto-render the attribute value if it's a string.
+
+    Returns:
+        float | None: The float attribute value.
+    """
     val = get_attr(shared_ctx, attr, default, auto_render)
     if isinstance(val, str):
         return float(val)
@@ -85,6 +155,19 @@ def get_attr(
     default: AnyAttr,
     auto_render: bool = True,
 ) -> Any | None:
+    """
+    Retrieve an attribute value from shared context, handling callables and rendering.
+
+    Args:
+        shared_ctx (AnySharedContext): The shared context object.
+        attr (AnyAttr): The attribute to retrieve. Can be a value, a callable,
+            or a string to render.
+        default (AnyAttr): The default value if the attribute is None.
+        auto_render (bool): Whether to auto-render the attribute value if it's a string.
+
+    Returns:
+        Any | None: The retrieved attribute value or the default value.
+    """
     if attr is None:
         if callable(default):
             return default(shared_ctx)

zrb/util/cli/style.py

@@ -121,6 +121,15 @@ ICONS = [
 
 
 def remove_style(text):
+    """
+    Remove ANSI escape codes from a string.
+
+    Args:
+        text (str): The input string with potential ANSI escape codes.
+
+    Returns:
+        str: The string with ANSI escape codes removed.
+    """
     ansi_escape = re.compile(r"\x1B[@-_][0-?]*[ -/]*[@-~]")
     return ansi_escape.sub("", text)
 
@@ -131,6 +140,18 @@ def stylize(
     background: int | None = None,
     style: int | None = None,
 ):
+    """
+    Apply ANSI escape codes to a string for terminal styling.
+
+    Args:
+        text (str): The input string to stylize.
+        color (int | None): The foreground color code.
+        background (int | None): The background color code.
+        style (int | None): The text style code (e.g., bold, underline).
+
+    Returns:
+        str: The stylized string with ANSI escape codes.
+    """
     # Start constructing the ANSI escape code
     code_parts = []
     if style is not None and style in VALID_STYLES:
@@ -146,56 +167,182 @@ def stylize(
 
 
 def stylize_section_header(text: str):
+    """
+    Stylize text as a section header.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized section header string.
+    """
     return stylize(f" {text} ", color=BLACK, background=BG_WHITE, style=UNDERLINE)
 
 
 def stylize_green(text: str):
+    """
+    Stylize text with green foreground color.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize(text, color=GREEN)
 
 
 def stylize_blue(text: str):
+    """
+    Stylize text with blue foreground color.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize(text, color=BLUE)
 
 
 def stylize_cyan(text: str):
+    """
+    Stylize text with cyan foreground color.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize(text, color=CYAN)
 
 
 def stylize_magenta(text: str):
+    """
+    Stylize text with magenta foreground color.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize(text, color=MAGENTA)
 
 
 def stylize_yellow(text: str):
+    """
+    Stylize text with yellow foreground color.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize(text, color=YELLOW)
 
 
 def stylize_red(text: str):
+    """
+    Stylize text with red foreground color.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize(text, color=RED)
 
 
 def stylize_bold_green(text: str):
+    """
+    Stylize text with bold green foreground color.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize(text, color=GREEN, style=BOLD)
 
 
 def stylize_bold_yellow(text: str):
+    """
+    Stylize text with bold yellow foreground color.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize(text, color=YELLOW, style=BOLD)
 
 
 def stylize_bold_red(text: str):
+    """
+    Stylize text with bold red foreground color.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize(text, color=RED, style=BOLD)
 
 
 def stylize_faint(text: str):
+    """
+    Stylize text with faint style.
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize(text, style=FAINT)
 
 
 def stylize_log(text: str):
+    """
+    Stylize text for log messages (faint).
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize_faint(text)
 
 
 def stylize_warning(text: str):
+    """
+    Stylize text for warning messages (bold yellow).
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize_bold_yellow(text)
 
 
 def stylize_error(text: str):
+    """
+    Stylize text for error messages (bold red).
+
+    Args:
+        text (str): The input string.
+
+    Returns:
+        str: The stylized string.
+    """
     return stylize_bold_red(text)

zrb/util/cli/subcommand.py

@@ -4,16 +4,37 @@ from zrb.util.group import get_non_empty_subgroups, get_subtasks
 
 class SubCommand:
     def __init__(self, paths: list[str] = [], nexts: list[str] = []):
+        """
+        Initialize a SubCommand object.
+
+        Args:
+            paths (list[str]): The list of path components leading to this subcommand.
+            nexts (list[str]): The list of possible next subcommand or task names.
+        """
         self.paths = paths
         self.nexts = nexts
 
     def __repr__(self):
+        """
+        Return a string representation of the SubCommand object.
+        """
         return f"<{self.__class__.__name__} paths={self.paths} nexts={self.nexts}>"
 
 
 def get_group_subcommands(
-    group: AnyGroup, previous_path: str = [], subcommands: list[SubCommand] = []
+    group: AnyGroup, previous_path: list[str] = [], subcommands: list[SubCommand] = []
 ) -> list[SubCommand]:
+    """
+    Recursively get all possible subcommands within a group hierarchy.
+
+    Args:
+        group (AnyGroup): The current group to process.
+        previous_path (list[str]): The path leading to the current group.
+        subcommands (list[SubCommand]): The list to accumulate subcommands.
+
+    Returns:
+        list[SubCommand]: A list of all discovered subcommands.
+    """
     nexts = []
     for task_alias in get_subtasks(group):
         nexts.append(task_alias)

zrb/util/cmd/command.py

@@ -10,6 +10,16 @@ from zrb.cmd.cmd_result import CmdResult
 
 
 def check_unrecommended_commands(cmd_script: str) -> dict[str, str]:
+    """
+    Check a command script for the use of unrecommended or non-POSIX compliant commands.
+
+    Args:
+        cmd_script (str): The command script string to check.
+
+    Returns:
+        dict[str, str]: A dictionary where keys are the violating commands/patterns
+            and values are the reasons they are unrecommended.
+    """
     banned_commands = {
         "<(": "Process substitution isn't POSIX compliant and causes trouble",
         "column": "Command isn't included in Ubuntu packages and is not POSIX compliant",
@@ -96,6 +106,14 @@ async def run_command(
 
 
 def kill_pid(pid: int, print_method: Callable[..., None] | None = None):
+    """
+    Kill a process and its children given the parent process ID.
+
+    Args:
+        pid (int): The process ID of the parent process.
+        print_method (Callable[..., None] | None): A method to print status messages.
+            Defaults to the built-in print function.
+    """
     actual_print_method = print_method if print_method is not None else print
     parent = psutil.Process(pid)
     children = parent.children(recursive=True)