zrb 1.0.0a3__py3-none-any.whl → 1.0.0a4__py3-none-any.whl

zrb/task/llm_task.py

@@ -1,7 +1,10 @@
 import json
+import os
 from collections.abc import Callable
 from typing import Any
 
+from pydantic import BaseModel
+
 from zrb.attr.type import StrAttr
 from zrb.config import LLM_MODEL, LLM_SYSTEM_PROMPT
 from zrb.context.any_context import AnyContext
@@ -11,9 +14,16 @@ 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.llm.tool import callable_to_tool_schema
 
-DictList = list[dict[str, Any]]
+ListOfDict = list[dict[str, Any]]
+
+
+class AdditionalTool(BaseModel):
+    fn: Callable
+    name: str | None
+    description: str | None
 
 
 def scratchpad(thought: str) -> str:
@@ -33,13 +43,19 @@ class LLMTask(BaseTask):
         input: list[AnyInput | None] | AnyInput | None = None,
         env: list[AnyEnv | None] | AnyEnv | None = None,
         model: StrAttr | None = LLM_MODEL,
+        render_model: bool = True,
         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]]
         ) = {},
-        tool_schema: DictList | Callable[[AnySharedContext], DictList] = [],
-        history: DictList | Callable[[AnySharedContext], DictList] = [],
+        history: ListOfDict | Callable[[AnySharedContext], ListOfDict] = [],
+        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,
@@ -73,47 +89,64 @@ class LLMTask(BaseTask):
             fallback=fallback,
         )
         self._model = model
+        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._tool_schema = tool_schema
         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)
+        )
 
     async def _exec_action(self, ctx: AnyContext) -> Any:
-        from litellm import completion
+        from litellm import acompletion
 
+        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 = self._get_message(ctx)
-        ctx.log_debug("USER MESSAGE", user_message)
-        messages = (
-            [{"role": "system", "content": system_prompt}]
-            + history
-            + [{"role": "user", "content": user_message}]
-        )
+        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
-        tool_schema = self._get_tool_schema(ctx)
-        for tool_name, tool in available_tools.items():
-            matched_tool_schema = [
-                schema
-                for schema in tool_schema
-                if "function" in schema
-                and "name" in schema["function"]
-                and schema["function"]["name"] == tool_name
-            ]
-            if len(matched_tool_schema) == 0:
-                tool_schema.append(callable_to_tool_schema(tool))
+        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
+                )
+            )
         ctx.log_debug("TOOL SCHEMA", tool_schema)
+        history_file = self._get_history_file(ctx)
         while True:
-            response = completion(
-                model=self._get_model(ctx), messages=messages, tools=tool_schema
+            response = await acompletion(
+                model=self._get_model(ctx),
+                messages=[{"role": "system", "content": system_prompt}] + messages,
+                tools=tool_schema,
+                **model_kwargs,
             )
             response_message = response.choices[0].message
-            ctx.print(response_message)
-            messages.append(response_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
@@ -128,33 +161,55 @@ class LLMTask(BaseTask):
                         "name": function_name,
                         "content": function_response,
                     }
-                    ctx.print(tool_call_message)
+                    ctx.print(stylize_faint(f"{tool_call_message}"))
                     messages.append(tool_call_message)
                 continue
+            if history_file != "":
+                os.makedirs(os.path.dirname(history_file), exist_ok=True)
+                with open(history_file, "w") as f:
+                    f.write(json.dumps(messages, indent=2))
             return response_message.content
 
     def _get_model(self, ctx: AnyContext) -> str:
-        return get_str_attr(ctx, self._model, "ollama_chat/llama3.1", auto_render=True)
+        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=True
+            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)
 
+    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_tool_schema(self, ctx: AnyContext) -> DictList:
-        if callable(self._tool_schema):
-            return self._tool_schema(ctx)
-        return self._tool_schema
-
-    def _get_history(self, ctx: AnyContext) -> DictList:
+    def _get_history(self, ctx: AnyContext) -> ListOfDict:
         if callable(self._history):
             return self._history(ctx)
