zrb 1.21.29__py3-none-any.whl → 2.0.0a4__py3-none-any.whl

zrb/__init__.py

@@ -1,130 +1,119 @@
-import importlib
-from typing import TYPE_CHECKING, Any
-
-_LAZY_LOAD = {
-    "AnyAttr": "zrb.attr.type",
-    "BoolAttr": "zrb.attr.type",
-    "FloatAttr": "zrb.attr.type",
-    "IntAttr": "zrb.attr.type",
-    "StrAttr": "zrb.attr.type",
-    "StrDictAttr": "zrb.attr.type",
-    "fstring": "zrb.attr.type",
-    "AnyCallback": "zrb.callback.any_callback",
-    "Callback": "zrb.callback.callback",
-    "CmdResult": "zrb.cmd.cmd_result",
-    "Cmd": "zrb.cmd.cmd_val",
-    "CmdPath": "zrb.cmd.cmd_val",
-    "CFG": "zrb.config.config",
-    "AnyContentTransformer": "zrb.content_transformer.any_content_transformer",
-    "ContentTransformer": "zrb.content_transformer.content_transformer",
-    "AnyContext": "zrb.context.any_context",
-    "AnySharedContext": "zrb.context.any_shared_context",
-    "Context": "zrb.context.context",
-    "SharedContext": "zrb.context.shared_context",
-    "AnyEnv": "zrb.env.any_env",
-    "Env": "zrb.env.env",
-    "EnvFile": "zrb.env.env_file",
-    "EnvMap": "zrb.env.env_map",
-    "AnyGroup": "zrb.group.any_group",
-    "Group": "zrb.group.group",
-    "AnyInput": "zrb.input.any_input",
-    "BaseInput": "zrb.input.base_input",
-    "BoolInput": "zrb.input.bool_input",
-    "FloatInput": "zrb.input.float_input",
-    "IntInput": "zrb.input.int_input",
-    "OptionInput": "zrb.input.option_input",
-    "PasswordInput": "zrb.input.password_input",
-    "StrInput": "zrb.input.str_input",
-    "TextInput": "zrb.input.text_input",
-    "llm_config": "zrb.config.llm_config",
-    "llm_rate_limitter": "zrb.config.llm_rate_limitter",
-    "cli": "zrb.runner.cli",
-    "web_auth_config": "zrb.config.web_auth_config",
-    "User": "zrb.runner.web_schema.user",
-    "Session": "zrb.session.session",
-    "AnyTask": "zrb.task.any_task",
-    "BaseTask": "zrb.task.base_task",
-    "BaseTrigger": "zrb.task.base_trigger",
-    "CmdTask": "zrb.task.cmd_task",
-    "HttpCheck": "zrb.task.http_check",
-    "ConversationHistory": "zrb.task.llm.conversation_history",
-    "LLMTask": "zrb.task.llm_task",
-    "make_task": "zrb.task.make_task",
-    "RsyncTask": "zrb.task.rsync_task",
-    "Scaffolder": "zrb.task.scaffolder",
-    "Scheduler": "zrb.task.scheduler",
-    "Task": "zrb.task.task",
-    "TcpCheck": "zrb.task.tcp_check",
-    "load_file": "zrb.util.load",
-    "load_module": "zrb.util.load",
-    "Xcom": "zrb.xcom.xcom",
-}
-
-if TYPE_CHECKING:
-    from zrb import builtin
-    from zrb.attr.type import (
-        AnyAttr,
-        BoolAttr,
-        FloatAttr,
-        IntAttr,
-        StrAttr,
-        StrDictAttr,
-        fstring,
-    )
-    from zrb.callback.any_callback import AnyCallback
-    from zrb.callback.callback import Callback
-    from zrb.cmd.cmd_result import CmdResult
-    from zrb.cmd.cmd_val import Cmd, CmdPath
-    from zrb.config.config import CFG
-    from zrb.config.llm_config import llm_config
-    from zrb.config.llm_rate_limitter import llm_rate_limitter
-    from zrb.config.web_auth_config import web_auth_config
-    from zrb.content_transformer.any_content_transformer import AnyContentTransformer
-    from zrb.content_transformer.content_transformer import ContentTransformer
-    from zrb.context.any_context import AnyContext
-    from zrb.context.any_shared_context import AnySharedContext
-    from zrb.context.context import Context
-    from zrb.context.shared_context import SharedContext
-    from zrb.env.any_env import AnyEnv
-    from zrb.env.env import Env
-    from zrb.env.env_file import EnvFile
-    from zrb.env.env_map import EnvMap
-    from zrb.group.any_group import AnyGroup
-    from zrb.group.group import Group
-    from zrb.input.any_input import AnyInput
-    from zrb.input.base_input import BaseInput
-    from zrb.input.bool_input import BoolInput
-    from zrb.input.float_input import FloatInput
-    from zrb.input.int_input import IntInput
-    from zrb.input.option_input import OptionInput
-    from zrb.input.password_input import PasswordInput
-    from zrb.input.str_input import StrInput
-    from zrb.input.text_input import TextInput
-    from zrb.runner.cli import cli
-    from zrb.runner.web_schema.user import User
-    from zrb.session.session import Session
-    from zrb.task.any_task import AnyTask
-    from zrb.task.base_task import BaseTask
-    from zrb.task.base_trigger import BaseTrigger
-    from zrb.task.cmd_task import CmdTask
-    from zrb.task.http_check import HttpCheck
-    from zrb.task.llm.conversation_history import ConversationHistory
-    from zrb.task.llm_task import LLMTask
-    from zrb.task.make_task import make_task
-    from zrb.task.rsync_task import RsyncTask
-    from zrb.task.scaffolder import Scaffolder
-    from zrb.task.scheduler import Scheduler
-    from zrb.task.task import Task
-    from zrb.task.tcp_check import TcpCheck
-    from zrb.util.load import load_file, load_module
-    from zrb.xcom.xcom import Xcom
-
-
-def __getattr__(name: str) -> Any:
-    if name in _LAZY_LOAD:
-        module = importlib.import_module(_LAZY_LOAD[name])
-        return getattr(module, name)
-    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
-
-
 from zrb import builtin
+from zrb.attr.type import (
+    AnyAttr,
+    BoolAttr,
+    FloatAttr,
+    IntAttr,
+    StrAttr,
+    StrDictAttr,
+    fstring,
+)
+from zrb.callback.any_callback import AnyCallback
+from zrb.callback.callback import Callback
+from zrb.cmd.cmd_result import CmdResult
+from zrb.cmd.cmd_val import Cmd, CmdPath
+from zrb.config.config import CFG
+from zrb.config.web_auth_config import web_auth_config
+from zrb.content_transformer.any_content_transformer import AnyContentTransformer
+from zrb.content_transformer.content_transformer import ContentTransformer
+from zrb.context.any_context import AnyContext
+from zrb.context.any_shared_context import AnySharedContext
+from zrb.context.context import Context
+from zrb.context.shared_context import SharedContext
+from zrb.env.any_env import AnyEnv
+from zrb.env.env import Env
+from zrb.env.env_file import EnvFile
+from zrb.env.env_map import EnvMap
+from zrb.group.any_group import AnyGroup
+from zrb.group.group import Group
+from zrb.input.any_input import AnyInput
+from zrb.input.base_input import BaseInput
+from zrb.input.bool_input import BoolInput
+from zrb.input.float_input import FloatInput
+from zrb.input.int_input import IntInput
+from zrb.input.option_input import OptionInput
+from zrb.input.password_input import PasswordInput
+from zrb.input.str_input import StrInput
+from zrb.input.text_input import TextInput
+from zrb.llm.config.config import LLMConfig, llm_config
+from zrb.llm.config.limiter import LLMLimiter, llm_limiter
+from zrb.llm.task.llm_chat_task import LLMChatTask
+from zrb.llm.task.llm_task import LLMTask
+from zrb.runner.cli import cli
+from zrb.runner.web_schema.user import User
+from zrb.session.session import Session
+from zrb.task.any_task import AnyTask
+from zrb.task.base_task import BaseTask
+from zrb.task.base_trigger import BaseTrigger
+from zrb.task.cmd_task import CmdTask
+from zrb.task.http_check import HttpCheck
+from zrb.task.make_task import make_task
+from zrb.task.rsync_task import RsyncTask
+from zrb.task.scaffolder import Scaffolder
+from zrb.task.scheduler import Scheduler
+from zrb.task.task import Task
+from zrb.task.tcp_check import TcpCheck
+from zrb.util.load import load_file, load_module
+from zrb.xcom.xcom import Xcom
+
+__all__ = [
+    "builtin",
+    "AnyAttr",
+    "BoolAttr",
+    "FloatAttr",
+    "IntAttr",
+    "StrAttr",
+    "StrDictAttr",
+    "fstring",
+    "AnyCallback",
+    "Callback",
+    "CmdResult",
+    "Cmd",
+    "CmdPath",
+    "CFG",
+    "web_auth_config",
+    "AnyContentTransformer",
+    "ContentTransformer",
+    "AnyContext",
+    "AnySharedContext",
+    "Context",
+    "SharedContext",
+    "AnyEnv",
+    "Env",
+    "EnvFile",
+    "EnvMap",
+    "AnyGroup",
+    "Group",
+    "AnyInput",
+    "BaseInput",
+    "BoolInput",
+    "FloatInput",
+    "IntInput",
+    "OptionInput",
+    "PasswordInput",
+    "StrInput",
+    "TextInput",
+    "cli",
+    "User",
+    "Session",
+    "AnyTask",
+    "BaseTask",
+    "BaseTrigger",
+    "CmdTask",
+    "HttpCheck",
+    "make_task",
+    "RsyncTask",
+    "Scaffolder",
+    "Scheduler",
+    "Task",
+    "TcpCheck",
+    "load_file",
+    "load_module",
+    "Xcom",
+    "LLMTask",
+    "LLMChatTask",
+    "LLMConfig",
+    "llm_config",
+    "LLMLimiter",
+    "llm_limiter",
+]

