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

zrb/llm/prompt/note.py

@@ -0,0 +1,41 @@
+import os
+from typing import Callable
+
+from zrb.context.any_context import AnyContext
+from zrb.llm.note.manager import NoteManager
+from zrb.util.markdown import make_markdown_section
+
+
+def create_note_prompt(note_manager: NoteManager):
+    def note_prompt(
+        ctx: AnyContext,
+        current_prompt: str,
+        next_handler: Callable[[AnyContext, str], str],
+    ) -> str:
+        cwd = os.getcwd()
+        notes = note_manager.read_all(cwd)
+
+        if not notes:
+            return next_handler(ctx, current_prompt)
+
+        # Format notes
+        note_lines = []
+
+        # Sort keys for deterministic output
+        sorted_keys = sorted(notes.keys())
+
+        for key in sorted_keys:
+            content = notes[key].strip()
+            if not content:
+                continue
+            note_lines.append(f"**{key}**:\n{content}\n")
+
+        if not note_lines:
+            return next_handler(ctx, current_prompt)
+
+        full_note_content = "\n".join(note_lines)
+        note_block = make_markdown_section("Notes & Context", full_note_content)
+
+        return next_handler(ctx, f"{current_prompt}\n\n{note_block}")
+
+    return note_prompt

zrb/llm/prompt/system_context.py

@@ -0,0 +1,46 @@
+import os
+import platform
+import subprocess
+from datetime import datetime
+from typing import Callable
+
+from zrb.context.any_context import AnyContext
+from zrb.util.markdown import make_markdown_section
+
+
+def system_context(
+    ctx: AnyContext, current_prompt: str, next_handler: Callable[[AnyContext, str], str]
+) -> str:
+    # Time
+    now = datetime.now().strftime("%A, %B %d, %Y %H:%M:%S")
+
+    # System
+    os_info = f"{platform.system()} {platform.release()}"
+
+    # CWD
+    cwd = os.getcwd()
+
+    # Git
+    git_info = ""
+    try:
+        if os.path.isdir(os.path.join(cwd, ".git")):
+            # Get branch
+            res = subprocess.run(
+                ["git", "branch", "--show-current"], capture_output=True, text=True
+            )
+            if res.returncode == 0:
+                branch = res.stdout.strip()
+                if branch:
+                    git_info = f"\n- Git Branch: {branch}"
+    except Exception:
+        pass
+
+    info_content = "\n".join(
+        [
+            f"- Date: {now}",
+            f"- OS: {os_info}",
+            f"- Directory: {cwd}{git_info}",
+        ]
+    )
+    info_block = make_markdown_section("System Context", info_content)
+    return next_handler(ctx, f"{current_prompt}\n\n{info_block}")

zrb/llm/prompt/zrb.py

@@ -0,0 +1,41 @@
+from typing import Callable
+
+from zrb.config.config import CFG
+from zrb.context.any_context import AnyContext
+from zrb.runner.cli import cli
+from zrb.util.markdown import make_markdown_section
+
+
+def create_zrb_prompt():
+    def zrb_prompt(
+        ctx: AnyContext,
+        current_prompt: str,
+        next_handler: Callable[[AnyContext, str], str],
+    ) -> str:
+        additional_context = []
+        # Available Zrb Workflows (Skills)
+        zrb_cmd = CFG.ROOT_GROUP_NAME
+        zrb_context = []
+
+        # List top-level groups
+        if cli.subgroups:
+            zrb_context.append(f"Groups (Use 'list_{zrb_cmd}_tasks' to see details):")
+            for alias, grp in cli.subgroups.items():
+                zrb_context.append(f"- {alias}: {grp.description}")
+
+        # List top-level tasks
+        if cli.subtasks:
+            zrb_context.append("Tasks:")
+            for alias, task in cli.subtasks.items():
+                zrb_context.append(f"- {alias}: {task.description}")
+
+        additional_context.append(
+            make_markdown_section(
+                f"Available Skills ({zrb_cmd} Workflows)", "\n".join(zrb_context)
+            )
+        )
+
+        new_section = "\n\n".join(additional_context)
+        return next_handler(ctx, f"{current_prompt}\n\n{new_section}")
+
+    return zrb_prompt