+        history_file = self._get_history_file(ctx)
+        if (
+            len(self._history) == 0
+            and history_file != ""
+            and os.path.isfile(history_file)
+        ):
+            with open(history_file, "r") as f:
+                return json.loads(f.read())
         return self._history
+
+    def _get_history_file(self, ctx: AnyContext) -> str:
+        return get_str_attr(
+            ctx, self._history_file, "", auto_render=self._render_history_file
+        )

zrb/task/rsync_task.py

@@ -32,13 +32,13 @@ class RsyncTask(CmdTask):
         remote_ssh_key: StrAttr | None = None,
         auto_render_remote_ssh_key: bool = True,
         remote_source_path: StrAttr | None = None,
-        auto_render_remote_source_path: bool = True,
+        render_remote_source_path: bool = True,
         remote_destination_path: StrAttr | None = None,
-        auto_render_remote_destination_path: bool = True,
+        render_remote_destination_path: bool = True,
         local_source_path: StrAttr | None = None,
-        auto_render_local_source_path: bool = True,
+        render_local_source_path: bool = True,
         local_destination_path: StrAttr | None = None,
-        auto_render_local_destination_path: bool = True,
+        render_local_destination_path: bool = True,
         cwd: str | None = None,
         auto_render_cwd: bool = True,
         max_output_line: int = 1000,
@@ -59,19 +59,19 @@ class RsyncTask(CmdTask):
             input=input,
             env=env,
             shell=shell,
-            auto_render_shell=auto_render_shell,
+            render_shell=auto_render_shell,
             remote_host=remote_host,
-            auto_render_remote_host=auto_render_remote_host,
+            render_remote_host=auto_render_remote_host,
             remote_port=remote_port,
             auto_render_remote_port=auto_render_remote_port,
             remote_user=remote_user,
-            auto_render_remote_user=auto_render_remote_user,
+            render_remote_user=auto_render_remote_user,
             remote_password=remote_password,
-            auto_render_remote_password=auto_render_remote_password,
+            render_remote_password=auto_render_remote_password,
             remote_ssh_key=remote_ssh_key,
-            auto_render_remote_ssh_key=auto_render_remote_ssh_key,
+            render_remote_ssh_key=auto_render_remote_ssh_key,
             cwd=cwd,
-            auto_render_cwd=auto_render_cwd,
+            render_cwd=auto_render_cwd,
             max_output_line=max_output_line,
             max_error_line=max_error_line,
             execute_condition=execute_condition,
@@ -82,13 +82,13 @@ class RsyncTask(CmdTask):
             fallback=fallback,
         )
         self._remote_source_path = remote_source_path
-        self._auto_render_remote_source_path = auto_render_remote_source_path
+        self._render_remote_source_path = render_remote_source_path
         self._remote_destination_path = remote_destination_path
-        self._auto_render_remote_destination_path = auto_render_remote_destination_path
+        self._render_remote_destination_path = render_remote_destination_path
         self._local_source_path = local_source_path
-        self._auto_render_local_source_path = auto_render_local_source_path
+        self._render_local_source_path = render_local_source_path
         self._local_destination_path = local_destination_path
-        self._auto_render_local_destination_path = auto_render_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)
@@ -113,7 +113,7 @@ class RsyncTask(CmdTask):
             ctx,
             self._remote_source_path,
             "",
-            auto_render=self._auto_render_remote_source_path,
+            auto_render=self._render_remote_source_path,
         )
 
     def _get_remote_destination_path(self, ctx: AnyContext) -> str:
@@ -121,7 +121,7 @@ class RsyncTask(CmdTask):
             ctx,
             self._remote_destination_path,
             "",
-            auto_render=self._auto_render_remote_destination_path,
+            auto_render=self._render_remote_destination_path,
         )
 
     def _get_local_source_path(self, ctx: AnyContext) -> str:
@@ -129,7 +129,7 @@ class RsyncTask(CmdTask):
             ctx,
             self._local_source_path,
             "",
