zrb 1.0.0a21__py3-none-any.whl → 1.0.0b5__py3-none-any.whl

zrb/task/cmd_task.py

@@ -56,6 +56,7 @@ class CmdTask(BaseTask):
         monitor_readiness: bool = False,
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
         super().__init__(
             name=name,
@@ -76,6 +77,7 @@ class CmdTask(BaseTask):
             monitor_readiness=monitor_readiness,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._shell = shell
         self._render_shell = render_shell
@@ -233,17 +235,18 @@ class CmdTask(BaseTask):
 
     def _render_cmd_val(self, ctx: AnyContext, cmd_val: CmdVal) -> str:
         if isinstance(cmd_val, list):
+            cmd_val_list = [
+                self.__render_single_cmd_val(ctx, single_cmd_val)
+                for single_cmd_val in cmd_val
+            ]
             return "\n".join(
-                [
-                    self.__render_single_cmd_val(ctx, single_cmd_val)
-                    for single_cmd_val in cmd_val
-                ]
+                [cmd_val for cmd_val in cmd_val_list if cmd_val is not None]
             )
         return self.__render_single_cmd_val(ctx, cmd_val)
 
     def __render_single_cmd_val(
         self, ctx: AnyContext, single_cmd_val: SingleCmdVal
-    ) -> str:
+    ) -> str | None:
         if callable(single_cmd_val):
             return single_cmd_val(ctx)
         if isinstance(single_cmd_val, str):
@@ -252,6 +255,7 @@ class CmdTask(BaseTask):
             return single_cmd_val
         if isinstance(single_cmd_val, AnyCmdVal):
             return single_cmd_val.to_str(ctx)
+        return None
 
     def __get_multiline_repr(self, text: str) -> str:
         lines_repr: list[str] = []

zrb/task/http_check.py

@@ -28,6 +28,7 @@ class HttpCheck(BaseTask):
         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,
@@ -41,6 +42,7 @@ class HttpCheck(BaseTask):
             retries=0,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._url = url
         self._render_url = render_url

zrb/task/llm_task.py

@@ -3,7 +3,9 @@ import os
 from collections.abc import Callable
 from typing import Any
 
-from pydantic import BaseModel
+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
@@ -16,20 +18,10 @@ 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.llm.tool import callable_to_tool_schema
+from zrb.util.run import run_async
 
 ListOfDict = list[dict[str, Any]]
-
-
-class AdditionalTool(BaseModel):
-    fn: Callable
-    name: str | None
-    description: str | None
-
-
-def scratchpad(thought: str) -> str:
-    """Use this tool to note your thought and planning"""
-    return thought
+ToolOrCallable = Tool | Callable
 
 
 class LLMTask(BaseTask):
@@ -43,19 +35,28 @@ class LLMTask(BaseTask):
         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: (
-            dict[str, Callable] | Callable[[AnySharedContext], dict[str, Callable]]
-        ) = {},
-        history: ListOfDict | Callable[[AnySharedContext], ListOfDict] = [],
-        history_file: StrAttr | None = None,
+            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,
-        model_kwargs: (
-            dict[str, Any] | Callable[[AnySharedContext], dict[str, Any]]
-        ) = {},
         execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
         retries: int = 2,
         retry_period: float = 0,
@@ -65,8 +66,10 @@ class LLMTask(BaseTask):
         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,
@@ -87,92 +90,77 @@ class LLMTask(BaseTask):
             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._model_kwargs = model_kwargs
         self._system_prompt = system_prompt
         self._render_system_prompt = render_system_prompt
         self._message = message
         self._tools = tools
-        self._history = history
-        self._history_file = history_file
+        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._additional_tools: list[AdditionalTool] = []
+        self._max_call_iteration = max_call_iteration
 
-    def add_tool(
-        self, tool: Callable, name: str | None = None, description: str | None = None
-    ):
-        self._additional_tools.append(
-            AdditionalTool(fn=tool, name=name, description=description)
-        )
+    def add_tool(self, tool: ToolOrCallable):
+        self._additional_tools.append(tool)
 
     async def _exec_action(self, ctx: AnyContext) -> Any:
-        from litellm import acompletion, supports_function_calling
-
-        model = self._get_model(ctx)
-        try:
-            allow_function_call = supports_function_calling(model=model)
-        except Exception:
-            allow_function_call = False
-        model_kwargs = self._get_model_kwargs(ctx)
-        ctx.log_debug("MODEL KWARGS", model_kwargs)
-        system_prompt = self._get_system_prompt(ctx)
-        ctx.log_debug("SYSTEM PROMPT", system_prompt)
-        history = self._get_history(ctx)
-        ctx.log_debug("HISTORY PROMPT", history)
-        user_message = {"role": "user", "content": self._get_message(ctx)}
-        ctx.print(stylize_faint(f"{user_message}"))
-        messages = history + [user_message]
-        available_tools = self._get_tools(ctx)
-        available_tools["scratchpad"] = scratchpad
-        if allow_function_call:
-            tool_schema = [
-                callable_to_tool_schema(tool, name)
-                for name, tool in available_tools.items()
-            ]
-            for additional_tool in self._additional_tools:
-                fn = additional_tool.fn
-                tool_name = additional_tool.name or fn.__name__
-                tool_description = additional_tool.description
-                available_tools[tool_name] = additional_tool.fn
-                tool_schema.append(
-                    callable_to_tool_schema(
-                        fn, name=tool_name, description=tool_description
-                    )
-                )
-            model_kwargs["tools"] = tool_schema
-            ctx.log_debug("TOOL SCHEMA", tool_schema)
+        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)
-        while True:
-            response = await acompletion(
-                model=model,
-                messages=[{"role": "system", "content": system_prompt}] + messages,
-                **model_kwargs,
-            )
-            response_message = response.choices[0].message
-            ctx.print(stylize_faint(f"{response_message.to_dict()}"))
-            messages.append(response_message.to_dict())
-            tool_calls = response_message.tool_calls
-            if tool_calls:
-                # noqa Reference: https://docs.litellm.ai/docs/completion/function_call#full-code---parallel-function-calling-with-gpt-35-turbo-1106
-                for tool_call in tool_calls:
-                    function_name = tool_call.function.name
-                    function_to_call = available_tools[function_name]
-                    function_kwargs = json.loads(tool_call.function.arguments)
-                    function_response = function_to_call(**function_kwargs)
-                    tool_call_message = {
-                        "tool_call_id": tool_call.id,
-                        "role": "tool",
-                        "name": function_name,
-                        "content": function_response,
-                    }
-                    ctx.print(stylize_faint(f"{tool_call_message}"))
-                    messages.append(tool_call_message)
-                continue
-            if history_file != "":
-                write_file(history_file, json.dumps(messages, indent=2))
-            return response_message.content
+        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
+        ]
+        for tool in tools:
+            print("tool", tool)
+        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(
@@ -190,29 +178,24 @@ class LLMTask(BaseTask):
     def _get_message(self, ctx: AnyContext) -> str:
         return get_str_attr(ctx, self._message, "How are you?", auto_render=True)
 
-    def _get_model_kwargs(self, ctx: AnyContext) -> dict[str, Callable]:
-        if callable(self._model_kwargs):
-            return self._model_kwargs(ctx)
-        return {**self._model_kwargs}
-
-    def _get_tools(self, ctx: AnyContext) -> dict[str, Callable]:
-        if callable(self._tools):
-            return self._tools(ctx)
-        return self._tools
-
-    def _get_history(self, ctx: AnyContext) -> ListOfDict:
-        if callable(self._history):
-            return self._history(ctx)
+    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._history) == 0
+            len(self._conversation_history) == 0
             and history_file != ""
             and os.path.isfile(history_file)
         ):
             return json.loads(read_file(history_file))
-        return self._history
+        return self._conversation_history
 
     def _get_history_file(self, ctx: AnyContext) -> str:
         return get_str_attr(
-            ctx, self._history_file, "", auto_render=self._render_history_file
+            ctx,
+            self._conversation_history_file,
+            "",
+            auto_render=self._render_history_file,
         )

zrb/task/make_task.py

@@ -29,6 +29,7 @@ def make_task(
     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]:
