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

zrb/llm/util/prompt.py

@@ -0,0 +1,104 @@
+import os
+import re
+
+from zrb.util.file import list_files, read_file
+
+
+def expand_prompt(prompt: str) -> str:
+    """
+    Expands @reference patterns in the prompt into a Reference + Appendix style.
+    Example: "Check @main.py" -> "Check main.py (see Appendix)\n...[Appendix with content]..."
+    """
+    if not prompt:
+        return prompt
+
+    # Regex to capture @path.
+    # Matches @ followed by typical path chars.
+    # We'll allow alphanumeric, _, -, ., /, \, and ~ (home dir).
+    pattern = re.compile(r"@(?P<path>[\w~\-\./\\]+)")
+
+    matches = list(pattern.finditer(prompt))
+    if not matches:
+        return prompt
+
+    appendix_entries: list[str] = []
+    # We construct the new string by slicing.
+    last_idx = 0
+    parts = []
+
+    for match in matches:
+        # Add text before match
+        parts.append(prompt[last_idx : match.start()])
+
+        path_ref = match.group("path")
+        original_token = match.group(0)
+
+        # Check existence
+        expanded_path = os.path.expanduser(path_ref)
+        abs_path = os.path.abspath(expanded_path)
+
+        content = ""
+        header = ""
+        is_valid_ref = False
+
+        if os.path.isfile(abs_path):
+            try:
+                content = read_file(abs_path)
+                header = f"File Content: `{path_ref}`"
+                is_valid_ref = True
+            except Exception:
+                pass
+        elif os.path.isdir(abs_path):
+            try:
+                # Use list_files for directory structure
+                file_list = list_files(abs_path, depth=2)
+                content = "\n".join(file_list)
+                if not content:
+                    content = "(Empty directory)"
+                header = f"Directory Listing: `{path_ref}`"
+                is_valid_ref = True
+            except Exception:
+                pass
+
+        if not is_valid_ref:
+            # Fallback: leave original token if unreadable or not found
+            parts.append(original_token)
+            last_idx = match.end()
+            continue
+
+        # If we successfully got content
+        parts.append(f"`{path_ref}` (see Appendix)")
+
+        # Add to appendix with strict instructions
+        entry_lines = [
+            f"### {header}",
+            f"> **SYSTEM NOTE:** The content of `{path_ref}` is provided below.",
+            "> **DO NOT** use tools like `read_file` or `list_files` to read this path again.",
+            "> Use the content provided here directly.\n",
+            "```",
+            f"{content}",
+            "```",
+        ]
+        appendix_entries.append("\n".join(entry_lines))
+
+        last_idx = match.end()
+
+    # Add remaining text
+    parts.append(prompt[last_idx:])
+
+    new_prompt = "".join(parts)
+
+    if appendix_entries:
+        sep = "=" * 20
+        header_lines = [
+            f"\n\n{sep}APPENDIX: PRE-LOADED CONTEXT {sep}",
+            "⚠️ **SYSTEM INSTRUCTION**: The user has attached the following content.",
+            "You MUST use this provided content for your analysis.",
+            "**DO NOT** consume resources by calling `read_file` or `list_files`",
+            "or `run_shell_command` to read these specific paths again.\n",
+        ]
+        appendix_section = "\n".join(header_lines)
+        appendix_section += "\n\n".join(appendix_entries)
+        new_prompt += appendix_section
+
+    return new_prompt

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/session/any_session.py

@@ -15,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.
 

zrb/session/session.py

@@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Coroutine
 from zrb.context.any_shared_context import AnySharedContext
 from zrb.context.context import AnyContext, Context
 from zrb.group.any_group import AnyGroup
-from zrb.session.any_session import AnySession, TAnySession
+from zrb.session.any_session import AnySession
 from zrb.session_state_logger.any_session_state_logger import AnySessionStateLogger
 from zrb.session_state_logger.session_state_logger_factory import session_state_logger
 from zrb.task.any_task import AnyTask

zrb/task/base/context.py