zrb/llm/skill/__init__.py

@@ -0,0 +1,3 @@
+from zrb.llm.skill.manager import Skill, SkillManager, skill_manager
+
+__all__ = ["Skill", "SkillManager", "skill_manager"]

zrb/llm/skill/manager.py

@@ -0,0 +1,86 @@
+import os
+from typing import Dict, List, Optional
+
+IGNORE_DIRS = {
+    ".git",
+    "node_modules",
+    "__pycache__",
+    ".venv",
+    "venv",
+    "dist",
+    "build",
+    ".idea",
+    ".vscode",
+}
+
+
+class Skill:
+    def __init__(self, name: str, path: str, description: str):
+        self.name = name
+        self.path = path
+        self.description = description
+
+
+class SkillManager:
+    def __init__(self, root_dir: str = "."):
+        self.root_dir = root_dir
+        self._skills: Dict[str, Skill] = {}
+
+    def scan(self) -> List[Skill]:
+        self._skills = {}
+        for root, dirs, files in os.walk(self.root_dir):
+            dirs[:] = [d for d in dirs if d not in IGNORE_DIRS]
+
+            for file in files:
+                if file == "SKILL.md" or file.endswith(".skill.md"):
+                    full_path = os.path.join(root, file)
+                    rel_path = os.path.relpath(full_path, self.root_dir)
+                    self._load_skill(rel_path, full_path)
+
+        return list(self._skills.values())
+
+    def _load_skill(self, rel_path: str, full_path: str):
+        try:
+            with open(full_path, "r", encoding="utf-8") as f:
+                lines = f.readlines()
+
+            name = os.path.basename(os.path.dirname(full_path))
+            description = "No description"
+
+            # Parse Markdown for Header 1
+            for line in lines:
+                stripped = line.strip()
+                if stripped.startswith("# "):
+                    name = stripped[2:].strip()
+                    break
+
+            # Use name as key, handle duplicates
+            key = name
+            if key in self._skills:
+                key = f"{name} ({rel_path})"
+
+            self._skills[key] = Skill(name=key, path=rel_path, description=description)
+
+        except Exception:
+            pass
+
+    def get_skill_content(self, name: str) -> Optional[str]:
+        skill = self._skills.get(name)
+        if not skill:
+            # Try partial match or path match
+            for s in self._skills.values():
+                if s.name == name or s.path == name:
+                    skill = s
+                    break
+
+        if not skill:
+            return None
+
+        try:
+            with open(skill.path, "r", encoding="utf-8") as f:
+                return f.read()
+        except Exception as e:
+            return f"Error reading skill file: {e}"
+
+
+skill_manager = SkillManager()

zrb/llm/task/__init__.py

@@ -0,0 +1,4 @@
+from zrb.llm.task.llm_chat_task import LLMChatTask
+from zrb.llm.task.llm_task import LLMTask
+
+__all__ = ["LLMChatTask", "LLMTask"]

zrb/llm/task/llm_chat_task.py

