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

zrb/task/llm_task.py

@@ -17,6 +17,7 @@ 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]]
 
@@ -24,7 +25,6 @@ ListOfDict = list[dict[str, Any]]
 class AdditionalTool(BaseModel):
     fn: Callable
     name: str | None
-    description: str | None
 
 
 def scratchpad(thought: str) -> str:
@@ -32,6 +32,11 @@ def scratchpad(thought: str) -> str:
     return thought
 
 
+def end_conversation(final_answer: str) -> str:
+    """End conversation with a final answer containing all necessary information"""
+    return final_answer
+
+
 class LLMTask(BaseTask):
     def __init__(
         self,
@@ -47,9 +52,7 @@ class LLMTask(BaseTask):
         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]]
-        ) = {},
+        tools: list[Callable] | Callable[[AnySharedContext], list[Callable]] = [],
         history: ListOfDict | Callable[[AnySharedContext], ListOfDict] = [],
         history_file: StrAttr | None = None,
         render_history_file: bool = True,
@@ -67,6 +70,7 @@ class LLMTask(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,
@@ -87,6 +91,7 @@ class LLMTask(BaseTask):
             monitor_readiness=monitor_readiness,
             upstream=upstream,
             fallback=fallback,
+            successor=successor,
         )
         self._model = model
         self._render_model = render_model
@@ -98,81 +103,156 @@ class LLMTask(BaseTask):
         self._history = history
         self._history_file = history_file
         self._render_history_file = render_history_file
-        self._additional_tools: list[AdditionalTool] = []
 
-    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: Callable):
+        self._tools.append(tool)
 
     async def _exec_action(self, ctx: AnyContext) -> Any:
-        from litellm import acompletion, supports_function_calling
+        import litellm
+        from litellm.utils import supports_function_calling
 
+        user_message = {"role": "user", "content": self._get_message(ctx)}
+        ctx.print(stylize_faint(f"{user_message}"))
         model = self._get_model(ctx)
         try:
-            allow_function_call = supports_function_calling(model=model)
+            is_function_call_supported = supports_function_calling(model=model)
         except Exception:
-            allow_function_call = False
-        model_kwargs = self._get_model_kwargs(ctx)
+            is_function_call_supported = False
+            litellm.add_function_to_prompt = True
+        if not is_function_call_supported:
+            ctx.log_warning(f"Model {model} doesn't support function call")
+        available_tools = self._get_available_tools(
+            ctx, include_end_conversation=not is_function_call_supported
+        )
+        model_kwargs = self._get_model_kwargs(ctx, available_tools)
         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)
+        conversations = history + [user_message]
         history_file = self._get_history_file(ctx)
         while True:
-            response = await acompletion(
-                model=model,
-                messages=[{"role": "system", "content": system_prompt}] + messages,
-                **model_kwargs,
+            llm_response = await self._get_llm_response(
+                model, system_prompt, conversations, model_kwargs
+            )
+            llm_response_dict = llm_response.to_dict()
+            ctx.print(stylize_faint(f"{llm_response_dict}"))
+            conversations.append(llm_response_dict)
+            ctx.log_debug("RESPONSE MESSAGE", llm_response)
+            if is_function_call_supported:
+                if not llm_response.tool_calls:
+                    # No tool call, end conversation
+                    self._save_conversation(history_file, conversations)
+                    return llm_response.content
+                await self._handle_tool_calls(
+                    ctx, available_tools, conversations, llm_response
+                )
+            if not is_function_call_supported:
+                try:
+                    json_payload = json.loads(llm_response.content)
+                    function_name = _get_fallback_function_name(json_payload)
+                    function_kwargs = _get_fallback_function_kwargs(json_payload)
+                    tool_execution_message = (
+                        await self._create_fallback_tool_exec_message(
+                            available_tools, function_name, function_kwargs
+                        )
+                    )
+                    ctx.print(stylize_faint(f"{tool_execution_message}"))
+                    conversations.append(tool_execution_message)
+                    if function_name == "end_conversation":
+                        self._save_conversation(history_file, conversations)
+                        return function_kwargs.get("final_answer", "")
+                except Exception as e:
+                    ctx.log_error(e)
+                    tool_execution_message = self._create_exec_scratchpad_message(
+                        f"{e}"
+                    )
+                    conversations.append(tool_execution_message)
+
+    async def _handle_tool_calls(
+        self,
+        ctx: AnyContext,
+        available_tools: dict[str, Callable],
+        conversations: list[dict[str, Any]],
+        llm_response: Any,
+    ):
+        # noqa Reference: https://docs.litellm.ai/docs/completion/function_call#full-code---parallel-function-calling-with-gpt-35-turbo-1106
+        for tool_call in llm_response.tool_calls:
+            tool_execution_message = await self._create_tool_exec_message(
+                available_tools, tool_call
             )
-            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
+            ctx.print(stylize_faint(f"{tool_execution_message}"))
+            conversations.append(tool_execution_message)
+
+    def _save_conversation(self, history_file: str, conversations: list[Any]):
+        if history_file != "":
+            write_file(history_file, json.dumps(conversations, indent=2))
+
+    async def _get_llm_response(
+        self,
+        model: str,
+        system_prompt: str,
+        conversations: list[Any],
+        model_kwargs: dict[str, Any],
+    ) -> Any:
+        from litellm import acompletion
+
+        llm_response = await acompletion(
+            model=model,
+            messages=[{"role": "system", "content": system_prompt}] + conversations,
+            **model_kwargs,
+        )
+        return llm_response.choices[0].message
+
+    async def _create_tool_exec_message(
+        self, available_tools: dict[str, Callable], tool_call: Any
+    ) -> dict[str, Any]:
+        function_name = tool_call.function.name
+        function_kwargs = json.loads(tool_call.function.arguments)
+        return {
+            "tool_call_id": tool_call.id,
+            "role": "tool",
+            "name": function_name,
+            "content": await self._get_exec_tool_result(
+                available_tools, function_name, function_kwargs
+            ),
+        }
+
+    async def _create_fallback_tool_exec_message(
+        self,
+        available_tools: dict[str, Callable],
+        function_name: str,
+        function_kwargs: dict[str, Any],
+    ) -> dict[str, Any]:
+        result = await self._get_exec_tool_result(
+            available_tools, function_name, function_kwargs
+        )
+        return self._create_exec_scratchpad_message(
+            f"Result of {function_name} call: {result}"
+        )
+
+    def _create_exec_scratchpad_message(self, message: str) -> dict[str, Any]:
+        return {
+            "role": "assistant",
+            "content": json.dumps(
+                {"name": "scratchpad", "arguments": {"thought": message}}
+            ),
+        }
+
+    async def _get_exec_tool_result(
+        self,
+        available_tools: dict[str, Callable],
+        function_name: str,
+        function_kwargs: dict[str, Any],
+    ) -> str:
+        if function_name not in available_tools:
+            return f"[ERROR] Invalid tool: {function_name}"
+        function_to_call = available_tools[function_name]
+        try:
+            return await run_async(function_to_call(**function_kwargs))
+        except Exception as e:
+            return f"[ERROR] {e}"
 
     def _get_model(self, ctx: AnyContext) -> str:
         return get_str_attr(
@@ -190,15 +270,29 @@ 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]:
+    def _get_model_kwargs(
+        self, ctx: AnyContext, available_tools: dict[str, Callable]
+    ) -> dict[str, Any]:
+        model_kwargs = {}
         if callable(self._model_kwargs):
-            return self._model_kwargs(ctx)
-        return {**self._model_kwargs}
+            model_kwargs = self._model_kwargs(ctx)
+        else:
+            model_kwargs = self._model_kwargs
+        model_kwargs["tools"] = [
+            callable_to_tool_schema(tool) for tool in available_tools.values()
+        ]
+        return model_kwargs
 
-    def _get_tools(self, ctx: AnyContext) -> dict[str, Callable]:
-        if callable(self._tools):
-            return self._tools(ctx)
-        return self._tools
+    def _get_available_tools(
+        self, ctx: AnyContext, include_end_conversation: bool
+    ) -> dict[str, Callable]:
+        tools = {"scratchpad": scratchpad}
+        if include_end_conversation:
+            tools["end_conversation"] = end_conversation
+        tool_list = self._tools(ctx) if callable(self._tools) else self._tools
+        for tool in tool_list:
+            tools[tool.__name__] = tool
+        return tools
 
     def _get_history(self, ctx: AnyContext) -> ListOfDict:
         if callable(self._history):
@@ -216,3 +310,22 @@ class LLMTask(BaseTask):
         return get_str_attr(
             ctx, self._history_file, "", auto_render=self._render_history_file
         )
+
+
+def _get_fallback_function_name(json_payload: dict[str, Any]) -> str:
+    for key in ("name",):
+        if key in json_payload:
+            return json_payload[key]
+    raise ValueError("Function name not provided")
+
+
+def _get_fallback_function_kwargs(json_payload: dict[str, Any]) -> str:
+    for key in (
+        "arguments",
+        "args",
+        "parameters",
+        "params",
+    ):
+        if key in json_payload:
+            return json_payload[key]
+    raise ValueError("Function arguments not provided")

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

@@ -3,22 +3,18 @@ from collections.abc import Callable
 from typing import Any, 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)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.0.0a21
+Version: 1.0.0b3
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -26,6 +26,7 @@ Requires-Dist: pdfplumber (>=0.11.4,<0.12.0) ; extra == "rag"
 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