zrb 0.28.1__py3-none-any.whl → 1.0.0__py3-none-any.whl

zrb/task/http_check.py

@@ -0,0 +1,75 @@
+import asyncio
+from collections.abc import Callable
+
+from zrb.attr.type import StrAttr
+from zrb.context.any_context import AnyContext
+from zrb.context.context import Context
+from zrb.env.any_env import AnyEnv
+from zrb.input.any_input import AnyInput
+from zrb.task.any_task import AnyTask
+from zrb.task.base_task import BaseTask
+from zrb.util.attr import get_str_attr
+
+
+class HttpCheck(BaseTask):
+    def __init__(
+        self,
+        name: str,
+        color: int | None = None,
+        icon: str | None = None,
+        description: str | None = None,
+        cli_only: bool = False,
+        input: list[AnyInput] | AnyInput | None = None,
+        env: list[AnyEnv] | AnyEnv | None = None,
+        url: StrAttr = "http://localhost",
+        render_url: bool = True,
+        http_method: StrAttr = "GET",
+        interval: int = 5,
+        execute_condition: bool | str | Callable[[Context], bool] = True,
+        upstream: list[AnyTask] | AnyTask | None = None,
+        fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
+    ):
+        super().__init__(
+            name=name,
+            color=color,
+            icon=icon,
+            description=description,
+            cli_only=cli_only,
+            input=input,
+            env=env,
+            execute_condition=execute_condition,
+            retries=0,
+            upstream=upstream,
+            fallback=fallback,
+            successor=successor,
+        )
+        self._url = url
+        self._render_url = render_url
+        self._http_method = http_method
+        self._interval = interval
+
+    def _get_url(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx, self._url, "http://localhost", auto_render=self._render_url
+        )
+
+    def _get_http_method(self, ctx: AnyContext) -> str:
+        return get_str_attr(ctx, self._http_method, "GET", auto_render=True).upper()
+
+    async def _exec_action(self, ctx: AnyContext) -> bool:
+        import requests
+
+        url = self._get_url(ctx)
+        http_method = self._get_http_method(ctx)
+        while True:
+            try:
+                response = requests.request(http_method, url)
+                if response.status_code == 200:
+                    return response
+                ctx.log_info(f"HTTP Status code: {response.status_code}")
+            except asyncio.TimeoutError as e:
+                ctx.log_info(f"Timeout error {e}")
+            except Exception as e:
+                ctx.log_info(f"Error: {e}")
+            await asyncio.sleep(self._interval)

zrb/task/llm_task.py

@@ -0,0 +1,199 @@
+import json
+import os
+from collections.abc import Callable
+from typing import Any
+
+from pydantic_ai import Agent, Tool
+from pydantic_ai.messages import ModelMessagesTypeAdapter
+from pydantic_ai.settings import ModelSettings
+
+from zrb.attr.type import StrAttr
+from zrb.config import LLM_MODEL, LLM_SYSTEM_PROMPT
+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.task.any_task import AnyTask
+from zrb.task.base_task import BaseTask
+from zrb.util.attr import get_str_attr
+from zrb.util.cli.style import stylize_faint
+from zrb.util.file import read_file, write_file
+from zrb.util.run import run_async
+
+ListOfDict = list[dict[str, Any]]
+ToolOrCallable = Tool | Callable
+
+
+class LLMTask(BaseTask):
+    def __init__(
+        self,
+        name: str,
+        color: int | None = None,
+        icon: str | None = None,
+        description: str | None = None,
+        cli_only: bool = False,
+        input: list[AnyInput | None] | AnyInput | None = None,
+        env: list[AnyEnv | None] | AnyEnv | None = None,
+        model: StrAttr | None = LLM_MODEL,
+        model_settings: (
+            ModelSettings | Callable[[AnySharedContext], ModelSettings] | None
+        ) = None,
+        render_model: bool = True,
+        agent: Agent | Callable[[AnySharedContext], Agent] | None = None,
+        system_prompt: StrAttr | None = LLM_SYSTEM_PROMPT,
+        render_system_prompt: bool = True,
+        message: StrAttr | None = None,
+        tools: (
+            list[ToolOrCallable] | Callable[[AnySharedContext], list[ToolOrCallable]]
+        ) = [],
+        conversation_history: (
+            ListOfDict | Callable[[AnySharedContext], ListOfDict]
+        ) = [],
+        conversation_history_reader: (
+            Callable[[AnySharedContext], ListOfDict] | None
+        ) = None,
+        conversation_history_writer: (
+            Callable[[AnySharedContext, ListOfDict], None] | None
+        ) = None,
+        conversation_history_file: StrAttr | None = None,
+        render_history_file: bool = True,
+        execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
+        retries: int = 2,
+        retry_period: float = 0,
+        readiness_check: list[AnyTask] | AnyTask | None = None,
+        readiness_check_delay: float = 0.5,
+        readiness_check_period: float = 5,
+        readiness_failure_threshold: int = 1,
+        readiness_timeout: int = 60,
+        monitor_readiness: bool = False,
+        max_call_iteration: int = 20,
+        upstream: list[AnyTask] | AnyTask | None = None,
+        fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
+    ):
+        super().__init__(
+            name=name,
+            color=color,
+            icon=icon,
+            description=description,
+            cli_only=cli_only,
+            input=input,
+            env=env,
+            execute_condition=execute_condition,
+            retries=retries,
+            retry_period=retry_period,
+            readiness_check=readiness_check,
+            readiness_check_delay=readiness_check_delay,
+            readiness_check_period=readiness_check_period,
+            readiness_failure_threshold=readiness_failure_threshold,
+            readiness_timeout=readiness_timeout,
+            monitor_readiness=monitor_readiness,
+            upstream=upstream,
+            fallback=fallback,
+            successor=successor,
+        )
+        self._model = model
+        self._model_settings = (model_settings,)
+        self._agent = agent
+        self._render_model = render_model
+        self._system_prompt = system_prompt
+        self._render_system_prompt = render_system_prompt
+        self._message = message
+        self._tools = tools
+        self._additional_tools: list[ToolOrCallable] = []
+        self._conversation_history = conversation_history
+        self._conversation_history_reader = conversation_history_reader
+        self._conversation_history_writer = conversation_history_writer
+        self._conversation_history_file = conversation_history_file
+        self._render_history_file = render_history_file
+        self._max_call_iteration = max_call_iteration
+
+    def add_tool(self, tool: ToolOrCallable):
+        self._additional_tools.append(tool)
+
+    async def _exec_action(self, ctx: AnyContext) -> Any:
+        history = await self._read_conversation_history(ctx)
+        user_prompt = self._get_message(ctx)
+        agent = self._get_agent(ctx)
+        result = await agent.run(
+            user_prompt=user_prompt,
+            message_history=ModelMessagesTypeAdapter.validate_python(history),
+        )
+        new_history = json.loads(result.all_messages_json())
+        for history in new_history:
+            ctx.print(stylize_faint(json.dumps(history)))
+        await self._write_conversation_history(ctx, new_history)
+        return result.data
+
+    async def _write_conversation_history(
+        self, ctx: AnyContext, conversations: list[Any]
+    ):
+        if self._conversation_history_writer is not None:
+            await run_async(self._conversation_history_writer(ctx, conversations))
+        history_file = self._get_history_file(ctx)
+        if history_file != "":
+            write_file(history_file, json.dumps(conversations, indent=2))
+
+    def _get_model_settings(self, ctx: AnyContext) -> ModelSettings | None:
+        if isinstance(self._model_settings, ModelSettings):
+            return self._model_settings
+        if callable(self._model_settings):
+            return self._model_settings(ctx)
+        return None
+
+    def _get_agent(self, ctx: AnyContext) -> Agent:
+        if isinstance(self._agent, Agent):
+            return self._agent
+        if callable(self._agent):
+            return self._agent(ctx)
+        tools_or_callables = list(
+            self._tools(ctx) if callable(self._tools) else self._tools
+        )
+        tools_or_callables.extend(self._additional_tools)
+        tools = [
+            tool if isinstance(tool, Tool) else Tool(tool, takes_ctx=False)
+            for tool in tools_or_callables
+        ]
+        return Agent(
+            self._get_model(ctx),
+            system_prompt=self._get_system_prompt(ctx),
+            tools=tools,
+        )
+
+    def _get_model(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx, self._model, "ollama_chat/llama3.1", auto_render=self._render_model
+        )
+
+    def _get_system_prompt(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx,
+            self._system_prompt,
+            "You are a helpful assistant",
+            auto_render=self._render_system_prompt,
+        )
+
+    def _get_message(self, ctx: AnyContext) -> str:
+        return get_str_attr(ctx, self._message, "How are you?", auto_render=True)
+
+    async def _read_conversation_history(self, ctx: AnyContext) -> ListOfDict:
+        if self._conversation_history_reader is not None:
+            return await run_async(self._conversation_history_reader(ctx))
+        if callable(self._conversation_history):
+            return self._conversation_history(ctx)
+        history_file = self._get_history_file(ctx)
+        if (
+            len(self._conversation_history) == 0
+            and history_file != ""
+            and os.path.isfile(history_file)
+        ):
+            return json.loads(read_file(history_file))
+        return self._conversation_history
+
+    def _get_history_file(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx,
+            self._conversation_history_file,
+            "",
+            auto_render=self._render_history_file,
+        )

zrb/task/make_task.py

@@ -0,0 +1,63 @@
+from collections.abc import Callable
+from typing import Any
+
+from zrb.context.any_context import AnyContext
+from zrb.context.any_shared_context import AnySharedContext
+from zrb.env.any_env import AnyEnv
+from zrb.group.any_group import AnyGroup
+from zrb.input.any_input import AnyInput
+from zrb.task.any_task import AnyTask
+from zrb.task.base_task import BaseTask
+
+
+def make_task(
+    name: str,
+    color: int | None = None,
+    icon: str | None = None,
+    description: str | None = None,
+    cli_only: bool = False,
+    input: list[AnyInput | None] | AnyInput | None = None,
+    env: list[AnyEnv | None] | AnyEnv | None = None,
+    execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
+    retries: int = 2,
+    retry_period: float = 0,
+    readiness_check: list[AnyTask] | AnyTask | None = None,
+    readiness_check_delay: float = 0.5,
+    readiness_check_period: float = 5,
+    readiness_failure_threshold: int = 1,
+    readiness_timeout: int = 60,
+    monitor_readiness: bool = False,
+    upstream: list[AnyTask] | AnyTask | None = None,
+    fallback: list[AnyTask] | AnyTask | None = None,
+    successor: list[AnyTask] | AnyTask | None = None,
+    group: AnyGroup | None = None,
+    alias: str | None = None,
+) -> Callable[[Callable[[AnyContext], Any]], AnyTask]:
+    def _make_task(fn: Callable[[AnyContext], Any]) -> BaseTask:
+        task = BaseTask(
+            name=name,
+            color=color,
+            icon=icon,
+            description=description,
+            cli_only=cli_only,
+            input=input,
+            env=env,
+            action=fn,
+            execute_condition=execute_condition,
+            retries=retries,
+            retry_period=retry_period,
+            readiness_check=readiness_check,
+            readiness_check_delay=readiness_check_delay,
+            readiness_check_period=readiness_check_period,
+            readiness_failure_threshold=readiness_failure_threshold,
+            readiness_timeout=readiness_timeout,
+            monitor_readiness=monitor_readiness,
+            upstream=upstream,
+            fallback=fallback,
+            successor=successor,
+        )
+        if group is not None:
+            return group.add_task(task, alias=alias)
+        return task
+
+    return _make_task

zrb/task/rsync_task.py

@@ -1,144 +1,155 @@
-import os
-import pathlib
-from collections.abc import Callable, Iterable
-from typing import Any, Optional, Union
+from collections.abc import Callable
 
-from zrb.helper.accessories.color import colored
-from zrb.helper.log import logger
-from zrb.helper.typecheck import typechecked
-from zrb.helper.typing import JinjaTemplate
+from zrb.attr.type import IntAttr, StrAttr
+from zrb.context.any_context import AnyContext
+from zrb.env.any_env import AnyEnv
+from zrb.input.any_input import AnyInput
 from zrb.task.any_task import AnyTask
-from zrb.task.any_task_event_handler import (
-    OnFailed,
-    OnReady,
-    OnRetry,
-    OnSkipped,
-    OnStarted,
-    OnTriggered,
-    OnWaiting,
-)
 from zrb.task.cmd_task import CmdTask
-from zrb.task_env.env import Env, PrivateEnv
-from zrb.task_env.env_file import EnvFile
-from zrb.task_group.group import Group
-from zrb.task_input.any_input import AnyInput
+from zrb.util.attr import get_str_attr
 
-logger.debug(colored("Loading zrb.task.rsync_task", attrs=["dark"]))
 
-_CURRENT_DIR = os.path.dirname(__file__)
-_SHELL_SCRIPT_DIR = os.path.join(_CURRENT_DIR, "..", "shell-scripts")
-
-ensure_rsync_is_installed = CmdTask(
-    name="ensure-ssh-is-installed",
-    cmd_path=[
-        os.path.join(_SHELL_SCRIPT_DIR, "_common-util.sh"),
-        os.path.join(_SHELL_SCRIPT_DIR, "ensure-ssh-is-installed.sh"),
-        os.path.join(_SHELL_SCRIPT_DIR, "ensure-rsync-is-installed.sh"),
-    ],
-    preexec_fn=None,
-    should_print_cmd_result=False,
-    should_show_cmd=False,
-    should_show_working_directory=False,
-)
-
-
-@typechecked
 class RsyncTask(CmdTask):
     def __init__(
         self,
         name: str,
-        group: Optional[Group] = None,
-        inputs: Iterable[AnyInput] = [],
-        envs: Iterable[Env] = [],
-        env_files: Iterable[EnvFile] = [],
-        icon: Optional[str] = None,
-        color: Optional[str] = None,
-        description: str = "",
-        executable: Optional[str] = None,
-        remote_host: JinjaTemplate = "localhost",
-        remote_port: Union[JinjaTemplate, int] = 22,
-        remote_user: JinjaTemplate = "root",
-        remote_password: JinjaTemplate = "",
-        remote_ssh_key: JinjaTemplate = "",
-        src_path: JinjaTemplate = ".",
-        src_is_remote: bool = False,
-        dst_path: JinjaTemplate = ".",
-        dst_is_remote: bool = True,
-        cwd: Optional[Union[JinjaTemplate, pathlib.Path]] = None,
-        should_render_cwd: bool = True,
-        upstreams: Iterable[AnyTask] = [],
-        fallbacks: Iterable[AnyTask] = [],
-        on_triggered: Optional[OnTriggered] = None,
-        on_waiting: Optional[OnWaiting] = None,
-        on_skipped: Optional[OnSkipped] = None,
-        on_started: Optional[OnStarted] = None,
-        on_ready: Optional[OnReady] = None,
-        on_retry: Optional[OnRetry] = None,
-        on_failed: Optional[OnFailed] = None,
-        checkers: Iterable[AnyTask] = [],
-        checking_interval: Union[float, int] = 0.05,
-        retry: int = 2,
-        retry_interval: Union[float, int] = 1,
+        color: int | None = None,
+        icon: str | None = None,
+        description: str | None = None,
+        cli_only: bool = False,
+        input: list[AnyInput] | AnyInput | None = None,
+        env: list[AnyEnv] | AnyEnv | None = None,
+        shell: StrAttr | None = None,
+        auto_render_shell: bool = True,
+        remote_host: StrAttr | None = None,
+        render_remote_host: bool = True,
+        remote_port: IntAttr | None = None,
+        render_remote_port: bool = True,
+        remote_user: StrAttr | None = None,
+        render_remote_user: bool = True,
+        remote_password: StrAttr | None = None,
+        render_remote_password: bool = True,
+        remote_ssh_key: StrAttr | None = None,
+        render_remote_ssh_key: bool = True,
+        remote_source_path: StrAttr | None = None,
+        render_remote_source_path: bool = True,
+        remote_destination_path: StrAttr | None = None,
+        render_remote_destination_path: bool = True,
+        local_source_path: StrAttr | None = None,
+        render_local_source_path: bool = True,
+        local_destination_path: StrAttr | None = None,
+        render_local_destination_path: bool = True,
+        cwd: str | None = None,
+        render_cwd: bool = True,
+        plain_print: bool = False,
         max_output_line: int = 1000,
         max_error_line: int = 1000,
-        preexec_fn: Optional[Callable[[], Any]] = os.setsid,
-        should_execute: Union[bool, str, Callable[..., bool]] = True,
-        return_upstream_result: bool = False,
-        should_print_cmd_result: bool = True,
-        should_show_cmd: bool = True,
-        should_show_working_directory: bool = True,
+        execute_condition: bool | str | Callable[[AnyContext], bool] = True,
+        retries: int = 2,
+        retry_period: float = 0,
+        readiness_check: list[AnyTask] | AnyTask | None = None,
+        upstream: list[AnyTask] | AnyTask | None = None,
+        fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
-        CmdTask.__init__(
-            self,
+        super().__init__(
             name=name,
-            group=group,
-            inputs=inputs,
-            envs=envs + [PrivateEnv(name="_ZRB_SSH_PASSWORD", default=remote_password)],
-            env_files=env_files,
-            icon=icon,
             color=color,
+            icon=icon,
             description=description,
-            executable=executable,
+            cli_only=cli_only,
+            input=input,
+            env=env,
+            shell=shell,
+            render_shell=auto_render_shell,
+            remote_host=remote_host,
+            render_remote_host=render_remote_host,
+            remote_port=remote_port,
+            auto_render_remote_port=render_remote_port,
+            remote_user=remote_user,
+            render_remote_user=render_remote_user,
+            remote_password=remote_password,
+            render_remote_password=render_remote_password,
+            remote_ssh_key=remote_ssh_key,
+            render_remote_ssh_key=render_remote_ssh_key,
             cwd=cwd,
-            should_render_cwd=should_render_cwd,
-            upstreams=[ensure_rsync_is_installed] + upstreams,
-            fallbacks=fallbacks,
-            on_triggered=on_triggered,
-            on_waiting=on_waiting,
-            on_skipped=on_skipped,
-            on_started=on_started,
-            on_ready=on_ready,
-            on_retry=on_retry,
-            on_failed=on_failed,
-            checkers=checkers,
-            checking_interval=checking_interval,
-            retry=retry,
-            retry_interval=retry_interval,
+            render_cwd=render_cwd,
+            plain_print=plain_print,
             max_output_line=max_output_line,
             max_error_line=max_error_line,
-            preexec_fn=preexec_fn,
-            should_execute=should_execute,
-            return_upstream_result=return_upstream_result,
-            should_print_cmd_result=should_print_cmd_result,
-            should_show_cmd=should_show_cmd,
-            should_show_working_directory=should_show_working_directory,
+            execute_condition=execute_condition,
+            retries=retries,
+            retry_period=retry_period,
+            readiness_check=readiness_check,
+            upstream=upstream,
+            fallback=fallback,
+            successor=successor,
+        )
+        self._remote_source_path = remote_source_path
+        self._render_remote_source_path = render_remote_source_path
+        self._remote_destination_path = remote_destination_path
+        self._render_remote_destination_path = render_remote_destination_path
+        self._local_source_path = local_source_path
+        self._render_local_source_path = render_local_source_path
+        self._local_destination_path = local_destination_path
+        self._render_local_destination_path = render_local_destination_path
+
+    def _get_source_path(self, ctx: AnyContext) -> str:
+        local_source_path = self._get_local_source_path(ctx)
+        if local_source_path != "":
+            return local_source_path
+        remote_source_path = self._get_remote_source_path(ctx)
+        host = self._get_remote_host(ctx)
+        user = self._get_remote_user(ctx)
+        return f"{user}@{host}:{remote_source_path}"
+
+    def _get_destination_path(self, ctx: AnyContext) -> str:
+        local_destination_path = self._get_local_destination_path(ctx)
+        if local_destination_path != "":
+            return local_destination_path
+        remote_destination_path = self._get_remote_destination_path(ctx)
+        host = self._get_remote_host(ctx)
+        user = self._get_remote_user(ctx)
+        return f"{user}@{host}:{remote_destination_path}"
+
+    def _get_remote_source_path(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx,
+            self._remote_source_path,
+            "",
+            auto_render=self._render_remote_source_path,
+        )
+
+    def _get_remote_destination_path(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx,
+            self._remote_destination_path,
+            "",
+            auto_render=self._render_remote_destination_path,
+        )
+
+    def _get_local_source_path(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx,
+            self._local_source_path,
+            "",
+            auto_render=self._render_local_source_path,
+        )
+
+    def _get_local_destination_path(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx,
+            self._local_destination_path,
+            "",
+            auto_render=self._render_local_destination_path,
         )
-        self._remote_host = remote_host
-        self._remote_port = remote_port
-        self._remote_user = remote_user
-        self._remote_password = remote_password
-        self._remote_ssh_key = remote_ssh_key
-        self._src_path = src_path
-        self._src_is_remote = src_is_remote
-        self._dst_path = dst_path
-        self._dst_is_remote = dst_is_remote
 
-    def get_cmd_script(self, *args: Any, **kwargs: Any) -> str:
-        port = self.render_str(self._remote_port)
-        password = self.render_str(self._remote_password)
-        key = self.render_str(self._remote_ssh_key)
-        src = self._get_path(self._src_path, self._src_is_remote)
-        dst = self._get_path(self._dst_path, self._dst_is_remote)
+    def _get_cmd_script(self, ctx: AnyContext) -> str:
+        port = self._get_remote_port(ctx)
+        password = self._get_remote_password(ctx)
+        key = self._get_remote_ssh_key(ctx)
+        src = self._get_source_path(ctx)
+        dst = self._get_destination_path(ctx)
         if key != "" and password != "":
             return f'sshpass -p "$_ZRB_SSH_PASSWORD" rsync --mkpath -avz -e "ssh -i {key} -p {port}" {src} {dst}'  # noqa
         if key != "":
@@ -146,11 +157,3 @@ class RsyncTask(CmdTask):
         if password != "":
             return f'sshpass -p "$_ZRB_SSH_PASSWORD" rsync --mkpath -avz -e "ssh -p {port}" {src} {dst}'  # noqa
         return f'rsync --mkpath -avz -e "ssh -p {port}" {src} {dst}'
-
-    def _get_path(self, resource_path: str, is_remote: bool) -> str:
-        rendered_path = self.render_str(resource_path)
-        if is_remote:
-            host = self.render_str(self._remote_host)
-            user = self.render_str(self._remote_user)
-            return f"{user}@{host}:{rendered_path}"
-        return rendered_path