-            auto_render=self._auto_render_local_source_path,
+            auto_render=self._render_local_source_path,
         )
 
     def _get_local_destination_path(self, ctx: AnyContext) -> str:
@@ -137,7 +137,7 @@ class RsyncTask(CmdTask):
             ctx,
             self._local_destination_path,
             "",
-            auto_render=self._auto_render_local_destination_path,
+            auto_render=self._render_local_destination_path,
         )
 
     def _get_cmd_script(self, ctx: AnyContext) -> str:

zrb/task/scaffolder.py

@@ -26,9 +26,9 @@ class Scaffolder(BaseTask):
         input: list[AnyInput] | AnyInput | None = None,
         env: list[AnyEnv] | AnyEnv | None = None,
         source_path: StrAttr | None = None,
-        auto_render_source_path: bool = True,
+        render_source_path: bool = True,
         destination_path: StrAttr | None = None,
-        auto_render_destination_path: bool = True,
+        render_destination_path: bool = True,
         transform_path: TransformConfig = {},
         transform_content: (
             list[AnyContentTransformer] | AnyContentTransformer | TransformConfig
@@ -66,9 +66,9 @@ class Scaffolder(BaseTask):
             fallback=fallback,
         )
         self._source_path = source_path
-        self._auto_render_source_path = auto_render_source_path
+        self._render_source_path = render_source_path
         self._destination_path = destination_path
-        self._auto_render_destination_path = auto_render_destination_path
+        self._render_destination_path = render_destination_path
         self._content_transformers = transform_content
         self._path_transformer = transform_path
 

zrb/task/tcp_check.py

@@ -22,7 +22,7 @@ class TcpCheck(BaseTask):
         input: list[AnyInput] | AnyInput | None = None,
         env: list[AnyEnv] | AnyEnv | None = None,
         host: StrAttr = "localhost",
-        auto_render_host: bool = True,
+        render_host: bool = True,
         port: IntAttr = 80,
         interval: int = 5,
         execute_condition: bool | str | Callable[[Context], bool] = True,
@@ -43,14 +43,12 @@ class TcpCheck(BaseTask):
             fallback=fallback,
         )
         self._host = host
-        self._auto_render_host = auto_render_host
+        self._render_host = render_host
         self._port = port
         self._interval = interval
 
     def _get_host(self, ctx: AnyContext) -> str:
-        return get_str_attr(
-            ctx, self._host, "localhost", auto_render=self._auto_render_host
-        )
+        return get_str_attr(ctx, self._host, "localhost", auto_render=self._render_host)
 
     def _get_port(self, ctx: AnyContext) -> str:
         return get_int_attr(ctx, self._port, 80, auto_render=True)

zrb/util/todo.py

@@ -3,8 +3,16 @@ import re
 
 from pydantic import BaseModel, Field, model_validator
 
+from zrb.util.cli.style import (
+    stylize_bold_green,
+    stylize_cyan,
+    stylize_magenta,
+    stylize_yellow,
+)
+from zrb.util.string.name import get_random_name
+
 
-class TodoTask(BaseModel):
+class TodoTaskModel(BaseModel):
     priority: str | None = Field("D", pattern=r"^[A-Z]$")  # Priority like A, B, ...
     completed: bool = False  # True if completed, False otherwise
     description: str  # Main task description
@@ -34,16 +42,46 @@ TODO_TXT_PATTERN = re.compile(
 )
 
 
-def read_todo_from_file(todo_file_path: str) -> list[TodoTask]:
+def cascade_todo_task(todo_task: TodoTaskModel):
+    if todo_task.creation_date is None:
+        todo_task.creation_date = datetime.date.today()
+    if "id" not in todo_task.keyval:
+        todo_task.keyval["id"] = get_random_name()
+    return todo_task
+
+
+def select_todo_task(
+    todo_list: list[TodoTaskModel], keyword: str
+) -> TodoTaskModel | None:
+    for todo_task in todo_list:
+        id = todo_task.keyval.get("id", "")
+        if keyword.lower().strip() == id.lower().strip():
+            return todo_task
+    for todo_task in todo_list:
+        description = todo_task.description
+        if keyword.lower().strip() == description.lower().strip():
+            return todo_task
+    for todo_task in todo_list:
+        id = todo_task.keyval.get("id", "")
+        if keyword.lower().strip() in id.lower().strip():
+            return todo_task
+    for todo_task in todo_list:
+        description = todo_task.description
+        if keyword.lower().strip() in description.lower().strip():
+            return todo_task
+    return None
+
+
+def load_todo_list(todo_file_path: str) -> list[TodoTaskModel]:
     with open(todo_file_path, "r") as f:
         todo_lines = f.read().strip().split("\n")
-    todo_tasks: list[TodoTask] = []
+    todo_list: list[TodoTaskModel] = []
     for todo_line in todo_lines:
         todo_line = todo_line.strip()
         if todo_line == "":
             continue
-        todo_tasks.append(parse_todo_line(todo_line))
-    todo_tasks.sort(
+        todo_list.append(line_to_todo_task(todo_line))
+    todo_list.sort(
         key=lambda task: (
             task.completed,
             task.priority if task.priority else "Z",
@@ -51,16 +89,16 @@ def read_todo_from_file(todo_file_path: str) -> list[TodoTask]:
             task.creation_date if task.creation_date else datetime.date.max,
         )
     )
-    return todo_tasks
+    return todo_list
 
 
-def write_todo_to_file(todo_file_path: str, todo_task_list: list[TodoTask]):
+def save_todo_list(todo_file_path: str, todo_list: list[TodoTaskModel]):
     with open(todo_file_path, "w") as f:
-        for todo_task in todo_task_list:
-            f.write(todo_task_to_line(todo_task))
+        for todo_task in todo_list:
+            f.write(todo_task_to_line(todo_task) + "\n")
 
 
-def parse_todo_line(line: str) -> TodoTask:
+def line_to_todo_task(line: str) -> TodoTaskModel:
     """Parses a single todo.txt line into a TodoTask model."""
     match = TODO_TXT_PATTERN.match(line)
     if not match:
@@ -69,8 +107,8 @@ def parse_todo_line(line: str) -> TodoTask:
     # Extract completion status
     is_completed = groups["status"] == "x"
     # Extract dates
-    date1 = parse_date(groups["date1"])
-    date2 = parse_date(groups["date2"])
+    date1 = _parse_date(groups["date1"])
+    date2 = _parse_date(groups["date2"])
     # Determine creation_date and completion_date
     completion_date, creation_date = None, None
     if date2 is None:
@@ -87,7 +125,7 @@ def parse_todo_line(line: str) -> TodoTask:
         key, val = keyval_str.split(":", 1)
         keyval[key] = val
     description = re.sub(r"\s*\+\S+|\s*@\S+|\s*\S+:\S+", "", raw_description).strip()
-    return TodoTask(
+    return TodoTaskModel(
         priority=groups["priority"],
         completed=is_completed,
         description=description,
@@ -99,14 +137,14 @@ def parse_todo_line(line: str) -> TodoTask:
     )
 
 
-def parse_date(date_str: str | None) -> datetime.date | None:
+def _parse_date(date_str: str | None) -> datetime.date | None:
     """Parses a date string in the format YYYY-MM-DD."""
     if date_str:
         return datetime.date.fromisoformat(date_str)
     return None
 
 
-def todo_task_to_line(task: TodoTask) -> str:
+def todo_task_to_line(task: TodoTaskModel) -> str:
     """Converts a TodoTask instance back into a todo.txt formatted line."""
     parts = []
     # Add completion mark if task is completed
@@ -133,3 +171,89 @@ def todo_task_to_line(task: TodoTask) -> str:
         parts.append(f"{key}:{val}")
     # Join all parts with a space
     return " ".join(parts)
+
+
+def get_visual_todo_list(todo_list: list[TodoTaskModel]) -> str:
+    if len(todo_list) == 0:
+        return "\n".join(["", "  Empty todo list... 🌵🦖", ""])
+    max_desc_name_length = max(len(todo_task.description) for todo_task in todo_list)
+    if max_desc_name_length < len("DESCRIPTION"):
+        max_desc_name_length = len("DESCRIPTION")
+    # Headers
+    results = [
+        stylize_bold_green(
+            "  ".join(
+                [
+                    "".ljust(3),  # priority
+                    "".ljust(3),  # completed
+                    "COMPLETED AT".rjust(14),  # completed date
+                    "CREATED AT".rjust(14),  # completed date
+                    "DESCRIPTION".ljust(max_desc_name_length),
+                    "PROJECT/CONTEXT/OTHERS",
+                ]
+            )
+        )
+    ]
+    for todo_task in todo_list:
+        completed = "[x]" if todo_task.completed else "[ ]"
+        priority = "   " if todo_task.priority is None else f"({todo_task.priority})"
+        completion_date = stylize_yellow(_date_to_str(todo_task.completion_date))
+        creation_date = stylize_cyan(_date_to_str(todo_task.creation_date))
+        description = todo_task.description.ljust(max_desc_name_length)
+        additions = ", ".join(
+            [stylize_yellow(f"+{project}") for project in todo_task.projects]
+            + [stylize_cyan(f"@{context}") for context in todo_task.contexts]
+            + [stylize_magenta(f"{key}:{val}") for key, val in todo_task.keyval.items()]
+        )
+        results.append(
+            "  ".join(
+                [
+                    completed,
+                    priority,
+                    completion_date,
+                    creation_date,
+                    description,
+                    additions,
+                ]
+            )
+        )
+    return "\n".join(results)
+
+
+def _date_to_str(date: datetime.date | None) -> str:
+    if date is None:
+        return "".ljust(14)
+    return date.strftime("%a %Y-%m-%d")
+
+
+def add_durations(duration1: str, duration2: str) -> str:
+    total_seconds = _parse_duration(duration1) + _parse_duration(duration2)
+    # Format and return the result
+    return _format_duration(total_seconds)
+
+
+def _parse_duration(duration: str) -> int:
+    """Parse a duration string into total seconds."""
+    units = {"M": 2592000, "w": 604800, "d": 86400, "h": 3600, "m": 60, "s": 1}
+    total_seconds = 0
+    match = re.findall(r"(\d+)([Mwdhms])", duration)
+    for value, unit in match:
+        total_seconds += int(value) * units[unit]
+    return total_seconds
+
+
+def _format_duration(total_seconds: int) -> str:
+    """Format total seconds into a duration string."""
+    units = [
+        ("w", 604800),  # 7 days in a week
+        ("d", 86400),  # 24 hours in a day
+        ("h", 3600),  # 60 minutes in an hour
+        ("m", 60),  # 60 seconds in a minute
+        ("s", 1),  # seconds
+    ]
+    result = []
+    for unit, value_in_seconds in units:
+        if total_seconds >= value_in_seconds:
+            amount, total_seconds = divmod(total_seconds, value_in_seconds)
+            result.append(f"{amount}{unit}")
+    return "".join(result) if result else "0s"

{zrb-1.0.0a3.dist-info → zrb-1.0.0a4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: zrb
-Version: 1.0.0a3
+Version: 1.0.0a4
 Summary: Your Automation Powerhouse
 Home-page: https://github.com/state-alchemists/zrb
 License: AGPL-3.0-or-later
@@ -13,11 +13,16 @@ Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Provides-Extra: rag
 Requires-Dist: autopep8 (>=2.0.4,<3.0.0)
+Requires-Dist: beautifulsoup4 (>=4.12.3,<5.0.0)
 Requires-Dist: black (>=24.10.0,<24.11.0)
+Requires-Dist: chromadb (>=0.5.20,<0.6.0) ; extra == "rag"
+Requires-Dist: fastapi[standard] (>=0.115.5,<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: python-dotenv (>=1.0.1,<2.0.0)
 Requires-Dist: requests (>=2.32.3,<3.0.0)
 Project-URL: Documentation, https://github.com/state-alchemists/zrb