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

zrb/task/llm/tool_wrapper.py

@@ -1,35 +1,61 @@
 import functools
 import inspect
+import signal
 import traceback
 import typing
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
+from zrb.config.config import CFG
+from zrb.config.llm_rate_limitter import llm_rate_limitter
 from zrb.context.any_context import AnyContext
 from zrb.task.llm.error import ToolExecutionError
+from zrb.task.llm.file_replacement import edit_replacement, is_single_path_replacement
+from zrb.util.callable import get_callable_name
+from zrb.util.cli.markdown import render_markdown
+from zrb.util.cli.style import (
+    stylize_blue,
+    stylize_error,
+    stylize_faint,
+    stylize_green,
+    stylize_yellow,
+)
+from zrb.util.cli.text import edit_text
 from zrb.util.run import run_async
+from zrb.util.string.conversion import to_boolean
+from zrb.util.yaml import edit_obj, yaml_dump
 
 if TYPE_CHECKING:
     from pydantic_ai import Tool
-else:
-    Tool = Any
 
 
-def wrap_tool(func: Callable, ctx: AnyContext) -> Tool:
+class ToolExecutionCancelled(ValueError):
+    pass
+
+
+def wrap_tool(func: Callable, ctx: AnyContext, yolo_mode: bool | list[str]) -> "Tool":
     """Wraps a tool function to handle exceptions and context propagation."""
     from pydantic_ai import RunContext, Tool
 
     original_sig = inspect.signature(func)
-    # Use helper function for clarity
     needs_run_context_for_pydantic = _has_context_parameter(original_sig, RunContext)
+    wrapper = wrap_func(func, ctx, yolo_mode)
+    return Tool(wrapper, takes_ctx=needs_run_context_for_pydantic)
+
+
+def wrap_func(func: Callable, ctx: AnyContext, yolo_mode: bool | list[str]) -> Callable:
+    original_sig = inspect.signature(func)
     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)
+    wrapper = _create_wrapper(
+        func=func,
+        original_signature=original_sig,
+        ctx=ctx,
+        needs_any_context_for_injection=needs_any_context_for_injection,
+        yolo_mode=yolo_mode,
+    )
+    _adjust_signature(wrapper, original_sig)
+    return wrapper
 
 
 def _has_context_parameter(original_sig: inspect.Signature, context_type: type) -> bool:
@@ -63,9 +89,10 @@ def _is_annotated_with_context(param_annotation, context_type):
 
 def _create_wrapper(
     func: Callable,
-    original_sig: inspect.Signature,
-    ctx: AnyContext,  # Accept ctx
+    original_signature: inspect.Signature,
+    ctx: AnyContext,
     needs_any_context_for_injection: bool,
+    yolo_mode: bool | list[str],
 ) -> Callable:
     """Creates the core wrapper function."""
 
