zrb 1.15.3__py3-none-any.whl → 1.21.29__py3-none-any.whl

zrb/task/llm_task.py

@@ -3,40 +3,44 @@ from collections.abc import Callable
 from typing import TYPE_CHECKING, Any
 
 from zrb.attr.type import BoolAttr, IntAttr, StrAttr, StrListAttr, fstring
-from zrb.config.llm_rate_limitter import LLMRateLimiter
+from zrb.config.llm_rate_limitter import LLMRateLimitter
 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.task.llm.agent import get_agent, run_agent_iteration
+from zrb.task.llm.agent import get_agent
+from zrb.task.llm.agent_runner import run_agent_iteration
 from zrb.task.llm.config import (
-    get_is_yolo_mode,
     get_model,
     get_model_settings,
+    get_yolo_mode,
 )
 from zrb.task.llm.conversation_history import (
+    inject_conversation_history_notes,
     read_conversation_history,
     write_conversation_history,
 )
 from zrb.task.llm.conversation_history_model import ConversationHistory
-from zrb.task.llm.history_summarization import maybe_summarize_history
+from zrb.task.llm.history_summarization import get_history_summarization_token_threshold
 from zrb.task.llm.prompt import (
     get_attachments,
     get_summarization_system_prompt,
     get_system_and_user_prompt,
     get_user_message,
 )
+from zrb.task.llm.subagent_conversation_history import (
+    extract_subagent_conversation_history_from_ctx,
+    inject_subagent_conversation_history_into_ctx,
+)
+from zrb.task.llm.workflow import load_workflow
 from zrb.util.cli.style import stylize_faint
 from zrb.xcom.xcom import Xcom
 
 if TYPE_CHECKING:
-    from pydantic_ai import Agent, Tool
-    from pydantic_ai.messages import UserContent
+    from pydantic_ai import AbstractToolset, Agent, Tool, UserContent
     from pydantic_ai.models import Model
     from pydantic_ai.settings import ModelSettings
-    from pydantic_ai.toolsets import AbstractToolset
 
     ToolOrCallable = Tool | Callable
 
@@ -51,62 +55,74 @@ class LLMTask(BaseTask):
         cli_only: bool = False,
         input: list[AnyInput | None] | AnyInput | None = None,
         env: list[AnyEnv | None] | AnyEnv | None = None,
-        model: (
-            "Callable[[AnySharedContext], Model | str | fstring] | Model | None"
-        ) = None,
+        model: "Callable[[AnyContext], Model | str | fstring | None] | Model | None" = None,
         render_model: bool = True,
-        model_base_url: StrAttr | None = None,
+        model_base_url: "Callable[[AnyContext], str | None] | str | None" = None,
         render_model_base_url: bool = True,
-        model_api_key: StrAttr | None = None,
+        model_api_key: "Callable[[AnyContext], str | None] | str | None" = None,
         render_model_api_key: bool = True,
         model_settings: (
-            "ModelSettings | Callable[[AnySharedContext], ModelSettings] | None"
+            "ModelSettings | Callable[[AnyContext], ModelSettings] | None"
+        ) = None,
+        small_model: (
+            "Callable[[AnyContext], Model | str | fstring] | Model | None"
         ) = None,
-        agent: "Agent | Callable[[AnySharedContext], Agent] | None" = None,
-        persona: StrAttr | None = None,
+        render_small_model: bool = True,
+        small_model_base_url: StrAttr | None = None,
+        render_small_model_base_url: bool = True,
+        small_model_api_key: StrAttr | None = None,
+        render_small_model_api_key: bool = True,
+        small_model_settings: (
+            "ModelSettings | Callable[[AnyContext], ModelSettings] | None"
+        ) = None,
+        persona: "Callable[[AnyContext], str | None] | str | None" = None,
         render_persona: bool = False,
-        system_prompt: StrAttr | None = None,
+        system_prompt: "Callable[[AnyContext], str | None] | str | None" = None,
         render_system_prompt: bool = False,
-        special_instruction_prompt: StrAttr | None = None,
+        special_instruction_prompt: "Callable[[AnyContext], str | None] | str | None" = None,
         render_special_instruction_prompt: bool = False,
-        modes: StrListAttr | None = None,
-        render_modes: bool = True,
+        workflows: StrListAttr | None = None,
+        render_workflows: bool = True,
         message: StrAttr | None = None,
-        attachment: "UserContent | list[UserContent] | Callable[[AnySharedContext], UserContent | list[UserContent]] | None" = None,  # noqa
+        attachment: "UserContent | list[UserContent] | Callable[[AnyContext], UserContent | list[UserContent]] | None" = None,  # noqa
         render_message: bool = True,
         tools: (
-            list["ToolOrCallable"]
-            | Callable[[AnySharedContext], list["ToolOrCallable"]]
+            list["ToolOrCallable"] | Callable[[AnyContext], list["ToolOrCallable"]]
         ) = [],
         toolsets: (
-            list["AbstractToolset[Agent]"] | Callable[[AnySharedContext], list["Tool"]]
+            list["AbstractToolset[None] | str"]
+            | Callable[[AnyContext], list["AbstractToolset[None] | str"]]
         ) = [],
         conversation_history: (
             ConversationHistory
-            | Callable[[AnySharedContext], ConversationHistory | dict | list]
+            | Callable[[AnyContext], ConversationHistory | dict | list]
             | dict
             | list
         ) = ConversationHistory(),
         conversation_history_reader: (
-            Callable[[AnySharedContext], ConversationHistory | dict | list | None]
-            | None
+            Callable[[AnyContext], ConversationHistory | dict | list | None] | None
         ) = None,
         conversation_history_writer: (
-            Callable[[AnySharedContext, ConversationHistory], None] | None
+            Callable[[AnyContext, ConversationHistory], None] | None
         ) = None,
         conversation_history_file: StrAttr | None = None,
         render_history_file: bool = True,
         summarize_history: BoolAttr | None = None,
         render_summarize_history: bool = True,
-        summarization_prompt: StrAttr | None = None,
+        summarization_prompt: "Callable[[AnyContext], str | None] | str | None" = None,
         render_summarization_prompt: bool = False,
         history_summarization_token_threshold: IntAttr | None = None,
         render_history_summarization_token_threshold: bool = True,
-        rate_limitter: LLMRateLimiter | None = None,
-        execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
+        rate_limitter: LLMRateLimitter | None = None,
+        execute_condition: bool | str | Callable[[AnyContext], bool] = True,
         retries: int = 2,
         retry_period: float = 0,
-        is_yolo_mode: bool | None = None,
+        yolo_mode: (
+            Callable[[AnyContext], list[str] | bool | None]
+            | StrListAttr
+            | BoolAttr
+            | None
+        ) = None,
         render_yolo_mode: bool = True,
         readiness_check: list[AnyTask] | AnyTask | None = None,
         readiness_check_delay: float = 0.5,
@@ -119,7 +135,7 @@ class LLMTask(BaseTask):
         fallback: list[AnyTask] | AnyTask | None = None,
         successor: list[AnyTask] | AnyTask | None = None,
         conversation_context: (
-            dict[str, Any] | Callable[[AnySharedContext], dict[str, Any]] | None
+            dict[str, Any] | Callable[[AnyContext], dict[str, Any]] | None
         ) = None,
     ):
         super().__init__(
@@ -150,24 +166,30 @@ class LLMTask(BaseTask):
         self._model_api_key = model_api_key
         self._render_model_api_key = render_model_api_key
         self._model_settings = model_settings
-        self._agent = agent
+        self._small_model = small_model
+        self._render_small_model = render_small_model
+        self._small_model_base_url = small_model_base_url
+        self._render_small_model_base_url = render_small_model_base_url
+        self._small_model_api_key = small_model_api_key
+        self._render_small_model_api_key = render_small_model_api_key
+        self._small_model_settings = small_model_settings
         self._persona = persona
         self._render_persona = render_persona
         self._system_prompt = system_prompt
         self._render_system_prompt = render_system_prompt
         self._special_instruction_prompt = special_instruction_prompt
         self._render_special_instruction_prompt = render_special_instruction_prompt
-        self._modes = modes
-        self._render_modes = render_modes
+        self._workflows = workflows
+        self._render_workflows = render_workflows
         self._message = message
         self._render_message = render_message
         self._summarization_prompt = summarization_prompt
         self._render_summarization_prompt = render_summarization_prompt
         self._tools = tools
         self._rate_limitter = rate_limitter
-        self._additional_tools: list["ToolOrCallable"] = []
+        self._additional_tools: list["ToolOrCallable"] = [load_workflow]
         self._toolsets = toolsets
-        self._additional_toolsets: list["AbstractToolset[Agent]"] = []
+        self._additional_toolsets: list["AbstractToolset[None] | str"] = []
         self._conversation_history = conversation_history
         self._conversation_history_reader = conversation_history_reader
         self._conversation_history_writer = conversation_history_writer
@@ -183,7 +205,7 @@ class LLMTask(BaseTask):
         )
         self._max_call_iteration = max_call_iteration
         self._conversation_context = conversation_context
-        self._is_yolo_mode = is_yolo_mode
+        self._yolo_mode = yolo_mode
         self._render_yolo_mode = render_yolo_mode
         self._attachment = attachment
 
@@ -194,10 +216,10 @@ class LLMTask(BaseTask):
         for single_tool in tool:
             self._additional_tools.append(single_tool)
 
-    def add_toolset(self, *toolset: "AbstractToolset[Agent]"):
+    def add_toolset(self, *toolset: "AbstractToolset[None] | str"):
         self.append_toolset(*toolset)
 
-    def append_toolset(self, *toolset: "AbstractToolset[Agent]"):
+    def append_toolset(self, *toolset: "AbstractToolset[None] | str"):
         for single_toolset in toolset:
             self._additional_toolsets.append(single_toolset)
 
@@ -209,6 +231,12 @@ class LLMTask(BaseTask):
     ):
         self._history_summarization_token_threshold = summarization_token_threshold
 
+    def set_workflows(self, workflows: StrListAttr):
+        self._workflows = workflows
+
+    def set_yolo_mode(self, yolo_mode: StrListAttr | BoolAttr):
+        self._yolo_mode = yolo_mode
+
     async def _exec_action(self, ctx: AnyContext) -> Any:
         # Get dependent configurations first
         model_settings = get_model_settings(ctx, self._model_settings)
@@ -221,9 +249,9 @@ class LLMTask(BaseTask):
             model_api_key_attr=self._model_api_key,
             render_model_api_key=self._render_model_api_key,
         )
