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

zrb/task/llm/tool_wrapper.py

@@ -1,45 +1,60 @@
 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
 
 
-def wrap_tool(func: Callable, ctx: AnyContext, is_yolo_mode: bool) -> "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)
     needs_run_context_for_pydantic = _has_context_parameter(original_sig, RunContext)
-    wrapper = wrap_func(func, ctx, is_yolo_mode)
+    wrapper = wrap_func(func, ctx, yolo_mode)
     return Tool(wrapper, takes_ctx=needs_run_context_for_pydantic)
 
 
-def wrap_func(func: Callable, ctx: AnyContext, is_yolo_mode: bool) -> Callable:
+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, is_yolo_mode
+        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, takes_no_args)
+    _adjust_signature(wrapper, original_sig)
     return wrapper
 
 
@@ -74,10 +89,10 @@ def _is_annotated_with_context(param_annotation, context_type):
 
 def _create_wrapper(
     func: Callable,
-    original_sig: inspect.Signature,
+    original_signature: inspect.Signature,
     ctx: AnyContext,
     needs_any_context_for_injection: bool,
-    is_yolo_mode: bool,
+    yolo_mode: bool | list[str],
 ) -> Callable:
     """Creates the core wrapper function."""
 
@@ -86,7 +101,7 @@ def _create_wrapper(
         # 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
@@ -99,32 +114,167 @@ 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:
-            if not is_yolo_mode and not ctx.is_web_mode and ctx.is_tty:
-                approval, reason = await _ask_for_approval(ctx, func, args, kwargs)
-                if not approval:
-                    raise ValueError(f"User disapproving: {reason}")
-            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
 
 
-async def _ask_for_approval(
-    ctx: AnyContext, func: Callable, args: list[Any], kwargs: dict[str, Any]
-) -> tuple[bool, str]:
+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 = []
@@ -133,58 +283,67 @@ async def _ask_for_approval(
         normalized_kwargs.append(
             f"{stylize_yellow(key)}={stylize_green(truncated_val)}"
         )
-    func_param_str = ",".join(normalized_args + normalized_kwargs)
-    func_call_str = (
-        f"{stylize_blue(func_name + '(')}{func_param_str}{stylize_blue(')')}"
-    )
-    while True:
-        ctx.print(
-            f"\n✅ >> Allow to run tool: {func_call_str} (Yes | No, <reason>) ",
-            plain=True,
-        )
-        user_input = await _read_line()
-        ctx.print("", plain=True)
-        user_responses = [val.strip() for val in user_input.split(",", maxsplit=1)]
-        while len(user_responses) < 2:
-            user_responses.append("")
-        approval_str, reason = user_responses
-        try:
-            approved = True if approval_str.strip() == "" else to_boolean(approval_str)
-            if not approved and reason == "":
-                ctx.print(
-                    stylize_error(
-                        f"You must specify rejection reason (i.e., No, <why>) for {func_call_str}"  # noqa
-                    ),
-                    plain=True,
-                )
-                continue
-            return approved, reason
-        except Exception:
-            ctx.print(
-                stylize_error(
-                    f"Invalid approval value for {func_call_str}: {approval_str}"
-                ),
-                plain=True,
-            )
-            continue
+    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:
-    if len(arg) > length:
-        return f"{arg[:length-4]} ..."
-    return arg
+    normalized_arg = arg.replace("\n", "\\n")
+    if len(normalized_arg) > length:
+        return f"{normalized_arg[:length-4]} ..."
+    return normalized_arg
 
 
-async def _read_line():
+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()
+    return await reader.prompt_async(completer=completer)
 
 
-def _adjust_signature(
-    wrapper: Callable, original_sig: inspect.Signature, takes_no_args: bool
-):
+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
@@ -199,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