zrb 1.15.3__py3-none-any.whl → 2.0.0a4__py3-none-any.whl

zrb/llm/util/stream_response.py

@@ -0,0 +1,178 @@
+import json
+from typing import TYPE_CHECKING, Any, Callable
+
+from zrb.context.any_context import AnyContext
+from zrb.util.cli.style import stylize_faint
+
+if TYPE_CHECKING:
+    from pydantic_ai import AgentStreamEvent
+
+
+def create_event_handler(
+    print_event: Callable[..., None],
+    indent_level: int = 1,
+    show_tool_call_detail: bool = False,
+    show_tool_result: bool = False,
+):
+    from pydantic_ai import (
+        AgentRunResultEvent,
+        FinalResultEvent,
+        FunctionToolCallEvent,
+        FunctionToolResultEvent,
+        PartDeltaEvent,
+        PartStartEvent,
+        TextPartDelta,
+        ThinkingPartDelta,
+        ToolCallPartDelta,
+    )
+
+    indentation = indent_level * 2 * " "
+    progress_char_list = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
+    progress_char_index = 0
+    was_tool_call_delta = False
+    event_prefix = indentation
+
+    def fprint(content: str, preserve_leading_newline: bool = False):
+        if preserve_leading_newline and content.startswith("\n"):
+            return print_event("\n" + content[1:].replace("\n", f"\n{indentation}   "))
+        return print_event(content.replace("\n", f"\n{indentation}   "))
+
+    async def handle_event(event: "AgentStreamEvent"):
+        from pydantic_ai import ToolCallPart
+
+        nonlocal progress_char_index, was_tool_call_delta, event_prefix
+        if isinstance(event, PartStartEvent):
+            # Skip ToolCallPart start, we handle it in Deltas/CallEvent
+            if isinstance(event.part, ToolCallPart):
+                return
+            content = _get_event_part_content(event)
+            # Use preserve_leading_newline=True because event_prefix contains the correctly indented newline
+            fprint(f"{event_prefix}🧠 {content}", preserve_leading_newline=True)
+            was_tool_call_delta = False
+        elif isinstance(event, PartDeltaEvent):
+            if isinstance(event.delta, TextPartDelta):
+                # Standard fprint for deltas to ensure wrapping indentation
+                fprint(f"{event.delta.content_delta}")
+                was_tool_call_delta = False
+            elif isinstance(event.delta, ThinkingPartDelta):
+                fprint(f"{event.delta.content_delta}")
+                was_tool_call_delta = False
+            elif isinstance(event.delta, ToolCallPartDelta):
+                if show_tool_call_detail:
+                    fprint(f"{event.delta.args_delta}")
+                else:
+                    progress_char = progress_char_list[progress_char_index]
+                    if not was_tool_call_delta:
+                        # Print newline for tool param spinner
+                        fprint("\n")
+
+                    # Split \r to avoid UI._append_to_output stripping the ANSI start code along with the line
+                    print_event("\r")
+                    print_event(
+                        f"{indentation}🔄 Prepare tool parameters {progress_char}"
+                    )
+                    progress_char_index += 1
+                    if progress_char_index >= len(progress_char_list):
+                        progress_char_index = 0
+                    was_tool_call_delta = True
+        elif isinstance(event, FunctionToolCallEvent):
+            args = _get_truncated_event_part_args(event)
+            # Use preserve_leading_newline=True for the block header
+            fprint(
+                f"{event_prefix}🧰 {event.part.tool_call_id} | {event.part.tool_name} {args}",
+                preserve_leading_newline=True,
+            )
+            was_tool_call_delta = False
+        elif isinstance(event, FunctionToolResultEvent):
+            if show_tool_result:
+                fprint(
+                    f"{event_prefix}🔠 {event.tool_call_id} | Return {event.result.content}",
+                    preserve_leading_newline=True,
+                )
+            else:
+                fprint(
+                    f"{event_prefix}🔠 {event.tool_call_id} Executed",
+                    preserve_leading_newline=True,
+                )
+            was_tool_call_delta = False
+        elif isinstance(event, AgentRunResultEvent):
+            usage = event.result.usage()
+            usage_msg = " ".join(
+                [
+                    "💸",
+                    f"(Requests: {usage.requests} |",
+                    f"Tool Calls: {usage.tool_calls} |",
+                    f"Total: {usage.total_tokens})",
+                    f"Input: {usage.input_tokens} |",
+                    f"Audio Input: {usage.input_audio_tokens} |",
+                    f"Output: {usage.output_tokens} |",
+                    f"Audio Output: {usage.output_audio_tokens} |",
+                    f"Cache Read: {usage.cache_read_tokens} |",
+                    f"Cache Write: {usage.cache_write_tokens} |",
+                    f"Details: {usage.details}",
+                ]
+            )
+            fprint(
+                f"{event_prefix}{stylize_faint(usage_msg)}\n",
+                preserve_leading_newline=True,
+            )
+            was_tool_call_delta = False
+        elif isinstance(event, FinalResultEvent):
+            was_tool_call_delta = False
+        event_prefix = f"\n{indentation}"
+
+    return handle_event
+
+
+def create_faint_printer(ctx: AnyContext):
+    def faint_print(*values: object):
+        message = stylize_faint(" ".join([f"{value}" for value in values]))
+        ctx.print(message, end="", plain=True)
+
+    return faint_print
+
+
+def _get_truncated_event_part_args(event: "AgentStreamEvent") -> Any:
+    # Handle empty arguments across different providers
+    if not hasattr(event, "part"):
+        return {}
+    part = getattr(event, "part")
+    if not hasattr(part, "args"):
+        return {}
+    args = getattr(part, "args")
+    if args == "" or args is None:
+        return {}
+    if isinstance(args, str):
+        # Some providers might send "null" or "{}" as a string
+        if args.strip() in ["null", "{}"]:
+            return {}
+        try:
+            obj = json.loads(args)
+            if isinstance(obj, dict):
+                return _truncate_kwargs(obj)
+        except json.JSONDecodeError:
+            pass
+    # Handle dummy property if present (from our schema sanitization)
+    if isinstance(args, dict):
+        return _truncate_kwargs(args)
+    return args
+
+
+def _truncate_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
+    return {key: _truncate_arg(val) for key, val in kwargs.items()}
+
+
+def _truncate_arg(arg: str, length: int = 30) -> str:
+    if isinstance(arg, str) and len(arg) > length:
+        return f"{arg[:length-4]} ..."
+    return arg
+
+
+def _get_event_part_content(event: "AgentStreamEvent") -> str:
+    if not hasattr(event, "part"):
+        return ""
+    part = getattr(event, "part")
+    if hasattr(part, "content"):
+        return getattr(part, "content")
+    # For parts without content (like ToolCallPart, though we skip it now), return empty or simple repr
+    return ""