@@ -0,0 +1,316 @@
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
+
+from zrb.attr.type import BoolAttr, StrAttr, fstring
+from zrb.config.config import CFG
+from zrb.context.any_context import AnyContext
+from zrb.context.print_fn import PrintFn
+from zrb.context.shared_context import SharedContext
+from zrb.env.any_env import AnyEnv
+from zrb.input.any_input import AnyInput
+from zrb.input.bool_input import BoolInput
+from zrb.input.str_input import StrInput
+from zrb.llm.app.confirmation.handler import ConfirmationMiddleware
+from zrb.llm.config.config import LLMConfig
+from zrb.llm.config.config import llm_config as default_llm_config
+from zrb.llm.config.limiter import LLMLimiter
+from zrb.llm.history_manager.any_history_manager import AnyHistoryManager
+from zrb.llm.history_manager.file_history_manager import FileHistoryManager
+from zrb.llm.prompt.compose import PromptManager
+from zrb.llm.task.llm_task import LLMTask
+from zrb.llm.util.attachment import get_attachments
+from zrb.session.session import Session
+from zrb.task.any_task import AnyTask
+from zrb.task.base_task import BaseTask
+from zrb.util.attr import get_attr, get_bool_attr, get_str_attr
+from zrb.util.string.name import get_random_name
+
+if TYPE_CHECKING:
+    from pydantic_ai import Tool, UserContent
+    from pydantic_ai._agent_graph import HistoryProcessor
+    from pydantic_ai.models import Model
+    from pydantic_ai.settings import ModelSettings
+    from pydantic_ai.tools import ToolFuncEither
+    from pydantic_ai.toolsets import AbstractToolset
+    from rich.theme import Theme
+
+
+class LLMChatTask(BaseTask):
+
+    def __init__(
+        self,
+        name: str,
+        color: int | None = None,
+        icon: str | None = None,
+        description: str | None = None,
+        cli_only: bool = False,
+        input: list[AnyInput | None] | AnyInput | None = None,
+        env: list[AnyEnv | None] | AnyEnv | None = None,
+        system_prompt: (
+            "Callable[[AnyContext], str | fstring | None] | str | None"
+        ) = None,
+        render_system_prompt: bool = False,
+        prompt_manager: PromptManager | None = None,
+        tools: list["Tool | ToolFuncEither"] = [],
+        toolsets: list["AbstractToolset[None]"] = [],
+        message: StrAttr | None = None,
+        render_message: bool = True,
+        attachment: "UserContent | list[UserContent] | Callable[[AnyContext], UserContent | list[UserContent]] | None" = None,  # noqa
+        history_processors: list["HistoryProcessor"] = [],
+        llm_config: LLMConfig | None = None,
+        llm_limitter: LLMLimiter | None = None,
+        model: (
+            "Callable[[AnyContext], Model | str | fstring | None] | Model | None"
+        ) = None,
+        render_model: bool = True,
+        model_settings: (
+            "ModelSettings | Callable[[AnyContext], ModelSettings] | None"
+        ) = None,
+        conversation_name: StrAttr | None = None,
+        render_conversation_name: bool = True,
+        history_manager: AnyHistoryManager | None = None,
+        tool_confirmation: Callable[[Any], Any] | None = None,
+        yolo: BoolAttr = False,
+        ui_summarize_commands: list[str] = [],
+        ui_attach_commands: list[str] = [],
+        ui_exit_commands: list[str] = [],
+        ui_info_commands: list[str] = [],
+        ui_save_commands: list[str] = [],
+        ui_load_commands: list[str] = [],
+        ui_redirect_output_commands: list[str] = [],
+        ui_yolo_toggle_commands: list[str] = [],
+        ui_exec_commands: list[str] = [],
+        ui_greeting: StrAttr | None = None,
+        render_ui_greeting: bool = True,
+        ui_assistant_name: StrAttr | None = None,
+        render_ui_assistant_name: bool = True,
+        ui_jargon: StrAttr | None = None,
+        render_ui_jargon: bool = True,
+        ui_ascii_art: StrAttr | None = None,
+        render_ui_ascii_art_name: bool = True,
+        triggers: list[Callable[[], Any]] = [],
+        confirmation_middlewares: list[ConfirmationMiddleware] = [],
+        markdown_theme: "Theme | None" = None,
+        interactive: BoolAttr = True,
+        execute_condition: bool | str | Callable[[AnyContext], bool] = True,
+        retries: int = 2,
+        retry_period: float = 0,
+        readiness_check: list[AnyTask] | AnyTask | None = None,
+        readiness_check_delay: float = 0.5,
+        readiness_check_period: float = 5,
+        readiness_failure_threshold: int = 1,
+        readiness_timeout: int = 60,
+        monitor_readiness: bool = False,
+        upstream: list[AnyTask] | AnyTask | None = None,
+        fallback: list[AnyTask] | AnyTask | None = None,
+        successor: list[AnyTask] | AnyTask | None = None,
+        print_fn: PrintFn | None = None,
+    ):
+        super().__init__(
+            name=name,
+            color=color,
+            icon=icon,
+            description=description,
+            cli_only=cli_only,
+            input=input,
+            env=env,
+            execute_condition=execute_condition,
+            retries=retries,
+            retry_period=retry_period,
+            readiness_check=readiness_check,
+            readiness_check_delay=readiness_check_delay,
+            readiness_check_period=readiness_check_period,
+            readiness_failure_threshold=readiness_failure_threshold,
+            readiness_timeout=readiness_timeout,
+            monitor_readiness=monitor_readiness,
+            upstream=upstream,
+            fallback=fallback,
+            successor=successor,
+            print_fn=print_fn,
+        )
+        self._llm_config = default_llm_config if llm_config is None else llm_config
+        self._llm_limitter = llm_limitter
+        self._system_prompt = system_prompt
+        self._render_system_prompt = render_system_prompt
+        self._prompt_manager = prompt_manager
+        self._tools = tools
+        self._toolsets = toolsets
+        self._message = message
+        self._render_message = render_message
+        self._attachment = attachment
+        self._history_processors = history_processors
+        self._model = model
+        self._render_model = render_model
+        self._model_settings = model_settings
+        self._conversation_name = conversation_name
+        self._render_conversation_name = render_conversation_name
+        self._history_manager = (
+            FileHistoryManager(history_dir=CFG.LLM_HISTORY_DIR)
+            if history_manager is None
+            else history_manager
+        )
+        self._tool_confirmation = tool_confirmation
+        self._yolo = yolo
+        self._ui_summarize_commands = ui_summarize_commands
+        self._ui_attach_commands = ui_attach_commands
+        self._ui_exit_commands = ui_exit_commands
+        self._ui_info_commands = ui_info_commands
+        self._ui_save_commands = ui_save_commands
+        self._ui_load_commands = ui_load_commands
+        self._ui_redirect_output_commands = ui_redirect_output_commands
+        self._ui_yolo_toggle_commands = ui_yolo_toggle_commands
+        self._ui_exec_commands = ui_exec_commands
+        self._ui_greeting = ui_greeting
+        self._render_ui_greeting = render_ui_greeting
+        self._ui_assistant_name = ui_assistant_name
+        self._render_ui_assistant_name = render_ui_assistant_name
+        self._ui_jargon = ui_jargon
+        self._render_ui_jargon = render_ui_jargon
+        self._ui_ascii_art_name = ui_ascii_art
+        self._render_ui_ascii_art_name = render_ui_ascii_art_name
+        self._triggers = triggers
+        self._confirmation_middlewares = confirmation_middlewares
+        self._markdown_theme = markdown_theme
+        self._interactive = interactive
+
+    @property
+    def prompt_manager(self) -> PromptManager:
+        if self._prompt_manager is None:
+            raise ValueError(f"Task {self.name} doesn't have prompt_manager")
+        return self._prompt_manager
+
+    def add_toolset(self, *toolset: "AbstractToolset"):
+        self.append_toolset(*toolset)
+
+    def append_toolset(self, *toolset: "AbstractToolset"):
+        self._toolsets += list(toolset)
+
+    def add_tool(self, *tool: "Tool | ToolFuncEither"):
+        self.append_tool(*tool)
+
+    def append_tool(self, *tool: "Tool | ToolFuncEither"):
+        self._tools += list(tool)
+
+    def add_history_processor(self, *processor: "HistoryProcessor"):
+        self.append_history_processor(*processor)
+
+    def append_history_processor(self, *processor: "HistoryProcessor"):
+        self._history_processors += list(processor)
+
+    def add_confirmation_middleware(self, *middleware: ConfirmationMiddleware):
+        self.prepend_confirmation_middleware(*middleware)
+
+    def prepend_confirmation_middleware(self, *middleware: ConfirmationMiddleware):
+        self._confirmation_middlewares = (
+            list(middleware) + self._confirmation_middlewares
+        )
+
+    def add_trigger(self, *trigger: Callable[[], Any]):
+        self.append_trigger(*trigger)
+
+    def append_trigger(self, *trigger: Callable[[], Any]):
+        self._triggers += trigger
+
+    async def _exec_action(self, ctx: AnyContext) -> Any:
+        from zrb.llm.app.lexer import CLIStyleLexer
+        from zrb.llm.app.ui import UI
+
+        initial_conversation_name = self._get_conversation_name(ctx)
+        initial_yolo = get_bool_attr(ctx, self._yolo, False)
+        initial_message = get_attr(ctx, self._message, "", self._render_message)
+        initial_attachments = get_attachments(ctx, self._attachment)
+        ui_greeting = get_str_attr(ctx, self._ui_greeting, "", self._render_ui_greeting)
+        ui_assistant_name = get_str_attr(
+            ctx, self._ui_assistant_name, "", self._render_ui_assistant_name
+        )
+        ui_jargon = get_str_attr(ctx, self._ui_jargon, "", self._render_ui_jargon)
+        ascii_art = get_str_attr(
+            ctx, self._ui_ascii_art_name, "", self._render_ui_ascii_art_name
+        )
+        interactive = get_bool_attr(ctx, self._interactive, True)
+
+        llm_task_core = LLMTask(
+            name=f"{self.name}-process",
+            input=[
+                StrInput("message", "Message"),
+                StrInput("session", "Conversation Session"),
+                BoolInput("yolo", "YOLO Mode"),
+                StrInput("attachments", "Attachments"),
+            ],
+            env=self.envs,
+            system_prompt=self._system_prompt,
+            render_system_prompt=self._render_system_prompt,
+            prompt_manager=self._prompt_manager,
+            tools=self._tools,
+            toolsets=self._toolsets,
+            history_processors=self._history_processors,
+            llm_config=self._llm_config,
+            llm_limitter=self._llm_limitter,
+            model=self._model,
+            render_model=self._render_model,
+            model_settings=self._model_settings,
+            history_manager=self._history_manager,
+            message="{ctx.input.message}",
+            conversation_name="{ctx.input.session}",
+            yolo="{ctx.input.yolo}",
+            attachment=lambda ctx: ctx.input.attachments,
+            summarize_command=self._ui_summarize_commands,
+        )
+
+        if not interactive:
+            session_input = {
+                "message": initial_message,
+                "session": initial_conversation_name,
+                "yolo": initial_yolo,
+                "attachments": initial_attachments,
+            }
+            shared_ctx = SharedContext(
+                input=session_input,
+                print_fn=ctx.shared_print,  # Use current task's print function
+            )
+            session = Session(shared_ctx)
+            return await llm_task_core.async_run(session)
+
+        ui = UI(
+            greeting=ui_greeting,
+            assistant_name=ui_assistant_name,
+            ascii_art=ascii_art,
+            jargon=ui_jargon,
+            output_lexer=CLIStyleLexer(),
+            llm_task=llm_task_core,
+            history_manager=self._history_manager,
+            initial_message=initial_message,
+            initial_attachments=initial_attachments,
+            conversation_session_name=initial_conversation_name,
+            yolo=initial_yolo,
+            triggers=self._triggers,
+            confirmation_middlewares=self._confirmation_middlewares,
+            markdown_theme=self._markdown_theme,
+            summarize_commands=self._ui_summarize_commands,
+            attach_commands=self._ui_attach_commands,
+            exit_commands=self._ui_exit_commands,
+            info_commands=self._ui_info_commands,
+            save_commands=self._ui_save_commands,
+            load_commands=self._ui_load_commands,
+            yolo_toggle_commands=self._ui_yolo_toggle_commands,
+            redirect_output_commands=self._ui_redirect_output_commands,
+            exec_commands=self._ui_exec_commands,
+            model=self._get_model(ctx),
+        )
+        await ui.run_async()
+        return ui.last_output
+
+    def _get_conversation_name(self, ctx: AnyContext) -> str:
+        conversation_name = str(
+            get_attr(ctx, self._conversation_name, "", self._render_conversation_name)
+        )
+        if conversation_name.strip() == "":
+            conversation_name = get_random_name()
+        return conversation_name
+
+    def _get_model(self, ctx: AnyContext) -> "str | Model":
+        model = self._model
+        rendered_model = get_attr(ctx, model, None, auto_render=self._render_model)
+        if rendered_model is not None:
+            return rendered_model
+        return self._llm_config.model