@@ -79,24 +79,36 @@ def combine_inputs(
             input_names.append(task_input.name)  # Update names list
 
 
+def combine_envs(
+    existing_envs: list[AnyEnv],
+    new_envs: list[AnyEnv | None] | AnyEnv | None,
+):
+    """
+    Combines new envs into an existing list.
+    Modifies the existing_envs list in place.
+    """
+    if isinstance(new_envs, AnyEnv):
+        existing_envs.append(new_envs)
+    elif new_envs is None:
+        pass
+    else:
+        # new_envs is a list
+        for env in new_envs:
+            if env is not None:
+                existing_envs.append(env)
+
+
 def get_combined_envs(task: "BaseTask") -> list[AnyEnv]:
     """
     Aggregates environment variables from the task and its upstreams.
     """
-    envs = []
+    envs: list[AnyEnv] = []
     for upstream in task.upstreams:
-        envs.extend(upstream.envs)  # Use extend for list concatenation
-
-    # Access _envs directly as task is BaseTask
-    task_envs: list[AnyEnv | None] | AnyEnv | None = task._envs
-    if isinstance(task_envs, AnyEnv):
-        envs.append(task_envs)
-    elif isinstance(task_envs, list):
-        # Filter out None while extending
-        envs.extend(env for env in task_envs if env is not None)
-
-    # Filter out None values efficiently from the combined list
-    return [env for env in envs if env is not None]
+        combine_envs(envs, upstream.envs)
+
+    combine_envs(envs, task._envs)
+
+    return envs
 
 
 def get_combined_inputs(task: "BaseTask") -> list[AnyInput]:

zrb/task/base/execution.py

@@ -88,56 +88,61 @@ async def execute_action_until_ready(task: "BaseTask", session: AnySession):
         run_async(execute_action_with_retry(task, session))
     )
 
-    await asyncio.sleep(readiness_check_delay)
-
-    readiness_check_coros = [
-        run_async(check.exec_chain(session)) for check in readiness_checks
-    ]
-
-    # Wait primarily for readiness checks to complete
-    ctx.log_info("Waiting for readiness checks")
-    readiness_passed = False
     try:
-        # Gather results, but primarily interested in completion/errors
-        await asyncio.gather(*readiness_check_coros)
-        # Check if all readiness tasks actually completed successfully
-        all_readiness_completed = all(
-            session.get_task_status(check).is_completed for check in readiness_checks
-        )
-        if all_readiness_completed:
-            ctx.log_info("Readiness checks completed successfully")
-            readiness_passed = True
-            # Mark task as ready only if checks passed and action didn't fail during checks
-            if not session.get_task_status(task).is_failed:
-                ctx.log_info("Marked as ready")
-                session.get_task_status(task).mark_as_ready()
-        else:
-            ctx.log_warning(
-                "One or more readiness checks did not complete successfully."
-            )
+        await asyncio.sleep(readiness_check_delay)
 
-    except Exception as e:
-        ctx.log_error(f"Readiness check failed with exception: {e}")
-        # If readiness checks fail with an exception, the task is not ready.
-        # The action_coro might still be running or have failed.
-        # execute_action_with_retry handles marking the main task status.
-
-    # Defer the main action coroutine; it will be awaited later if needed
-    session.defer_action(task, action_coro)
-
-    # Start monitoring only if readiness passed and monitoring is enabled
-    if readiness_passed and monitor_readiness:
-        # Import dynamically to avoid circular dependency if monitoring imports execution
-        from zrb.task.base.monitoring import monitor_task_readiness
+        readiness_check_coros = [
+            run_async(check.exec_chain(session)) for check in readiness_checks
+        ]
 