-        is_yolo_mode = get_is_yolo_mode(
+        yolo_mode = get_yolo_mode(
             ctx=ctx,
-            is_yolo_mode_attr=self._is_yolo_mode,
+            yolo_mode_attr=self._yolo_mode,
             render_yolo_mode=self._render_yolo_mode,
         )
         summarization_prompt = get_summarization_system_prompt(
@@ -241,9 +269,10 @@ class LLMTask(BaseTask):
             render_history_file=self._render_history_file,
             conversation_history_attr=self._conversation_history,
         )
-        conversation_history.fetch_newest_notes()
+        inject_conversation_history_notes(conversation_history)
+        inject_subagent_conversation_history_into_ctx(ctx, conversation_history)
         # 2. Get system prompt and user prompt
-        system_prompt, user_message = get_system_and_user_prompt(
+        system_prompt, user_prompt = get_system_and_user_prompt(
             ctx=ctx,
             user_message=user_message,
             persona_attr=self._persona,
@@ -252,50 +281,59 @@ class LLMTask(BaseTask):
             render_system_prompt=self._render_system_prompt,
             special_instruction_prompt_attr=self._special_instruction_prompt,
             render_special_instruction_prompt=self._render_special_instruction_prompt,
-            modes_attr=self._modes,
-            render_modes=self._render_modes,
+            workflows_attr=self._workflows,
+            render_workflows=self._render_workflows,
             conversation_history=conversation_history,
         )
-        ctx.log_debug(f"SYSTEM PROMPT: {system_prompt}")
-        # 3. Get the agent instance
+        # 3. Summarization
+        small_model = get_model(
+            ctx=ctx,
+            model_attr=self._small_model,
+            render_model=self._render_small_model,
+            model_base_url_attr=self._small_model_base_url,
+            render_model_base_url=self._render_small_model_base_url,
+            model_api_key_attr=self._small_model_api_key,
+            render_model_api_key=self._render_small_model_api_key,
+        )
+        small_model_settings = get_model_settings(ctx, self._small_model_settings)
+        summarization_token_threshold = get_history_summarization_token_threshold(
+            ctx,
+            self._history_summarization_token_threshold,
+            self._render_history_summarization_token_threshold,
+        )
+        # 4. Get the agent instance
+        ctx.log_info(f"SYSTEM PROMPT:\n{system_prompt}")
+        ctx.log_info(f"USER PROMPT:\n{user_prompt}")
         agent = get_agent(
             ctx=ctx,
-            agent_attr=self._agent,
             model=model,
+            rate_limitter=self._rate_limitter,
             system_prompt=system_prompt,
             model_settings=model_settings,
             tools_attr=self._tools,
             additional_tools=self._additional_tools,
             toolsets_attr=self._toolsets,
             additional_toolsets=self._additional_toolsets,
-            is_yolo_mode=is_yolo_mode,
+            yolo_mode=yolo_mode,
+            summarization_model=small_model,
+            summarization_model_settings=small_model_settings,
+            summarization_system_prompt=summarization_prompt,
+            summarization_retries=2,  # TODO: make this a property
+            summarization_token_threshold=summarization_token_threshold,
+            history_processors=[],  # TODO: make this a property
         )
-        # 4. Run the agent iteration and save the results/history
+        # 5. Run the agent iteration and save the results/history
         result = await self._execute_agent(
             ctx=ctx,
             agent=agent,
-            user_prompt=user_message,
+            user_prompt=user_prompt,
             attachments=attachments,
             conversation_history=conversation_history,
         )
-        # 5. Summarize
-        conversation_history = await maybe_summarize_history(
-            ctx=ctx,
-            conversation_history=conversation_history,
-            should_summarize_history_attr=self._should_summarize_history,
-            render_summarize_history=self._render_summarize_history,
-            history_summarization_token_threshold_attr=(
-                self._history_summarization_token_threshold
-            ),
-            render_history_summarization_token_threshold=(
-                self._render_history_summarization_token_threshold
-            ),
-            model=model,
-            model_settings=model_settings,
-            summarization_prompt=summarization_prompt,
-            rate_limitter=self._rate_limitter,
-        )
         # 6. Write conversation history
+        conversation_history.subagent_history = (
+            extract_subagent_conversation_history_from_ctx(ctx)
+        )
         await write_conversation_history(
             ctx=ctx,
             history_data=conversation_history,
@@ -331,7 +369,7 @@ class LLMTask(BaseTask):
                     ctx.xcom[xcom_usage_key] = Xcom([])
                 usage = agent_run.result.usage()
                 ctx.xcom[xcom_usage_key].push(usage)
-                ctx.print(stylize_faint(f" 💸 Token: {usage}"), plain=True)
+                ctx.print(stylize_faint(f"  💸 Token: {usage}"), plain=True)
                 return agent_run.result.output
             else:
                 ctx.log_warning("Agent run did not produce a result.")

zrb/task/make_task.py

@@ -2,7 +2,6 @@ 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
@@ -18,7 +17,7 @@ def make_task(
     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,
+    execute_condition: bool | str | Callable[[AnyContext], bool] = True,
     retries: int = 2,
     retry_period: float = 0,
     readiness_check: list[AnyTask] | AnyTask | None = None,
@@ -33,7 +32,7 @@ def make_task(
     group: AnyGroup | None = None,
     alias: str | None = None,
 ) -> Callable[[Callable[[AnyContext], Any]], AnyTask]:
-    def _make_task(fn: Callable[[AnyContext], Any]) -> BaseTask:
+    def _make_task(fn: Callable[[AnyContext], Any]) -> AnyTask:
         task = BaseTask(
             name=name,
             color=color,

zrb/task/rsync_task.py

@@ -1,6 +1,4 @@
-from collections.abc import Callable
-
-from zrb.attr.type import IntAttr, StrAttr
+from zrb.attr.type import BoolAttr, IntAttr, StrAttr
 from zrb.context.any_context import AnyContext
 from zrb.env.any_env import AnyEnv
 from zrb.input.any_input import AnyInput
@@ -17,8 +15,8 @@ class RsyncTask(CmdTask):
         icon: str | None = None,
         description: str | None = None,
         cli_only: bool = False,
-        input: list[AnyInput] | AnyInput | None = None,
-        env: list[AnyEnv] | AnyEnv | None = None,
+        input: list[AnyInput | None] | AnyInput | None = None,
+        env: list[AnyEnv | None] | AnyEnv | None = None,
         shell: StrAttr | None = None,
         auto_render_shell: bool = True,
         remote_host: StrAttr | None = None,
@@ -39,12 +37,14 @@ class RsyncTask(CmdTask):
         render_local_source_path: bool = True,
         local_destination_path: StrAttr | None = None,
         render_local_destination_path: bool = True,
+        exclude_from: StrAttr | None = None,
+        render_exclude_from: bool = True,
         cwd: str | None = None,
         render_cwd: bool = True,
         plain_print: bool = False,
         max_output_line: int = 1000,
         max_error_line: int = 1000,
-        execute_condition: bool | str | Callable[[AnyContext], bool] = True,
+        execute_condition: BoolAttr = True,
         retries: int = 2,
         retry_period: float = 0,
         readiness_check: list[AnyTask] | AnyTask | None = None,
@@ -93,6 +93,8 @@ class RsyncTask(CmdTask):
         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
+        self._exclude_from = exclude_from
+        self._render_exclude_from = render_exclude_from
 
     def _get_source_path(self, ctx: AnyContext) -> str:
         local_source_path = self._get_local_source_path(ctx)
@@ -144,16 +146,29 @@ class RsyncTask(CmdTask):
             auto_render=self._render_local_destination_path,
         )
 
+    def _get_exclude_from_param(self, ctx: AnyContext) -> str:
+        exclude_from = get_str_attr(
+            ctx,
+            self._exclude_from,
+            "",
+            auto_render=self._render_exclude_from,
+        ).strip()
+        if exclude_from == "":
+            return ""
+        return f"--exclude-from='{exclude_from}'"
+
     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)
+        exclude_from = self._get_exclude_from_param(ctx)
+        exclude_from_with_space = f"{exclude_from} " if exclude_from != "" else ""
         if key != "" and password != "":
-            return f'sshpass -p "$_ZRB_SSH_PASSWORD" rsync --mkpath -avz -e "ssh -i {key} -p {port}" {src} {dst}'  # noqa
+            return f'sshpass -p "$_ZRB_SSH_PASSWORD" rsync --mkpath -avz -e "ssh -i {key} -p {port}" {exclude_from_with_space}{src} {dst}'  # noqa
         if key != "":
-            return f'rsync --mkpath -avz -e "ssh -i {key} -p {port}" {src} {dst}'
+            return f'rsync --mkpath -avz -e "ssh -i {key} -p {port}" {exclude_from_with_space}{src} {dst}'  # noqa
         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}'
+            return f'sshpass -p "$_ZRB_SSH_PASSWORD" rsync --mkpath -avz -e "ssh -p {port}" {exclude_from_with_space}{src} {dst}'  # noqa
+        return f'rsync --mkpath -avz -e "ssh -p {port}" {exclude_from_with_space}{src} {dst}'

zrb/task/scheduler.py

@@ -24,8 +24,8 @@ class Scheduler(BaseTrigger):
         cli_only: bool = False,
         input: list[AnyInput | None] | AnyInput | None = None,
         env: list[AnyEnv | None] | AnyEnv | None = None,
-        schedule: StrAttr = None,
-        execute_condition: bool | str | Callable[[AnySharedContext], bool] = True,
+        schedule: StrAttr | None = None,
+        execute_condition: bool | str | Callable[[AnyContext], bool] = True,
         queue_name: fstring | None = None,
         callback: list[AnyCallback] | AnyCallback = [],
         retries: int = 2,
@@ -76,6 +76,6 @@ class Scheduler(BaseTrigger):
             ctx.print(f"Current time: {now}")
             if match_cron(cron_pattern, now):
                 ctx.print(f"Matching {now} with pattern: {cron_pattern}")
-                xcom = self.get_exchange_xcom(ctx.session)
-                xcom.push(now)
+                if ctx.session is not None:
+                    self.push_exchange_xcom(ctx.session, now)
             await asyncio.sleep(60)