zrb/runner/cli.py

@@ -7,7 +7,7 @@ from zrb.context.any_context import AnyContext
 from zrb.context.shared_context import SharedContext
 from zrb.group.any_group import AnyGroup
 from zrb.group.group import Group
-from zrb.runner.common_util import get_run_kwargs
+from zrb.runner.common_util import get_task_str_kwargs
 from zrb.session.session import Session
 from zrb.session_state_logger.session_state_logger_factory import session_state_logger
 from zrb.task.any_task import AnyTask
@@ -38,23 +38,25 @@ class Cli(Group):
     def banner(self) -> str:
         return CFG.BANNER
 
-    def run(self, args: list[str] = []):
-        kwargs, args = self._extract_kwargs_from_args(args)
-        node, node_path, args = extract_node_from_args(self, args)
+    def run(self, str_args: list[str] = []):
+        str_kwargs, str_args = self._extract_kwargs_from_args(str_args)
+        node, node_path, str_args = extract_node_from_args(self, str_args)
         if isinstance(node, AnyGroup):
             self._show_group_info(node)
             return
-        if "h" in kwargs or "help" in kwargs:
+        if "h" in str_kwargs or "help" in str_kwargs:
             self._show_task_info(node)
             return
-        run_kwargs = get_run_kwargs(task=node, args=args, kwargs=kwargs, cli_mode=True)
+        task_str_kwargs = get_task_str_kwargs(
+            task=node, str_args=str_args, str_kwargs=str_kwargs, cli_mode=True
+        )
         try:
-            result = self._run_task(node, args, run_kwargs)
+            result = self._run_task(node, str_args, task_str_kwargs)
             if result is not None:
                 print(result)
             return result
         finally:
-            run_command = self._get_run_command(node_path, run_kwargs)
+            run_command = self._get_run_command(node_path, task_str_kwargs)
             self._print_run_command(run_command)
 
     def _print_run_command(self, run_command: str):
@@ -64,11 +66,14 @@ class Cli(Group):
             file=sys.stderr,
         )
 
-    def _get_run_command(self, node_path: list[str], run_kwargs: dict[str, str]) -> str:
+    def _get_run_command(
+        self, node_path: list[str], task_str_kwargs: dict[str, str]
+    ) -> str:
         parts = [self.name] + node_path
-        if len(run_kwargs) > 0:
+        if len(task_str_kwargs) > 0:
             parts += [
-                self._get_run_command_param(key, val) for key, val in run_kwargs.items()
+                self._get_run_command_param(key, val)
+                for key, val in task_str_kwargs.items()
             ]
         return " ".join(parts)
 
@@ -81,13 +86,9 @@ class Cli(Group):
         self, task: AnyTask, args: list[str], run_kwargs: dict[str, str]
     ) -> tuple[Any]:
         shared_ctx = SharedContext(args=args)
-        for task_input in task.inputs:
-            if task_input.name in run_kwargs:
-                task_input.update_shared_context(
-                    shared_ctx, run_kwargs[task_input.name]
-                )
-                continue
-        return task.run(Session(shared_ctx=shared_ctx, root_group=self))
+        return task.run(
+            Session(shared_ctx=shared_ctx, root_group=self), str_kwargs=run_kwargs
+        )
 
     def _show_task_info(self, task: AnyTask):
         description = task.description
@@ -150,11 +151,11 @@ class Cli(Group):
                         kwargs[key] = args[i + 1]
                         i += 1  # Skip the next argument as it's a value
                     else:
-                        kwargs[key] = True
+                        kwargs[key] = "true"
             elif arg.startswith("-"):
                 # Handle short flags like -t or -n
                 key = arg[1:]
-                kwargs[key] = True
+                kwargs[key] = "true"
             else:
                 # Anything else is considered a positional argument
                 residual_args.append(arg)

zrb/runner/common_util.py

@@ -1,31 +1,36 @@
-from typing import Any
-
 from zrb.context.shared_context import SharedContext
 from zrb.task.any_task import AnyTask
 
 