zrb/builtin/__init__.py

@@ -9,8 +9,7 @@ from zrb.builtin.git import (
 from zrb.builtin.git_subtree import git_add_subtree, git_pull_subtree, git_push_subtree
 from zrb.builtin.http import generate_curl, http_request
 from zrb.builtin.jwt import decode_jwt, encode_jwt, validate_jwt
-from zrb.builtin.llm.chat_trigger import llm_chat_trigger
-from zrb.builtin.llm.llm_ask import llm_ask
+from zrb.builtin.llm.chat import llm_chat
 from zrb.builtin.md5 import hash_md5, sum_md5, validate_md5
 from zrb.builtin.project.add.fastapp.fastapp_task import add_fastapp_to_project
 from zrb.builtin.project.create.project_task import create_project
@@ -45,3 +44,56 @@ from zrb.builtin.uuid import (
     validate_uuid_v4,
     validate_uuid_v5,
 )
+
+__all__ = [
+    "decode_base64",
+    "encode_base64",
+    "validate_base64",
+    "get_git_diff",
+    "git_commit",
+    "git_pull",
+    "git_push",
+    "prune_local_branches",
+    "git_add_subtree",
+    "git_pull_subtree",
+    "git_push_subtree",
+    "generate_curl",
+    "http_request",
+    "decode_jwt",
+    "encode_jwt",
+    "validate_jwt",
+    "llm_chat",
+    "hash_md5",
+    "sum_md5",
+    "validate_md5",
+    "add_fastapp_to_project",
+    "create_project",
+    "format_python_code",
+    "shuffle_values",
+    "throw_dice",
+    "start_searxng",
+    "setup_asdf",
+    "setup_latex_on_ubuntu",
+    "setup_tmux",
+    "setup_ubuntu",
+    "setup_zsh",
+    "make_bash_autocomplete",
+    "get_shell_subcommands",
+    "make_zsh_autocomplete",
+    "add_todo",
+    "archive_todo",
+    "complete_todo",
+    "edit_todo",
+    "list_todo",
+    "log_todo",
+    "show_todo",
+    "generate_uuid_v1",
+    "generate_uuid_v3",
+    "generate_uuid_v4",
+    "generate_uuid_v5",
+    "validate_uuid",
+    "validate_uuid_v1",
+    "validate_uuid_v3",
+    "validate_uuid_v4",
+    "validate_uuid_v5",
+]

zrb/builtin/llm/chat.py

@@ -0,0 +1,147 @@
+from zrb.builtin.group import llm_group
+from zrb.config.config import CFG
+from zrb.input.bool_input import BoolInput
+from zrb.input.str_input import StrInput
+from zrb.llm.app.confirmation.replace_confirmation import replace_confirmation
+from zrb.llm.history_processor.summarizer import create_summarizer_history_processor
+from zrb.llm.note.manager import NoteManager
+from zrb.llm.prompt.claude_compatibility import (
+    create_claude_compatibility_prompt,
+)
+from zrb.llm.prompt.compose import PromptManager, new_prompt
+from zrb.llm.prompt.default import (
+    get_mandate_prompt,
+    get_persona_prompt,
+)
+from zrb.llm.prompt.note import create_note_prompt
+from zrb.llm.prompt.system_context import system_context
+from zrb.llm.prompt.zrb import create_zrb_prompt
+from zrb.llm.skill.manager import SkillManager
+from zrb.llm.task.llm_chat_task import LLMChatTask
+from zrb.llm.tool.bash import run_shell_command
+from zrb.llm.tool.code import analyze_code
+from zrb.llm.tool.file import (
+    analyze_file,
+    list_files,
+    read_file,
+    read_files,
+    replace_in_file,
+    search_files,
+    write_file,
+    write_files,
+)
+from zrb.llm.tool.note import create_note_tools
+from zrb.llm.tool.skill import create_activate_skill_tool
+from zrb.llm.tool.sub_agent import create_sub_agent_tool
+from zrb.llm.tool.web import open_web_page, search_internet
+from zrb.llm.tool.zrb_task import create_list_zrb_task_tool, create_run_zrb_task_tool
+from zrb.runner.cli import cli
+
+skill_manager = SkillManager()
+note_manager = NoteManager()
+
+llm_chat = LLMChatTask(
+    name="chat",
+    description="🤖 Chat with your AI Assistant",
+    input=[
+        StrInput("message", "Message", allow_empty=True, always_prompt=False),
+        StrInput(
+            "session", "Conversation Session", allow_empty=True, always_prompt=False
+        ),
+        BoolInput(
+            "yolo", "YOLO Mode", default=False, allow_empty=True, always_prompt=False
+        ),
+        StrInput("attach", "Attachments", allow_empty=True, always_prompt=False),
+        BoolInput(
+            "interactive",
+            "Interactive Mode",
+            default=True,
+            allow_empty=True,
+            always_prompt=False,
+        ),
+    ],
+    yolo="{ctx.input.yolo}",
+    message="{ctx.input.message}",
+    conversation_name="{ctx.input.session}",
+    interactive="{ctx.input.interactive}",
+    history_processors=[
+        create_summarizer_history_processor(
+            token_threshold=CFG.LLM_HISTORY_SUMMARIZATION_TOKEN_THRESHOLD,
+            summary_window=CFG.LLM_HISTORY_SUMMARIZATION_WINDOW,
+        )
+    ],
+    prompt_manager=PromptManager(),
+    ui_ascii_art=lambda ctx: CFG.LLM_ASSISTANT_ASCII_ART,
+    ui_assistant_name=lambda ctx: CFG.LLM_ASSISTANT_NAME,
+    ui_greeting=lambda ctx: f"{CFG.LLM_ASSISTANT_NAME}\n{CFG.LLM_ASSISTANT_JARGON}",
+    ui_jargon=lambda ctx: CFG.LLM_ASSISTANT_JARGON,
+    ui_summarize_commands=["/compress", "/compact"],
+    ui_attach_commands=["/attach"],
+    ui_exit_commands=["/q", "/bye", "/quit", "/exit"],
+    ui_info_commands=["/info", "/help"],
+    ui_save_commands=["/save"],
+    ui_load_commands=["/load"],
+    ui_yolo_toggle_commands=["/yolo"],
+    ui_redirect_output_commands=[">", "/redirect"],
+    ui_exec_commands=["!", "/exec"],
+)
+llm_group.add_task(llm_chat)
+cli.add_task(llm_chat)
+
+
+async def roll_dice() -> str:
+    """Roll a six-sided die and return the result."""
+    import asyncio
+    import random
+
+    await asyncio.sleep(3)
+    return str(random.randint(1, 6))
+
+
+async def get_current_time() -> str:
+    """Get the current time."""
+    from datetime import datetime
+
+    return datetime.now().strftime("%H:%M:%S")
+
+
+joke_agent = create_sub_agent_tool(
+    name="joke_agent",
+    description="Generates jokes about the current directory content.",
+    system_prompt=(
+        "You are a comedian. Use the 'run_shell_command' tool to list files "
+        "(ls -la) and make a funny joke about the project structure."
+    ),
+    tools=[run_shell_command],
+)
+
+llm_chat.prompt_manager.add_middleware(
+    new_prompt(get_persona_prompt(CFG.LLM_ASSISTANT_NAME)),
+    new_prompt(get_mandate_prompt()),
+    system_context,
+    create_note_prompt(note_manager),
+    create_claude_compatibility_prompt(skill_manager),
+    create_zrb_prompt(),
+)
+llm_chat.add_confirmation_middleware(replace_confirmation)
+llm_chat.add_tool(
+    roll_dice,
+    joke_agent,
+    run_shell_command,
+    list_files,
+    read_file,
+    read_files,
+    write_file,
+    write_files,
+    replace_in_file,
+    search_files,
+    analyze_file,
+    analyze_code,
+    search_internet,
+    open_web_page,
+    create_list_zrb_task_tool(),
+    create_run_zrb_task_tool(),
+    create_activate_skill_tool(skill_manager),
+    *create_note_tools(note_manager),
+)
+llm_chat.add_tool(get_current_time)

zrb/callback/callback.py

@@ -6,7 +6,6 @@ from zrb.callback.any_callback import AnyCallback
 from zrb.session.any_session import AnySession
 from zrb.task.any_task import AnyTask
 from zrb.util.attr import get_str_dict_attr
-from zrb.util.cli.style import stylize_faint
 from zrb.util.string.conversion import to_snake_case
 from zrb.xcom.xcom import Xcom
 
@@ -24,6 +23,7 @@ class Callback(AnyCallback):
         task: AnyTask,
         input_mapping: StrDictAttr,
         render_input_mapping: bool = True,
+        xcom_mapping: dict[str, str] | None = None,
         result_queue: str | None = None,
         error_queue: str | None = None,
         session_name_queue: str | None = None,
@@ -36,6 +36,7 @@ class Callback(AnyCallback):
             input_mapping: A dictionary or attribute mapping to prepare inputs for the task.
             render_input_mapping: Whether to render the input mapping using
                 f-string like syntax.
+            xcom_mapping: Map of parent session's xcom names to current session's xcom names
             result_queue: The name of the XCom queue in the parent session
                 to publish the task result.
             result_queue: The name of the Xcom queue in the parent session
@@ -46,6 +47,7 @@ class Callback(AnyCallback):
         self._task = task
         self._input_mapping = input_mapping
         self._render_input_mapping = render_input_mapping
+        self._xcom_mapping = xcom_mapping
         self._result_queue = result_queue
         self._error_queue = error_queue
         self._session_name_queue = session_name_queue
@@ -63,6 +65,11 @@ class Callback(AnyCallback):
         for name, value in inputs.items():
             session.shared_ctx.input[name] = value
             session.shared_ctx.input[to_snake_case(name)] = value
+        # map xcom
+        if self._xcom_mapping is not None:
+            for parent_xcom_name, current_xcom_name in self._xcom_mapping.items():
+                parent_xcom = parent_session.shared_ctx.xcom[parent_xcom_name]
+                session.shared_ctx.xcom[current_xcom_name] = parent_xcom
         # run task and get result
         try:
             result = await self._task.async_run(session)

zrb/cmd/cmd_result.py

@@ -1,7 +1,8 @@
 class CmdResult:
-    def __init__(self, output: str, error: str):
+    def __init__(self, output: str, error: str, display: str):
         self.output = output
         self.error = error
+        self.display = display
 
     def __repr__(self):
         class_name = self.__class__.__name__