zrb 1.5.5__py3-none-any.whl → 1.5.6__py3-none-any.whl

zrb/task/llm/print_node.py

@@ -0,0 +1,98 @@
+from collections.abc import Callable
+from typing import Any
+
+from pydantic_ai import Agent
+from pydantic_ai.messages import (
+    FinalResultEvent,
+    FunctionToolCallEvent,
+    FunctionToolResultEvent,
+    PartDeltaEvent,
+    PartStartEvent,
+    TextPartDelta,
+    ToolCallPartDelta,
+)
+
+from zrb.util.cli.style import stylize_faint
+
+
+async def print_node(print_func: Callable, agent_run: Any, node: Any):
+    """Prints the details of an agent execution node using a provided print function."""
+    if Agent.is_user_prompt_node(node):
+        # A user prompt node => The user has provided input
+        print_func(stylize_faint(f">> UserPromptNode: {node.user_prompt}"))
+    elif Agent.is_model_request_node(node):
+        # A model request node => We can stream tokens from the model's request
+        print_func(
+            stylize_faint(">> ModelRequestNode: streaming partial request tokens")
+        )
+        async with node.stream(agent_run.ctx) as request_stream:
+            is_streaming = False
+            async for event in request_stream:
+                if isinstance(event, PartStartEvent):
+                    if is_streaming:
+                        print_func("", plain=True)
+                    print_func(
+                        stylize_faint(
+                            f"[Request] Starting part {event.index}: {event.part!r}"
+                        ),
+                    )
+                    is_streaming = False
+                elif isinstance(event, PartDeltaEvent):
+                    if isinstance(event.delta, TextPartDelta):
+                        print_func(
+                            stylize_faint(f"{event.delta.content_delta}"),
+                            end="",
+                            plain=is_streaming,
+                        )
+                    elif isinstance(event.delta, ToolCallPartDelta):
+                        print_func(
+                            stylize_faint(f"{event.delta.args_delta}"),
+                            end="",
+                            plain=is_streaming,
+                        )
+                    is_streaming = True
+                elif isinstance(event, FinalResultEvent):
+                    if is_streaming:
+                        print_func("", plain=True)
+                    print_func(
+                        stylize_faint(f"[Result] tool_name={event.tool_name}"),
+                    )
+                    is_streaming = False
+            if is_streaming:
+                print_func("", plain=True)
+    elif Agent.is_call_tools_node(node):
+        # A handle-response node => The model returned some data, potentially calls a tool
+        print_func(
+            stylize_faint(">> CallToolsNode: streaming partial response & tool usage")
+        )
+        async with node.stream(agent_run.ctx) as handle_stream:
+            async for event in handle_stream:
+                if isinstance(event, FunctionToolCallEvent):
+                    # Handle empty arguments across different providers
+                    if event.part.args == "" or event.part.args is None:
+                        event.part.args = {}
+                    elif isinstance(
+                        event.part.args, str
+                    ) and event.part.args.strip() in ["null", "{}"]:
+                        # Some providers might send "null" or "{}" as a string
+                        event.part.args = {}
+                    # Handle dummy property if present (from our schema sanitization)
+                    if (
+                        isinstance(event.part.args, dict)
+                        and "_dummy" in event.part.args
+                    ):
+                        del event.part.args["_dummy"]
+                    print_func(
+                        stylize_faint(
+                            f"[Tools] The LLM calls tool={event.part.tool_name!r} with args={event.part.args} (tool_call_id={event.part.tool_call_id!r})"  # noqa
+                        )
+                    )
+                elif isinstance(event, FunctionToolResultEvent):
+                    print_func(
+                        stylize_faint(
+                            f"[Tools] Tool call {event.tool_call_id!r} returned => {event.result.content}"  # noqa
+                        )
+                    )
+    elif Agent.is_end_node(node):
+        # Once an End node is reached, the agent run is complete
+        print_func(stylize_faint(f"{agent_run.result.data}"))

zrb/task/llm/tool_wrapper.py

@@ -0,0 +1,88 @@
+import functools
+import inspect
+import traceback
+from collections.abc import Callable
+
+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.
+
+    - 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.
+
+    Args:
+        func: The tool function (sync or async) to wrap.
+
+    Returns:
+        An async wrapper function.
+    """
+    original_sig = inspect.signature(func)
+    needs_ctx = "ctx" in original_sig.parameters
+    takes_no_args = len(original_sig.parameters) == 0
+
+    @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")
+
+        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))
+        except Exception as e:
+            error_model = ToolExecutionError(
+                tool_name=func.__name__,
+                error_type=type(e).__name__,
+                message=str(e),
+                details=traceback.format_exc(),
+            )
+            return error_model.model_dump_json()
+
+    if takes_no_args:
+        # Inject dummy parameter for schema generation if original func took no args
+        new_sig = inspect.Signature(
+            parameters=[
+                inspect.Parameter(
+                    "_dummy", inspect.Parameter.POSITIONAL_OR_KEYWORD, default=None
+                )
+            ]
+        )
+        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
+
+    return wrapper