-def get_run_kwargs(
-    task: AnyTask, args: list[str], kwargs: dict[str, str], cli_mode: bool
+def get_task_str_kwargs(
+    task: AnyTask, str_args: list[str], str_kwargs: dict[str, str], cli_mode: bool
 ) -> dict[str, str]:
     arg_index = 0
-    str_kwargs = {key: f"{val}" for key, val in kwargs.items()}
-    run_kwargs = {**str_kwargs}
-    shared_ctx = SharedContext(args=args)
+    dummmy_shared_ctx = SharedContext()
+    task_str_kwargs = {}
     for task_input in task.inputs:
+        task_name = task_input.name
         if task_input.name in str_kwargs:
-            # Update shared context for next input default value
-            task_input.update_shared_context(shared_ctx, str_kwargs[task_input.name])
-        elif arg_index < len(args) and task_input.allow_positional_parsing:
-            run_kwargs[task_input.name] = args[arg_index]
-            # Update shared context for next input default value
-            task_input.update_shared_context(shared_ctx, run_kwargs[task_input.name])
+            task_str_kwargs[task_input.name] = str_kwargs[task_name]
+            # Update dummy shared context for next input default value
+            task_input.update_shared_context(
+                dummmy_shared_ctx, str_value=str_kwargs[task_name]
+            )
+        elif arg_index < len(str_args) and task_input.allow_positional_parsing:
+            task_str_kwargs[task_name] = str_args[arg_index]
+            # Update dummy shared context for next input default value
+            task_input.update_shared_context(
+                dummmy_shared_ctx, str_value=task_str_kwargs[task_name]
+            )
             arg_index += 1
         else:
             if cli_mode and task_input.always_prompt:
-                str_value = task_input.prompt_cli_str(shared_ctx)
+                str_value = task_input.prompt_cli_str(dummmy_shared_ctx)
             else:
-                str_value = task_input.get_default_str(shared_ctx)
-            run_kwargs[task_input.name] = str_value
-            # Update shared context for next input default value
-            task_input.update_shared_context(shared_ctx, run_kwargs[task_input.name])
-    return run_kwargs
+                str_value = task_input.get_default_str(dummmy_shared_ctx)
+            task_str_kwargs[task_name] = str_value
+            # Update dummy shared context for next input default value
+            task_input.update_shared_context(
+                dummmy_shared_ctx, str_value=task_str_kwargs[task_name]
+            )
+    return task_str_kwargs

zrb/runner/web_route/task_input_api_route.py

@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING
 
 from zrb.config.web_auth_config import WebAuthConfig
 from zrb.group.any_group import AnyGroup
-from zrb.runner.common_util import get_run_kwargs
+from zrb.runner.common_util import get_task_str_kwargs
 from zrb.runner.web_util.user import get_user_from_request
 from zrb.task.any_task import AnyTask
 from zrb.util.group import NodeNotFoundError, extract_node_from_args
@@ -39,9 +39,9 @@ def serve_task_input_api(
         if isinstance(task, AnyTask):
             if not user.can_access_task(task):
                 return JSONResponse(content={"detail": "Forbidden"}, status_code=403)
-            query_dict = json.loads(query)
-            run_kwargs = get_run_kwargs(
-                task=task, args=[], kwargs=query_dict, cli_mode=False
+            str_kwargs = json.loads(query)
+            task_str_kwargs = get_task_str_kwargs(
+                task=task, str_args=[], str_kwargs=str_kwargs, cli_mode=False
             )
-            return run_kwargs
+            return task_str_kwargs
         return JSONResponse(content={"detail": "Not found"}, status_code=404)

zrb/runner/web_util/user.py

@@ -19,7 +19,7 @@ def get_user_by_credentials(
 
 async def get_user_from_request(
     web_auth_config: WebAuthConfig, request: "Request"
-) -> User | None:
+) -> User:
     from fastapi.security import OAuth2PasswordBearer
 
     if not web_auth_config.enable_auth:
@@ -45,7 +45,11 @@ def _get_user_from_cookie(
     return None
 
 
-def _get_user_from_token(web_auth_config: WebAuthConfig, token: str) -> User | None:
+def _get_user_from_token(
+    web_auth_config: WebAuthConfig, token: str | None
+) -> User | None:
+    if token is None:
+        return None
     try:
         from jose import jwt
 
@@ -54,7 +58,7 @@ def _get_user_from_token(web_auth_config: WebAuthConfig, token: str) -> User | N
             web_auth_config.secret_key,
             options={"require_sub": True, "require_exp": True},
         )
-        username: str = payload.get("sub")
+        username: str | None = payload.get("sub")
         if username is None:
             return None
         user = web_auth_config.find_user_by_username(username)

zrb/session/any_session.py

@@ -1,5 +1,6 @@
 from __future__ import annotations  # Enables forward references
 
+import asyncio
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Coroutine, TypeVar
 
@@ -14,9 +15,6 @@ if TYPE_CHECKING:
     from zrb.task.any_task import AnyTask
 
 
-TAnySession = TypeVar("TAnySession", bound="AnySession")
-
-
 class AnySession(ABC):
     """Abstract base class for managing task execution and context in a session.
 
@@ -62,12 +60,13 @@ class AnySession(ABC):
 
     @property
     @abstractmethod
-    def parent(self) -> TAnySession | None:
+    def parent(self) -> "AnySession | None":
         """Parent session"""
         pass
 
+    @property
     @abstractmethod
-    def task_path(self) -> str:
+    def task_path(self) -> list[str]:
         """Main task's path"""
         pass
 
@@ -105,7 +104,9 @@ class AnySession(ABC):
         pass
 
     @abstractmethod
-    def defer_monitoring(self, task: "AnyTask", coro: Coroutine):
+    def defer_monitoring(
+        self, task: "AnyTask", coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]
+    ):
         """Defers the execution of a task's monitoring coroutine for later processing.
 
         Args:
@@ -115,7 +116,9 @@ class AnySession(ABC):
         pass
 
     @abstractmethod
-    def defer_action(self, task: "AnyTask", coro: Coroutine):
+    def defer_action(
+        self, task: "AnyTask", coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]
+    ):
         """Defers the execution of a task's coroutine for later processing.
 
         Args:
@@ -125,7 +128,7 @@ class AnySession(ABC):
         pass
 
     @abstractmethod
-    def defer_coro(self, coro: Coroutine):
+    def defer_coro(self, coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]):
         """Defers the execution of a coroutine for later processing.
 
         Args:
@@ -185,7 +188,7 @@ class AnySession(ABC):
         pass
 
     @abstractmethod
-    def is_allowed_to_run(self, task: "AnyTask"):
+    def is_allowed_to_run(self, task: "AnyTask") -> bool:
         """Determines if the specified task is allowed to run based on its current state.
 
         Args:

zrb/session/session.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import asyncio
 from typing import TYPE_CHECKING, Any, Coroutine
 
@@ -48,10 +50,10 @@ class Session(AnySession):
         self._context: dict[AnyTask, Context] = {}
         self._shared_ctx = shared_ctx
         self._shared_ctx.set_session(self)
-        self._parent = parent
-        self._action_coros: dict[AnyTask, asyncio.Task] = {}
-        self._monitoring_coros: dict[AnyTask, asyncio.Task] = {}
-        self._coros: list[asyncio.Task] = []
+        self._parent: AnySession | None = parent
+        self._action_coros: dict[AnyTask, asyncio.Task[Any]] = {}
+        self._monitoring_coros: dict[AnyTask, asyncio.Task[Any]] = {}
+        self._coros: list[asyncio.Task[Any]] = []
         self._colors = [
             GREEN,
             YELLOW,
@@ -114,11 +116,13 @@ class Session(AnySession):
         return self._parent
 
     @property
-    def task_path(self) -> str:
+    def task_path(self) -> list[str]:
         return self._main_task_path
 
     @property
     def final_result(self) -> Any:
+        if self._main_task is None:
+            return None
         xcom: Xcom = self.shared_ctx.xcom[self._main_task.name]
         try:
             return xcom.peek()
@@ -134,7 +138,11 @@ class Session(AnySession):
     def set_main_task(self, main_task: AnyTask):
         self.register_task(main_task)
         self._main_task = main_task
-        main_task_path = get_node_path(self._root_group, main_task)
+        main_task_path = (
+            None
+            if self._root_group is None
+            else get_node_path(self._root_group, main_task)
+        )
         self._main_task_path = [] if main_task_path is None else main_task_path
 
     def as_state_log(self) -> "SessionStateLog":
@@ -171,7 +179,7 @@ class Session(AnySession):
         return SessionStateLog(
             name=self.name,
             start_time=log_start_time,
-            main_task_name=self._main_task.name,
+            main_task_name="" if self._main_task is None else self._main_task.name,
             path=self.task_path,
             final_result=(
                 remove_style(f"{self.final_result}")
@@ -188,16 +196,29 @@ class Session(AnySession):
         self._register_single_task(task)
         return self._context[task]
 
-    def defer_monitoring(self, task: AnyTask, coro: Coroutine):
+    def defer_monitoring(
+        self, task: AnyTask, coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]
+    ):
         self._register_single_task(task)
-        self._monitoring_coros[task] = coro
+        if isinstance(coro, asyncio.Task):
+            self._monitoring_coros[task] = coro
+        else:
+            self._monitoring_coros[task] = asyncio.create_task(coro)
 
-    def defer_action(self, task: AnyTask, coro: Coroutine):
+    def defer_action(
+        self, task: AnyTask, coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]
+    ):
         self._register_single_task(task)
-        self._action_coros[task] = coro
+        if isinstance(coro, asyncio.Task):
+            self._action_coros[task] = coro
+        else:
+            self._action_coros[task] = asyncio.create_task(coro)
 
-    def defer_coro(self, coro: Coroutine):
-        self._coros.append(coro)
+    def defer_coro(self, coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]):
+        if isinstance(coro, asyncio.Task):
+            self._coros.append(coro)
+        else:
+            self._coros.append(asyncio.create_task(coro))
         self._coros = [
             existing_coro for existing_coro in self._coros if not existing_coro.done()
         ]
@@ -246,15 +267,15 @@ class Session(AnySession):
 
     def get_next_tasks(self, task: AnyTask) -> list[AnyTask]:
         self._register_single_task(task)
-        return self._downstreams.get(task)
+        return self._downstreams.get(task, [])
 
     def get_task_status(self, task: AnyTask) -> TaskStatus:
         self._register_single_task(task)
         return self._task_status[task]
 
     def _register_single_task(self, task: AnyTask):
-        if task.name not in self._shared_ctx._xcom:
-            self._shared_ctx._xcom[task.name] = Xcom([])
+        if task.name not in self._shared_ctx.xcom:
+            self._shared_ctx.xcom[task.name] = Xcom([])
         if task not in self._context:
             self._context[task] = Context(
                 shared_ctx=self._shared_ctx,
@@ -278,7 +299,7 @@ class Session(AnySession):
             self._color_index = 0
         return chosen
 
-    def _get_icon(self, task: AnyTask) -> int:
+    def _get_icon(self, task: AnyTask) -> str:
         if task.icon is not None:
             return task.icon
         chosen = self._icons[self._icon_index]

zrb/task/any_task.py

@@ -1,7 +1,7 @@
 from __future__ import annotations  # Enables forward references
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Callable
 
 from zrb.env.any_env import AnyEnv
 from zrb.input.any_input import AnyInput
@@ -36,6 +36,14 @@ class AnyTask(ABC):
     the actual implementation for these abstract members.
     """
 
+    @abstractmethod
+    def __rshift__(self, other: "AnyTask | list[AnyTask]") -> "AnyTask | list[AnyTask]":
+        pass
+
+    @abstractmethod
+    def __lshift__(self, other: "AnyTask | list[AnyTask]") -> "AnyTask":
+        pass
+
     @property
     @abstractmethod
     def name(self) -> str:
@@ -148,13 +156,17 @@ class AnyTask(ABC):
 
     @abstractmethod
     def run(
-        self, session: "AnySession | None" = None, str_kwargs: dict[str, str] = {}
+        self,
+        session: "AnySession | None" = None,
+        str_kwargs: dict[str, str] | None = None,
+        kwargs: dict[str, Any] | None = None,
     ) -> Any:
         """Runs the task synchronously.
 
         Args:
             session (AnySession): The shared session.
             str_kwargs(dict[str, str]): The input string values.
+            kwargs(dict[str, Any]): The input values.
 
         Returns:
             Any: The result of the task execution.
@@ -163,13 +175,17 @@ class AnyTask(ABC):
 
     @abstractmethod
     async def async_run(
-        self, session: "AnySession | None" = None, str_kwargs: dict[str, str] = {}
+        self,
+        session: "AnySession | None" = None,
+        str_kwargs: dict[str, str] | None = None,
+        kwargs: dict[str, Any] | None = None,
     ) -> Any:
         """Runs the task asynchronously.
 
         Args:
             session (AnySession): The shared session.
             str_kwargs(dict[str, str]): The input string values.
+            kwargs(dict[str, Any]): The input values.
 
         Returns:
             Any: The result of the task execution.
@@ -203,3 +219,8 @@ class AnyTask(ABC):
             session (AnySession): The shared session.
         """
         pass
+
+    @abstractmethod
+    def to_function(self) -> Callable[..., Any]:
+        """Turn a task into a function"""
+        pass