@@ -73,13 +100,11 @@ def _create_wrapper(
     async def wrapper(*args, **kwargs):
         # 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():
+            for param in original_signature.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
@@ -89,40 +114,236 @@ def _create_wrapper(
             # 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"]
-
+        # We will need to overwrite SIGINT handler, so that when user press ctrl + c,
+        # the program won't immediately exit
+        original_sigint_handler = signal.getsignal(signal.SIGINT)
+        tool_name = get_callable_name(func)
         try:
-            # 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:
+            has_ever_edited = False
+            if not ctx.is_web_mode and ctx.is_tty:
+                if (
+                    isinstance(yolo_mode, list) and func.__name__ not in yolo_mode
+                ) or not yolo_mode:
+                    approval, reason, kwargs, has_ever_edited = (
+                        await _handle_user_response(ctx, func, args, kwargs)
+                    )
+                    if not approval:
+                        raise ToolExecutionCancelled(
+                            f"Tool execution cancelled. User disapproving: {reason}"
+                        )
+            signal.signal(signal.SIGINT, _tool_wrapper_sigint_handler)
+            ctx.print(stylize_faint(f"Run {tool_name}"), plain=True)
+            result = await run_async(func(*args, **kwargs))
+            _check_tool_call_result_limit(result)
+            if has_ever_edited:
+                return {
+                    "tool_call_result": result,
+                    "new_tool_parameters": kwargs,
+                    "message": "User correction: Tool was called with user's parameters",
+                }
+            return result
+        except BaseException as e:
             error_model = ToolExecutionError(
-                tool_name=func.__name__,
+                tool_name=tool_name,
                 error_type=type(e).__name__,
                 message=str(e),
                 details=traceback.format_exc(),
             )
             return error_model.model_dump_json()
+        finally:
+            signal.signal(signal.SIGINT, original_sigint_handler)
 
     return wrapper
 
 
-def _adjust_signature(
-    wrapper: Callable, original_sig: inspect.Signature, takes_no_args: bool
-):
+def _tool_wrapper_sigint_handler(signum, frame):
+    raise KeyboardInterrupt("SIGINT detected while running tool")
+
+
+def _check_tool_call_result_limit(result: Any):
+    if (
+        llm_rate_limitter.count_token(result)
+        > llm_rate_limitter.max_tokens_per_tool_call_result
+    ):
+        raise ValueError("Result value is too large, please adjust the parameter")
+
+
+async def _handle_user_response(
+    ctx: AnyContext,
+    func: Callable,
+    args: list[Any] | tuple[Any],
+    kwargs: dict[str, Any],
+) -> tuple[bool, str, dict[str, Any], bool]:
+    has_ever_edited = False
+    while True:
+        func_call_str = _get_func_call_str(func, args, kwargs)
+        complete_confirmation_message = "\n".join(
+            [
+                f"\n🎰 >> {func_call_str}",
+                _get_detail_func_param(args, kwargs),
+                f"🎰  >> {_get_run_func_confirmation(func)}",
+            ]
+        )
+        ctx.print(complete_confirmation_message, plain=True)
+        user_response = await _read_line(args, kwargs)
+        ctx.print("", plain=True)
+        new_kwargs, is_edited = _get_edited_kwargs(ctx, user_response, kwargs)
+        if is_edited:
+            kwargs = new_kwargs
+            has_ever_edited = True
+            continue
+        approval_and_reason = _get_user_approval_and_reason(
+            ctx, user_response, func_call_str
+        )
+        if approval_and_reason is None:
+            continue
+        approval, reason = approval_and_reason
+        return approval, reason, kwargs, has_ever_edited
+
+
+def _get_edited_kwargs(
+    ctx: AnyContext, user_response: str, kwargs: dict[str, Any]
+) -> tuple[dict[str, Any], bool]:
+    user_edit_responses = [val for val in user_response.split(" ", maxsplit=2)]
+    if len(user_edit_responses) >= 1 and user_edit_responses[0].lower() != "edit":
+        return kwargs, False
+    while len(user_edit_responses) < 3:
+        user_edit_responses.append("")
+    key, val_str = user_edit_responses[1:]
+    # Make sure first segment of the key is in kwargs
+    if key != "":
+        key_parts = key.split(".")
+        if len(key_parts) > 0 and key_parts[0] not in kwargs:
+            return kwargs, True
+    # Handle replacement edit
+    if len(kwargs) == 1:
+        kwarg_key = list(kwargs.keys())[0]
+        if is_single_path_replacement(kwargs[kwarg_key]) and (
+            key == "" or key == kwarg_key
+        ):
+            kwargs[kwarg_key], edited = edit_replacement(kwargs[kwarg_key])
+            return kwargs, True
+    # Handle other kind of edit
+    old_val_str = yaml_dump(kwargs, key)
+    if val_str == "":
+        val_str = edit_text(
+            prompt_message=f"# {key}" if key != "" else "",
+            value=old_val_str,
+            editor=CFG.DEFAULT_EDITOR,
+            extension=".yaml",
+        )
+    if old_val_str == val_str:
+        return kwargs, True
+    edited_kwargs = edit_obj(kwargs, key, val_str)
+    return edited_kwargs, True
+
+
+def _get_user_approval_and_reason(
+    ctx: AnyContext, user_response: str, func_call_str: str
+) -> tuple[bool, str] | None:
+    user_approval_responses = [
+        val.strip() for val in user_response.split(",", maxsplit=1)
+    ]
+    while len(user_approval_responses) < 2:
+        user_approval_responses.append("")
+    approval_str, reason = user_approval_responses
+    try:
+        approved = True if approval_str.strip() == "" else to_boolean(approval_str)
+        if not approved and reason == "":
+            reason = "User disapproving the tool execution"
+        return approved, reason
+    except Exception:
+        return False, user_response
+
+
+def _get_run_func_confirmation(func: Callable) -> str:
+    func_name = get_callable_name(func)
+    return render_markdown(
+        f"Allow to run `{func_name}`? (✅ `Yes` | ⛔ `No, <reason>` | 📝 `Edit <param> <value>`)"
+    ).strip()
+
+
+def _get_detail_func_param(args: list[Any] | tuple[Any], kwargs: dict[str, Any]) -> str:
+    if not kwargs:
+        return ""
+    yaml_str = yaml_dump(kwargs)
+    # Create the final markdown string
+    markdown = f"```yaml\n{yaml_str}\n```"
+    return render_markdown(markdown)
+
+
+def _get_func_call_str(
+    func: Callable, args: list[Any] | tuple[Any], kwargs: dict[str, Any]
+) -> str:
+    func_name = get_callable_name(func)
+    normalized_args = [stylize_green(_truncate_arg(arg)) for arg in args]
+    normalized_kwargs = []
+    for key, val in kwargs.items():
+        truncated_val = _truncate_arg(f"{val}")
+        normalized_kwargs.append(
+            f"{stylize_yellow(key)}={stylize_green(truncated_val)}"
+        )
+    func_param_str = ", ".join(normalized_args + normalized_kwargs)
+    return f"{stylize_blue(func_name + '(')}{func_param_str}{stylize_blue(')')}"
+
+
+def _truncate_arg(arg: str, length: int = 19) -> str:
+    normalized_arg = arg.replace("\n", "\\n")
+    if len(normalized_arg) > length:
+        return f"{normalized_arg[:length-4]} ..."
+    return normalized_arg
+
+
+async def _read_line(args: list[Any] | tuple[Any], kwargs: dict[str, Any]):
+    from prompt_toolkit import PromptSession
+    from prompt_toolkit.completion import Completer, Completion
+
+    class ToolConfirmationCompleter(Completer):
+        """Custom completer for tool confirmation that doesn't auto-complete partial words."""
+
+        def __init__(self, options, meta_dict):
+            self.options = options
+            self.meta_dict = meta_dict
+
+        def get_completions(self, document, complete_event):
+            text = document.text.strip()
+
+            # Only provide completions if:
+            # 1. Input is empty, OR
+            # 2. Input exactly matches the beginning of an option
+            if text == "":
+                # Show all options when nothing is typed
+                for option in self.options:
+                    yield Completion(
+                        option,
+                        start_position=0,
+                        display_meta=self.meta_dict.get(option, ""),
+                    )
+            else:
+                # Only complete if text exactly matches the beginning of an option
+                for option in self.options:
+                    if option.startswith(text):
+                        yield Completion(
+                            option,
+                            start_position=-len(text),
+                            display_meta=self.meta_dict.get(option, ""),
+                        )
+
+    options = ["yes", "no", "edit"]
+    meta_dict = {
+        "yes": "Approve the execution",
+        "no": "Disapprove the execution",
+        "edit": "Edit tool execution parameters",
+    }
+    for key in kwargs:
+        options.append(f"edit {key}")
+        meta_dict[f"edit {key}"] = f"Edit tool execution parameter: {key}"
+    completer = ToolConfirmationCompleter(options, meta_dict)
+    reader = PromptSession()
+    return await reader.prompt_async(completer=completer)
+
+
+def _adjust_signature(wrapper: Callable, original_sig: inspect.Signature):
     """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
@@ -137,22 +358,4 @@ def _adjust_signature(
         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(
-                    "_dummy", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=None
-                )
-            ]
-        )
-    else:
-        new_sig = inspect.Signature(parameters=params_for_schema)
-
-    wrapper.__signature__ = new_sig
+    wrapper.__signature__ = inspect.Signature(parameters=params_for_schema)

zrb/task/llm/workflow.py

@@ -0,0 +1,76 @@
+import os
+
+from zrb.config.config import CFG
+from zrb.config.llm_context.config import llm_context_config
+from zrb.config.llm_context.workflow import LLMWorkflow
+
+
+def load_workflow(workflow_name: str | list[str]) -> str:
+    """
+    Loads and formats one or more workflow documents for LLM consumption.
+
+    Retrieves workflows by name, formats with descriptive headers for LLM context injection.
+
+    Args:
+        workflow_name: Name or list of names of the workflow(s) to load
+
+    Returns:
+        Formatted workflow content as a string with headers
+
+    Raises:
+        ValueError: If any specified workflow name is not found
+    """
+    names = [workflow_name] if isinstance(workflow_name, str) else workflow_name
+    available_workflows = get_available_workflows()
+    contents = []
+    for name in names:
+        workflow = available_workflows.get(name.strip().lower())
+        if workflow is None:
+            raise ValueError(f"Workflow not found: {name}")
+        contents.append(
+            "\n".join(
+                [
+                    f"# {workflow.name}",
+                    f"> Workflow Location: `{workflow.path}`",
+                    workflow.content,
+                ]
+            )
+        )
+    return "\n".join(contents)
+
+
+def get_available_workflows() -> dict[str, LLMWorkflow]:
+    available_workflows = {
+        workflow_name.strip().lower(): workflow
+        for workflow_name, workflow in llm_context_config.get_workflows().items()
+    }
+    # Define builtin workflow locations in order of precedence
+    builtin_workflow_locations = [
+        os.path.expanduser(additional_builtin_workflow_path)
+        for additional_builtin_workflow_path in CFG.LLM_BUILTIN_WORKFLOW_PATHS
+        if os.path.isdir(os.path.expanduser(additional_builtin_workflow_path))
+    ]
+    builtin_workflow_locations.append(
+        os.path.join(os.path.dirname(__file__), "default_workflow")
+    )
+    # Load workflows from all locations
+    for workflow_location in builtin_workflow_locations:
+        if not os.path.isdir(workflow_location):
+            continue
+        for workflow_name in os.listdir(workflow_location):
+            workflow_dir = os.path.join(workflow_location, workflow_name)
+            workflow_file = os.path.join(workflow_dir, "workflow.md")
+            if not os.path.isfile(workflow_file):
+                workflow_file = os.path.join(workflow_dir, "SKILL.md")
+                if not os.path.isfile(path=workflow_file):
+                    continue
+            # Only add if not already defined (earlier locations have precedence)
+            if workflow_name not in available_workflows:
+                with open(workflow_file, "r") as f:
+                    workflow_content = f.read()
+                available_workflows[workflow_name] = LLMWorkflow(
+                    name=workflow_name,
+                    path=workflow_dir,
+                    content=workflow_content,
+                )
+    return available_workflows