zrb 1.5.8__tar.gz → 1.5.9__tar.gz

{zrb-1.5.8 → zrb-1.5.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.5.8
+Version: 1.5.9
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -26,7 +26,7 @@ Requires-Dist: openai (>=1.70.0,<2.0.0) ; extra == "rag" or extra == "all"
 Requires-Dist: pdfplumber (>=0.11.6,<0.12.0) ; extra == "rag" or extra == "all"
 Requires-Dist: playwright (>=1.51.0,<2.0.0) ; extra == "playwright" or extra == "all"
 Requires-Dist: psutil (>=7.0.0,<8.0.0)
-Requires-Dist: pydantic-ai (>=0.0.52,<0.0.53)
+Requires-Dist: pydantic-ai (>=0.1.2,<0.2.0)
 Requires-Dist: python-dotenv (>=1.1.0,<2.0.0)
 Requires-Dist: python-jose[cryptography] (>=3.4.0,<4.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)

{zrb-1.5.8 → zrb-1.5.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "zrb"
-version = "1.5.8"
+version = "1.5.9"
 description = "Your Automation Powerhouse"
 authors = ["Go Frendi Gunawan <gofrendiasgard@gmail.com>"]
 license = "AGPL-3.0-or-later"
@@ -49,7 +49,7 @@ beautifulsoup4 = "^4.13.3"
 fastapi = {extras = ["standard"], version = "^0.115.12"}
 python-jose = {extras = ["cryptography"], version = "^3.4.0"}
 ulid-py = "^1.1.0"
-pydantic-ai = "^0.0.52"
+pydantic-ai = "^0.1.2"
 openai = {version = "^1.70.0", optional = true}
 playwright = {version = "^1.51.0", optional = true}
 psutil = "^7.0.0"

{zrb-1.5.8 → zrb-1.5.9}/src/zrb/llm_config.py

@@ -73,6 +73,9 @@ class LLMConfig:
         default_system_prompt: str | None = None,
         default_summarization_prompt: str | None = None,
         default_context_enrichment_prompt: str | None = None,
+        default_summarize_history: bool | None = None,
+        default_history_summarization_threshold: int | None = None,
+        default_enrich_context: bool | None = None,
     ):
         self._default_model_name = (
             default_model_name
@@ -109,6 +112,21 @@ class LLMConfig:
             if default_context_enrichment_prompt is not None
             else os.getenv("ZRB_LLM_CONTEXT_ENRICHMENT_PROMPT", None)
         )
+        self._default_summarize_history = (
+            default_summarize_history
+            if default_summarize_history is not None
+            else os.getenv("ZRB_LLM_SUMMARIZE_HISTORY", "true").lower() == "true"
+        )
+        self._default_history_summarization_threshold = (
+            default_history_summarization_threshold
+            if default_history_summarization_threshold is not None
+            else int(os.getenv("ZRB_LLM_HISTORY_SUMMARIZATION_THRESHOLD", "5"))
+        )
+        self._default_enrich_context = (
+            default_enrich_context
+            if default_enrich_context is not None
+            else os.getenv("ZRB_LLM_ENRICH_CONTEXT", "false").lower() == "true"
+        )
         self._default_provider = None
         self._default_model = None
 
@@ -164,6 +182,15 @@ class LLMConfig:
             provider=self.get_default_model_provider(),
         )
 
+    def get_default_summarize_history(self) -> bool:
+        return self._default_summarize_history
+
+    def get_default_history_summarization_threshold(self) -> int:
+        return self._default_history_summarization_threshold
+
+    def get_default_enrich_context(self) -> bool:
+        return self._default_enrich_context
+
     def set_default_persona(self, persona: str):
         self._default_persona = persona
 
@@ -191,5 +218,16 @@ class LLMConfig:
     def set_default_model(self, model: Model | str | None):
         self._default_model = model
 
+    def set_default_summarize_history(self, summarize_history: bool):
+        self._default_summarize_history = summarize_history
+
+    def set_default_history_summarization_threshold(
+        self, history_summarization_threshold: int
+    ):
+        self._default_history_summarization_threshold = history_summarization_threshold
+
+    def set_default_enrich_context(self, enrich_context: bool):
+        self._default_enrich_context = enrich_context
+
 
 llm_config = LLMConfig()

{zrb-1.5.8 → zrb-1.5.9}/src/zrb/task/any_task.py

@@ -12,12 +12,28 @@ if TYPE_CHECKING:
 
 
 class AnyTask(ABC):
-    """Abstract base class representing a task in the system.
-
-    This class defines the interface for tasks, including methods for retrieving
-    task metadata (such as name, color, and description), handling inputs and
-    environment variables, and managing task dependencies (upstreams and fallbacks).
-    It also includes methods for running tasks synchronously or asynchronously.
+    """Abstract base class defining the interface for all executable tasks within Zrb.
+
+    This interface specifies the essential properties and methods that any concrete
+    task implementation must provide. It covers:
+
+    - **Metadata:** Properties like `name`, `description`, `color`, and `icon` for
+      identification and UI representation.
+    - **Configuration:** Properties for defining `inputs` (user-provided parameters)
+      and `envs` (environment variables).
+    - **Execution Control:** Properties like `cli_only` to restrict execution context.
+    - **Dependency Management:** Properties for `upstreams` (prerequisites),
+      `fallbacks` (error handling), `successors` (post-completion tasks), and
+      `readiness_checks` (for long-running tasks). Methods like `append_*` allow
+      programmatic modification of these dependencies.
+    - **Context Management:** `get_ctx` method to retrieve the task's specific
+      execution context.
+    - **Execution Methods:** `run` (synchronous) and `async_run` (asynchronous)
+      entry points, along with internal execution helpers (`exec_root_tasks`,
+      `exec_chain`, `exec`) defining the execution lifecycle.
+
+    Concrete task classes (like `BaseTask`) inherit from `AnyTask` and provide
+    the actual implementation for these abstract members.
     """
 
     @property

zrb-1.5.9/src/zrb/task/base/context.py

@@ -0,0 +1,108 @@
+import os
+from typing import TYPE_CHECKING
+
+from zrb.context.any_context import AnyContext
+from zrb.context.any_shared_context import AnySharedContext
+from zrb.env.any_env import AnyEnv
+from zrb.input.any_input import AnyInput
+from zrb.session.any_session import AnySession
+from zrb.task.any_task import AnyTask
+
+if TYPE_CHECKING:
+    # Keep BaseTask under TYPE_CHECKING for functions that need it
+    from zrb.task.base_task import BaseTask
+
+
+def build_task_context(task: AnyTask, session: AnySession) -> AnyContext:
+    """
+    Retrieves the context for the task from the session and enhances it
+    with the task's specific environment variables.
+    """
+    ctx = session.get_ctx(task)
+    # Enhance session ctx with current task env
+    for env in task.envs:
+        env.update_context(ctx)
+    return ctx
+
+
+def fill_shared_context_inputs(
+    task: AnyTask, shared_context: AnySharedContext, str_kwargs: dict[str, str] = {}
+):
+    """
+    Populates the shared context with input values provided via kwargs.
+    """
+    for task_input in task.inputs:
+        if task_input.name not in shared_context.input:
+            str_value = str_kwargs.get(task_input.name, None)
+            task_input.update_shared_context(shared_context, str_value)
+
+
+def fill_shared_context_envs(shared_context: AnySharedContext):
+    """
+    Injects OS environment variables into the shared context if they don't already exist.
+    """
+    os_env_map = {
+        key: val for key, val in os.environ.items() if key not in shared_context.env
+    }
+    shared_context.env.update(os_env_map)
+
+
+def combine_inputs(
+    existing_inputs: list[AnyInput],
+    new_inputs: list[AnyInput | None] | AnyInput | None,
+):
+    """
+    Combines new inputs into an existing list, avoiding duplicates by name.
+    Modifies the existing_inputs list in place.
+    """
+    input_names = [task_input.name for task_input in existing_inputs]
+    if isinstance(new_inputs, AnyInput):
+        new_inputs_list = [new_inputs]
+    elif new_inputs is None:
+        new_inputs_list = []
+    else:
+        new_inputs_list = new_inputs
+
+    for task_input in new_inputs_list:
+        if task_input is None:
+            continue
+        if task_input.name not in input_names:
+            existing_inputs.append(task_input)
+            input_names.append(task_input.name)  # Update names list
+
+
+def get_combined_envs(task: "BaseTask") -> list[AnyEnv]:
+    """
+    Aggregates environment variables from the task and its upstreams.
+    """
+    envs = []
+    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]
+
+
+def get_combined_inputs(task: "BaseTask") -> list[AnyInput]:
+    """
+    Aggregates inputs from the task and its upstreams, avoiding duplicates.
+    """
+    inputs: list[AnyInput] = []
+    for upstream in task.upstreams:
+        combine_inputs(inputs, upstream.inputs)
+
+    # Access _inputs directly as task is BaseTask
+    task_inputs: list[AnyInput | None] | AnyInput | None = task._inputs
+    if task_inputs is not None:
+        combine_inputs(inputs, task_inputs)
+
+    # Filter out None values (although combine_inputs should handle this)
+    return [task_input for task_input in inputs if task_input is not None]

zrb-1.5.9/src/zrb/task/base/dependencies.py

@@ -0,0 +1,57 @@
+from typing import TypeVar
+
+from zrb.task.any_task import AnyTask
+
+T = TypeVar("T", bound="AnyTask")
+
+
+def get_dependency_list(task: AnyTask, attr_name: str) -> list[AnyTask]:
+    """
+    Safely retrieves a list of dependencies (upstreams, fallbacks, etc.)
+    from a task attribute, handling None or single-task values.
+    """
+    dependencies = getattr(task, attr_name, None)
+    if dependencies is None:
+        return []
+    elif isinstance(dependencies, list):
+        # Ensure all elements are AnyTask (or compatible) if needed,
+        # but for now, assume the list contains tasks.
+        return dependencies
+    # Assume it's a single AnyTask instance
+    return [dependencies]
+
+
+def append_dependency(
+    task: T, attr_name: str, dependencies_to_add: AnyTask | list[AnyTask]
+) -> None:
+    """
+    Appends one or more dependencies to the specified attribute list of a task,
+    ensuring the attribute becomes a list and avoiding duplicates.
+
+    Modifies the attribute on the task object directly.
+    """
+    # Retrieve the current list, handling None or single item cases
+    current_deps_list = getattr(task, attr_name, None)
+    if current_deps_list is None:
+        current_deps_list = []
+    elif not isinstance(current_deps_list, list):
+        current_deps_list = [current_deps_list]
+    else:
+        # Create a copy to avoid modifying the original list if no changes occur
+        current_deps_list = list(current_deps_list)
+
+    if isinstance(dependencies_to_add, list):
+        new_deps = dependencies_to_add
+    else:
+        new_deps = [dependencies_to_add]  # Ensure it's always a list
+
+    added = False
+    for dep in new_deps:
+        # Check for existence before appending
+        if dep not in current_deps_list:
+            current_deps_list.append(dep)
+            added = True
+
+    # Update the attribute on the task object only if changes were made
+    if added:
+        setattr(task, attr_name, current_deps_list)

zrb-1.5.9/src/zrb/task/base/execution.py

@@ -0,0 +1,274 @@
+import asyncio
+from typing import TYPE_CHECKING, Any
+
+from zrb.context.any_context import AnyContext
+from zrb.session.any_session import AnySession
+from zrb.util.attr import get_bool_attr
+from zrb.util.run import run_async
+from zrb.xcom.xcom import Xcom
+
+if TYPE_CHECKING:
+    from zrb.task.base_task import BaseTask
+
+
+async def execute_task_chain(task: "BaseTask", session: AnySession):
+    """
+    Executes the task and its downstream successors if conditions are met.
+    """
+    if session.is_terminated or not session.is_allowed_to_run(task):
+        return
+    result = await execute_task_action(task, session)
+    # Get next tasks
+    nexts = session.get_next_tasks(task)
+    if session.is_terminated or len(nexts) == 0:
+        return result
+    # Run next tasks asynchronously
+    next_coros = [run_async(next_task.exec_chain(session)) for next_task in nexts]
+    # Wait for the next tasks to complete. The result of the current task is returned.
+    await asyncio.gather(*next_coros)
+    return result
+
+
+async def execute_task_action(task: "BaseTask", session: AnySession):
+    """
+    Executes a single task's action, handling conditions and readiness checks.
+    """
+    ctx = task.get_ctx(session)
+    if not session.is_allowed_to_run(task):
+        # Task is not allowed to run, skip it for now.
+        # This will be triggered later if dependencies are met.
+        ctx.log_info("Not allowed to run")
+        return
+    if not check_execute_condition(task, session):
+        # Skip the task based on its execute_condition
+        ctx.log_info("Marked as skipped (condition false)")
+        session.get_task_status(task).mark_as_skipped()
+        return
+    # Wait for task to be ready (handles action execution, readiness checks)
+    return await run_async(execute_action_until_ready(task, session))
+
+
+def check_execute_condition(task: "BaseTask", session: AnySession) -> bool:
+    """
+    Evaluates the task's execute_condition attribute.
+    """
+    ctx = task.get_ctx(session)
+    execute_condition_attr = getattr(task, "_execute_condition", True)
+    return get_bool_attr(ctx, execute_condition_attr, True, auto_render=True)
+
+
+async def execute_action_until_ready(task: "BaseTask", session: AnySession):
+    """
+    Manages the execution of the task's action, coordinating with readiness checks.
+    """
+    ctx = task.get_ctx(session)
+    readiness_checks = task.readiness_checks
+    readiness_check_delay = getattr(task, "_readiness_check_delay", 0.5)
+    monitor_readiness = getattr(task, "_monitor_readiness", False)
+
+    if not readiness_checks:  # Simplified check for empty list
+        ctx.log_info("No readiness checks")
+        # Task has no readiness check, execute action directly
+        result = await run_async(execute_action_with_retry(task, session))
+        # Mark ready only if the action completed successfully (not failed/cancelled)
+        if session.get_task_status(task).is_completed:
+            ctx.log_info("Marked as ready")
+            session.get_task_status(task).mark_as_ready()
+        return result
+
+    # Start the task action and readiness checks concurrently
+    ctx.log_info("Starting action and readiness checks")
+    action_coro = asyncio.create_task(
+        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."
+            )
+
+    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
+
+
+async def execute_action_with_retry(task: "BaseTask", session: AnySession) -> Any:
+    """
+    Executes the task's core action (`_exec_action`) with retry logic,
+    handling success (triggering successors) and failure (triggering fallbacks).
+    """
+    ctx = task.get_ctx(session)
+    retries = getattr(task, "_retries", 2)
+    retry_period = getattr(task, "_retry_period", 0)
+    max_attempt = retries + 1
+    ctx.set_max_attempt(max_attempt)
+
+    for attempt in range(max_attempt):
+        ctx.set_attempt(attempt + 1)
+        if attempt > 0:
+            ctx.log_info(f"Retrying in {retry_period}s...")
+            await asyncio.sleep(retry_period)
+
+        try:
+            ctx.log_info("Marked as started")
+            session.get_task_status(task).mark_as_started()
+
+            # Execute the underlying action (which might be overridden in subclasses)
+            # We call the task's _exec_action method directly here.
+            result = await run_async(task._exec_action(ctx))
+
+            ctx.log_info("Marked as completed")
+            session.get_task_status(task).mark_as_completed()
+
+            # Store result in XCom
+            task_xcom: Xcom = ctx.xcom.get(task.name)
+            task_xcom.push(result)
+
+            # Skip fallbacks and execute successors on success
+            skip_fallbacks(task, session)
+            await run_async(execute_successors(task, session))
+            return result
+
+        except (asyncio.CancelledError, KeyboardInterrupt):
+            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
+            return  # Or re-raise? Depends on desired cancellation behavior
+
+        except BaseException as e:
+            ctx.log_error(f"Attempt {attempt + 1}/{max_attempt} failed: {e}")
+            session.get_task_status(
+                task
+            ).mark_as_failed()  # Mark failed for this attempt
+
+            if attempt < max_attempt - 1:
+                # More retries available
+                continue
+            else:
+                # Final attempt failed
+                ctx.log_error("Marked as permanently failed")
+                session.get_task_status(task).mark_as_permanently_failed()
+                # Skip successors and execute fallbacks on permanent failure
+                skip_successors(task, session)
+                await run_async(execute_fallbacks(task, session))
+                raise e  # Re-raise the exception after handling fallbacks
+
+
+async def run_default_action(task: "BaseTask", ctx: AnyContext) -> Any:
+    """
+    Executes the specific action defined by the '_action' attribute for BaseTask.
+    This is the default implementation called by BaseTask._exec_action.
+    Subclasses like LLMTask override _exec_action with their own logic.
+    """
+    action = getattr(task, "_action", None)
+    if action is None:
+        ctx.log_debug("No action defined for this task.")
+        return None
+    if isinstance(action, str):
+        # Render f-string action
+        rendered_action = ctx.render(action)
+        ctx.log_debug(f"Rendered action string: {rendered_action}")
+        # Assuming string actions are meant to be returned as is.
+        # If they need execution (e.g., shell commands), that logic would go here.
+        return rendered_action
+    elif callable(action):
+        # Execute callable action
+        ctx.log_debug(f"Executing callable action: {action.__name__}")
+        return await run_async(action(ctx))
+    else:
+        ctx.log_warning(f"Unsupported action type: {type(action)}")
+        return None
+
+
+async def execute_successors(task: "BaseTask", session: "AnySession"):
+    """Executes all successor tasks."""
+    ctx = task.get_ctx(session)
+    successors = task.successors
+    if successors:
+        ctx.log_info(f"Executing {len(successors)} successor(s)")
+        successor_coros = [
+            run_async(successor.exec_chain(session)) for successor in successors
+        ]
+        await asyncio.gather(*successor_coros)
+    else:
+        ctx.log_debug("No successors to execute.")
+
+
+def skip_successors(task: "BaseTask", session: AnySession):
+    """Marks all successor tasks as skipped."""
+    ctx = task.get_ctx(session)
+    successors = task.successors
+    if successors:
+        ctx.log_info(f"Skipping {len(successors)} successor(s)")
+        for successor in successors:
+            # Check if already skipped to avoid redundant logging/state changes
+            if not session.get_task_status(successor).is_skipped:
+                session.get_task_status(successor).mark_as_skipped()
+
+
+async def execute_fallbacks(task: "BaseTask", session: AnySession):
+    """Executes all fallback tasks."""
+    ctx = task.get_ctx(session)
+    fallbacks = task.fallbacks
+    if fallbacks:
+        ctx.log_info(f"Executing {len(fallbacks)} fallback(s)")
+        fallback_coros = [
+            run_async(fallback.exec_chain(session)) for fallback in fallbacks
+        ]
+        await asyncio.gather(*fallback_coros)
+    else:
+        ctx.log_debug("No fallbacks to execute.")
+
+
+def skip_fallbacks(task: "BaseTask", session: AnySession):
+    """Marks all fallback tasks as skipped."""
+    ctx = task.get_ctx(session)
+    fallbacks = task.fallbacks
+    if fallbacks:
+        ctx.log_info(f"Skipping {len(fallbacks)} fallback(s)")
+        for fallback in fallbacks:
+            # Check if already skipped
+            if not session.get_task_status(fallback).is_skipped:
+                session.get_task_status(fallback).mark_as_skipped()