@@ -53,6 +54,7 @@ def make_task(
             monitor_readiness=monitor_readiness,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         if group is not None:
             return group.add_task(task, alias=alias)

zrb/task/rsync_task.py

@@ -49,6 +49,7 @@ class RsyncTask(CmdTask):
         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,
     ):
         super().__init__(
             name=name,
@@ -80,6 +81,7 @@ class RsyncTask(CmdTask):
             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

zrb/task/scaffolder.py

@@ -11,6 +11,7 @@ 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
 
 TransformConfig = dict[str, str] | Callable[[AnyContext, str], str]
 
@@ -46,6 +47,7 @@ class Scaffolder(BaseTask):
         monitor_readiness: bool = False,
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
         super().__init__(
             name=name,
@@ -66,6 +68,7 @@ class Scaffolder(BaseTask):
             monitor_readiness=monitor_readiness,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._source_path = source_path
         self._render_source_path = render_source_path
@@ -83,13 +86,12 @@ class Scaffolder(BaseTask):
         return get_str_attr(ctx, self._destination_path, "", auto_render=True)
 
     def _get_content_transformers(self) -> list[AnyContentTransformer]:
-        if callable(self._content_transformers):
-            return [
-                ContentTransformer(match=".*", transform=self._content_transformers)
-            ]
-        if isinstance(self._content_transformers, dict):
+        if callable(self._content_transformers) or isinstance(
+            self._content_transformers, dict
+        ):
             return [
                 ContentTransformer(
+                    name="default-transform",
                     match=".*",
                     transform=self._content_transformers,
                     auto_render=self._render_content_transformers,
@@ -109,6 +111,7 @@ class Scaffolder(BaseTask):
             for transformer in transformers:
                 if transformer.match(ctx, file_path):
                     try:
+                        ctx.print(stylize_faint(f"{transformer.name}: {file_path}"))
                         transformer.transform_file(ctx, file_path)
                     except UnicodeDecodeError:
                         pass

zrb/task/scheduler.py

@@ -39,6 +39,7 @@ class Scheduler(BaseTrigger):
         monitor_readiness: bool = False,
         upstream: list[AnyTask] | AnyTask | None = None,
         fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
     ):
         super().__init__(
             name=name,
@@ -61,6 +62,7 @@ class Scheduler(BaseTrigger):
             monitor_readiness=monitor_readiness,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._cron_pattern = schedule
 

zrb/task/tcp_check.py

@@ -28,6 +28,7 @@ class TcpCheck(BaseTask):
         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,
@@ -41,6 +42,7 @@ class TcpCheck(BaseTask):
             retries=0,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._host = host
         self._render_host = render_host

zrb/task_status/task_status.py

@@ -62,9 +62,10 @@ class TaskStatus:
         self._history.append((TASK_PERMANENTLY_FAILED, datetime.datetime.now()))
 
     def mark_as_terminated(self):
-        self._is_terminated = True
-        if not self.is_completed and not self.is_permanently_failed:
-            self._history.append((TASK_TERMINATED, datetime.datetime.now()))
+        if not self._is_terminated:
+            self._is_terminated = True
+            if not (self.is_skipped or self.is_completed or self.is_permanently_failed):
+                self._history.append((TASK_TERMINATED, datetime.datetime.now()))
 
     @property
     def is_started(self) -> bool:

zrb/util/cmd/command.py

@@ -62,6 +62,7 @@ async def run_command(
             log_method(line)
         return "\n".join(lines)
 
+    cmd_process = None
     try:
         if cwd is None:
             cwd = os.getcwd()

zrb/util/file.py

@@ -1,4 +1,5 @@
 import os
+import re
 
 
 def read_file(file_path: str, replace_map: dict[str, str] = {}) -> str:
@@ -14,5 +15,10 @@ def write_file(file_path: str, content: str | list[str]):
         content = "\n".join([line for line in content if line is not None])
     dir_path = os.path.dirname(file_path)
     os.makedirs(dir_path, exist_ok=True)
+    content = re.sub(r"\n{3,}$", "\n\n", content)
+    # Remove trailing newlines, but keep one if it exists
+    content = content.rstrip("\n")
+    if content.endswith("\n"):
+        content += "\n"
     with open(file_path, "w") as f:
-        f.write(content.strip())
+        f.write(content)

zrb/util/llm/tool.py

@@ -1,34 +1,34 @@
 import inspect
 from collections.abc import Callable
-from typing import Any, get_type_hints
+from typing import Annotated, Any, Literal, get_type_hints
 
 
-def callable_to_tool_schema(
-    callable_obj: Callable, name: str | None = None, description: str | None = None
-) -> dict[str, Any]:
+def callable_to_tool_schema(callable_obj: Callable) -> dict[str, Any]:
     """
     Convert a callable into a tool schema dictionary by deriving the parameter schema.
 
     :param callable_obj: The callable object (e.g., a function).
-    :param name: The name to assign to the function in the schema.
-    :param description: A description of the function.
     :return: A dictionary representing the tool schema.
     """
     if not callable(callable_obj):
         raise ValueError("Provided object is not callable")
     # Derive name and description
-    name = name or callable_obj.__name__
-    description = description or (callable_obj.__doc__ or "").strip()
+    name = callable_obj.__name__
+    description = (callable_obj.__doc__ or "").strip()
     # Get function signature
     sig = inspect.signature(callable_obj)
     hints = get_type_hints(callable_obj)
     # Build parameter schema
     param_schema = {"type": "object", "properties": {}, "required": []}
     for param_name, param in sig.parameters.items():
-        param_type = hints.get(param_name, str)  # Default type is string
-        param_schema["properties"][param_name] = {
-            "type": _python_type_to_json_type(param_type)
-        }
+        # Get the type hint or default to str
+        param_type = hints.get(param_name, str)
+
+        # Handle annotated types (e.g., Annotated[str, "description"])
+        json_type, param_metadata = _process_type_annotation(param_type)
+        param_schema["properties"][param_name] = param_metadata
+
+        # Mark required parameters
         if param.default is inspect.Parameter.empty:
             param_schema["required"].append(param_name)
     return {
@@ -41,6 +41,30 @@ def callable_to_tool_schema(
     }
 
 
+def _process_type_annotation(py_type: Any) -> tuple[str, dict]:
+    """
+    Process type annotations and return the JSON Schema type and metadata.
+
+    :param py_type: The type annotation.
+    :return: A tuple of (JSON type, parameter metadata).
+    """
+    if hasattr(py_type, "__origin__") and py_type.__origin__ is Literal:
+        # Handle Literal (enum)
+        enum_values = list(py_type.__args__)
+        return "string", {"type": "string", "enum": enum_values}
+
+    if hasattr(py_type, "__origin__") and py_type.__origin__ is Annotated:
+        # Handle Annotated types
+        base_type = py_type.__args__[0]
+        description = py_type.__args__[1]
+        json_type = _python_type_to_json_type(base_type)
+        return json_type, {"type": json_type, "description": description}
+
+    # Fallback to basic type conversion
+    json_type = _python_type_to_json_type(py_type)
+    return json_type, {"type": json_type}
+
+
 def _python_type_to_json_type(py_type):
     """
     Map Python types to JSON Schema types.

{zrb-1.0.0a21.dist-info → zrb-1.0.0b5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.0.0a21
+Version: 1.0.0b5
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -21,11 +21,12 @@ Requires-Dist: chromadb (>=0.5.20,<0.6.0) ; extra == "rag"
 Requires-Dist: fastapi[standard] (>=0.115.6,<0.116.0)
 Requires-Dist: isort (>=5.13.2,<5.14.0)
 Requires-Dist: libcst (>=1.5.0,<2.0.0)
-Requires-Dist: litellm (>=1.52.12,<2.0.0)
 Requires-Dist: pdfplumber (>=0.11.4,<0.12.0) ; extra == "rag"
+Requires-Dist: pydantic-ai (>=0.0.19,<0.0.20)
 Requires-Dist: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: python-jose[cryptography] (>=3.3.0,<4.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)
+Requires-Dist: ulid-py (>=1.1.0,<2.0.0)
 Project-URL: Documentation, https://github.com/state-alchemists/zrb
 Project-URL: Repository, https://github.com/state-alchemists/zrb
 Description-Content-Type: text/markdown