-        monitor_coro = asyncio.create_task(
-            run_async(monitor_task_readiness(task, session, action_coro))
-        )
-        session.defer_monitoring(task, monitor_coro)
+        # Wait primarily for readiness checks to complete
+        ctx.log_info("Waiting for readiness checks")
+        readiness_passed = False
+        try:
+            # Gather results, but primarily interested in completion/errors
+            await asyncio.gather(*readiness_check_coros)
+            # Check if all readiness tasks actually completed successfully
+            all_readiness_completed = all(
+                session.get_task_status(check).is_completed
+                for check in readiness_checks
+            )
+            if all_readiness_completed:
+                ctx.log_info("Readiness checks completed successfully")
+                readiness_passed = True
+                # Mark task as ready only if checks passed and action didn't fail during checks
+                if not session.get_task_status(task).is_failed:
+                    ctx.log_info("Marked as ready")
+                    session.get_task_status(task).mark_as_ready()
+            else:
+                ctx.log_warning(
+                    "One or more readiness checks did not complete successfully."
+                )
+
+        except Exception as e:
+            ctx.log_error(f"Readiness check failed with exception: {e}")
+            # If readiness checks fail with an exception, the task is not ready.
+            # The action_coro might still be running or have failed.
+            # execute_action_with_retry handles marking the main task status.
+
+        # Defer the main action coroutine; it will be awaited later if needed
+        session.defer_action(task, action_coro)
+
+        # Start monitoring only if readiness passed and monitoring is enabled
+        if readiness_passed and monitor_readiness:
+            # Import dynamically to avoid circular dependency if monitoring imports execution
+            from zrb.task.base.monitoring import monitor_task_readiness
+
+            monitor_coro = asyncio.create_task(
+                run_async(monitor_task_readiness(task, session, action_coro))
+            )
+            session.defer_monitoring(task, monitor_coro)
 
-    # The result here is primarily about readiness check completion.
-    # The actual task result is handled by the deferred action_coro.
-    return None
+        # The result here is primarily about readiness check completion.
+        # The actual task result is handled by the deferred action_coro.
+        return None
+    except (asyncio.CancelledError, KeyboardInterrupt, GeneratorExit):
+        action_coro.cancel()
+        raise
 
 
 async def execute_action_with_retry(task: "BaseTask", session: AnySession) -> Any:
@@ -178,7 +183,7 @@ async def execute_action_with_retry(task: "BaseTask", session: AnySession) -> An
             await run_async(execute_successors(task, session))
             return result
 
-        except (asyncio.CancelledError, KeyboardInterrupt):
+        except (asyncio.CancelledError, KeyboardInterrupt, GeneratorExit):
             ctx.log_warning("Task cancelled or interrupted")
             session.get_task_status(task).mark_as_failed()  # Mark as failed on cancel
             # Do not trigger fallbacks/successors on cancellation

zrb/task/base/lifecycle.py

@@ -1,6 +1,7 @@
 import asyncio
 from typing import Any
 
+from zrb.context.print_fn import PrintFn
 from zrb.context.shared_context import SharedContext
 from zrb.session.any_session import AnySession
 from zrb.session.session import Session
@@ -12,6 +13,7 @@ from zrb.util.run import run_async
 async def run_and_cleanup(
     task: AnyTask,
     session: AnySession | None = None,
+    print_fn: PrintFn | None = None,
     str_kwargs: dict[str, str] | None = None,
     kwargs: dict[str, Any] | None = None,
 ) -> Any:
@@ -21,11 +23,11 @@ async def run_and_cleanup(
     """
     # Ensure a session exists
     if session is None:
-        session = Session(shared_ctx=SharedContext())
+        session = Session(shared_ctx=SharedContext(print_fn=print_fn))
 
     # Create the main task execution coroutine
     main_task_coro = asyncio.create_task(
-        run_task_async(task, session, str_kwargs, kwargs)
+        run_task_async(task, session, print_fn, str_kwargs, kwargs)
     )
 
     try:
@@ -70,6 +72,7 @@ async def run_and_cleanup(
 async def run_task_async(
     task: AnyTask,
     session: AnySession | None = None,
+    print_fn: PrintFn | None = None,
     str_kwargs: dict[str, str] | None = None,
     kwargs: dict[str, Any] | None = None,
 ) -> Any:
@@ -78,7 +81,7 @@ async def run_task_async(
     Sets up the session and initiates the root task execution chain.
     """
     if session is None:
-        session = Session(shared_ctx=SharedContext())
+        session = Session(shared_ctx=SharedContext(print_fn=print_fn))
 
     # Populate shared context with inputs and environment variables
     fill_shared_context_inputs(session.shared_ctx, task, str_kwargs, kwargs)
@@ -176,7 +179,7 @@ async def log_session_state(task: AnyTask, session: AnySession):
     try:
         while not session.is_terminated:
             session.state_logger.write(session.as_state_log())
-            await asyncio.sleep(0.1)  # Log interval
+            await asyncio.sleep(0)  # Log interval
         # Log one final time after termination signal
         session.state_logger.write(session.as_state_log())
     except asyncio.CancelledError: