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

zrb/llm/note/manager.py

@@ -0,0 +1,122 @@
+import json
+import os
+from typing import Dict
+
+from zrb.config.config import CFG
+
+
+class NoteManager:
+    def __init__(self, context_file: str = CFG.LLM_NOTE_FILE):
+        self.context_file = os.path.abspath(os.path.expanduser(context_file))
+
+    def _get_normalized_path(self, path: str) -> str:
+        """
+        Convert path to absolute path, then try to make it relative to home.
+        If it's inside home, return '~/rel/path'.
+        Otherwise return absolute path.
+        """
+        abs_path = os.path.abspath(os.path.expanduser(path))
+        home = os.path.expanduser("~")
+
+        if abs_path.startswith(home):
+            rel_path = os.path.relpath(abs_path, home)
+            if rel_path == ".":
+                return "~"
+            return os.path.join("~", rel_path)
+
+        return abs_path
+
+    def _load_data(self) -> Dict[str, str]:
+        if not os.path.exists(self.context_file):
+            return {}
+        try:
+            with open(self.context_file, "r", encoding="utf-8") as f:
+                return json.load(f)
+        except (json.JSONDecodeError, OSError):
+            return {}
+
+    def _save_data(self, data: Dict[str, str]):
+        os.makedirs(os.path.dirname(self.context_file), exist_ok=True)
+        with open(self.context_file, "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=2, sort_keys=True)
+
+    def write(self, context_path: str, content: str):
+        """
+        Write content to context file for a specific path.
+        Turn context path to be relative to home directory (unless it is outside home directory).
+        """
+        key = self._get_normalized_path(context_path)
+        data = self._load_data()
+        data[key] = content
+        self._save_data(data)
+
+    def read(self, context_path: str) -> str:
+        """
+        Seek for any key in context file that match context path.
+        """
+        key = self._get_normalized_path(context_path)
+        data = self._load_data()
+        return data.get(key, "")
+
+    def read_all(self, context_path: str) -> Dict[str, str]:
+        """
+        Seek for all key in context file that match context path or any of its parent, return as key-value.
+        """
+        target_path = self._get_normalized_path(context_path)
+        data = self._load_data()
+
+        result = {}
+
+        # Check for root/home notes (~)
+        if "~" in data:
+            result["~"] = data["~"]
+
+        # If target_path is just "~", we are done
+        if target_path == "~":
+            return result
+
+        # We need to find all parents of target_path that are in data
+        # Example target: ~/zrb/src
+        # Parents: ~, ~/zrb
+
+        # Split the path parts
+        # If path starts with "~", we treat it specially
+
+        parts = []
+        if target_path.startswith("~" + os.sep) or target_path == "~":
+            # Remove ~
+            rel_part = target_path[2:]  # "zrb/src"
+            parts = rel_part.split(os.sep)
+            current = "~"
+        else:
+            # Absolute path like /opt/project
+            parts = target_path.split(os.sep)
+            # parts for /opt/project -> ['', 'opt', 'project']
+            current = parts[0]  # '' (root)
+            if current == "":
+                current = "/"
+            parts = parts[1:]
+
+        # Iterate and build up path
+        for part in parts:
+            if current == "~":
+                current = os.path.join("~", part)
+            elif current == "/":
+                current = os.path.join("/", part)
+            else:
+                current = os.path.join(current, part)
+
+            # Normalize just in case os.path.join changes things
+            # But wait, our keys in JSON are stored with separators.
+            # We should probably reconstruct the key carefully.
+
+            # Let's simplify:
+            # Check if 'current' exists in data
+            if current in data:
+                result[current] = data[current]
+
+        # Finally check the target_path itself if not covered
+        if target_path in data and target_path not in result:
+            result[target_path] = data[target_path]
+
+        return result

zrb/llm/prompt/__init__.py

@@ -0,0 +1,29 @@
+from zrb.llm.prompt.claude_compatibility import create_claude_compatibility_prompt
+from zrb.llm.prompt.compose import PromptManager, PromptMiddleware, new_prompt
+from zrb.llm.prompt.default import (
+    get_file_extractor_system_prompt,
+    get_mandate_prompt,
+    get_persona_prompt,
+    get_repo_extractor_system_prompt,
+    get_repo_summarizer_system_prompt,
+    get_summarizer_system_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
+
+__all__ = [
+    "create_claude_compatibility_prompt",
+    "PromptManager",
+    "PromptMiddleware",
+    "new_prompt",
+    "get_file_extractor_system_prompt",
+    "get_mandate_prompt",
+    "get_persona_prompt",
+    "get_repo_extractor_system_prompt",
+    "get_repo_summarizer_system_prompt",
+    "get_summarizer_system_prompt",
+    "create_note_prompt",
+    "system_context",
+    "create_zrb_prompt",
+]

zrb/llm/prompt/claude_compatibility.py

@@ -0,0 +1,92 @@
+from pathlib import Path
+from typing import Callable, List, Optional
+
+from zrb.context.any_context import AnyContext
+from zrb.llm.skill.manager import SkillManager
+from zrb.util.markdown import make_markdown_section
+
+
+def create_claude_compatibility_prompt(skill_manager: SkillManager):
+    def claude_compatibility(
+        ctx: AnyContext,
+        current_prompt: str,
+        next_handler: Callable[[AnyContext, str], str],
+    ) -> str:
+        search_dirs = _get_search_directories()
+        additional_context = []
+
+        # 1. CLAUDE.md
+        claude_content = _get_combined_content("CLAUDE.md", search_dirs)
+        if claude_content:
+            additional_context.append(
+                make_markdown_section(
+                    "Project Instructions (CLAUDE.md)", claude_content
+                )
+            )
+
+        # 2. AGENTS.md
+        agents_content = _get_combined_content("AGENTS.md", search_dirs)
+        if agents_content:
+            additional_context.append(
+                make_markdown_section("Agent Definitions (AGENTS.md)", agents_content)
+            )
+
+        # 3. Available Claude Skills
+        skills_section = _get_skills_section(skill_manager)
+        if skills_section:
+            additional_context.append(skills_section)
+
+        new_section = "\n\n".join(additional_context)
+        return next_handler(ctx, f"{current_prompt}\n\n{new_section}")
+
+    return claude_compatibility
+
+
+def _get_search_directories() -> List[Path]:
+    search_dirs: List[Path] = []
+    # 1. User global config (~/.claude)
+    try:
+        home = Path.home()
+        search_dirs.append(home / ".claude")
+    except Exception:
+        pass
+
+    # 2. Project directories (Root -> ... -> CWD)
+    try:
+        cwd = Path.cwd()
+        # Parents returns [parent, grandparent...]. We want reversed (Root first)
+        # This allows specific configs (closer to CWD) to override general ones
+        project_dirs = list(cwd.parents)[::-1] + [cwd]
+        search_dirs.extend(project_dirs)
+    except Exception:
+        pass
+    return search_dirs
+
+
+def _get_combined_content(filename: str, search_dirs: List[Path]) -> str:
+    contents = []
+    for directory in search_dirs:
+        file_path = directory / filename
+        if file_path.exists() and file_path.is_file():
+            try:
+                with open(file_path, "r", encoding="utf-8") as f:
+                    content = f.read().strip()
+                    if content:
+                        contents.append(content)
+            except Exception:
+                pass
+    return "\n\n".join(contents)
+
+
+def _get_skills_section(skill_manager: SkillManager) -> Optional[str]:
+    skills = skill_manager.scan()
+    if not skills:
+        return None
+
+    skills_context = ["Use 'activate_skill' to load instructions for a skill."]
+    for skill in skills:
+        skills_context.append(f"- {skill.name}")
+
+    return make_markdown_section(
+        "Available Skills (Claude Skills)", "\n".join(skills_context)
+    )

zrb/llm/prompt/compose.py

@@ -0,0 +1,55 @@
+from typing import Callable
+
+from zrb.context.any_context import AnyContext
+from zrb.util.attr import get_str_attr
+
+PromptMiddleware = Callable[[AnyContext, str, Callable[[AnyContext, str], str]], str]
+
+
+class PromptManager:
+    def __init__(self, middlewares: list[PromptMiddleware] = []):
+        self._middlewares = middlewares
+
+    def add_middleware(self, *middleware: PromptMiddleware):
+        self.append_middlewear(*middleware)
+
+    def append_middlewear(self, *middleware: PromptMiddleware):
+        self._middlewares += middleware
+
+    def compose_prompt(self) -> Callable[[AnyContext], str]:
+        """
+        Composes a list of prompt middlewares into a single prompt factory function.
+
+        Each middleware should have the signature:
+        (ctx: AnyContext, prompt: str, next: Callable[[AnyContext, str], str]) -> str
+
+        The resulting function takes an AnyContext and returns the final prompt string.
+        """
+
+        def composed_prompt_factory(ctx: AnyContext) -> str:
+            def dispatch(index: int, current_prompt: str) -> str:
+                if index >= len(self._middlewares):
+                    return current_prompt
+
+                middleware = self._middlewares[index]
+
+                def next_handler(c: AnyContext, p: str) -> str:
+                    return dispatch(index + 1, p)
+
+                return middleware(ctx, current_prompt, next_handler)
+
+            return dispatch(0, "")
+
+        return composed_prompt_factory
+
+
+def new_prompt(new_prompt: str, render: bool = False):
+    def new_prompt_middleware(
+        ctx: AnyContext, current_prompt: str, next: Callable[[AnyContext, str], str]
+    ):
+        effective_new_prompt = new_prompt
+        if render:
+            effective_new_prompt = get_str_attr(ctx, new_prompt, auto_render=True)
+        return next(ctx, f"{current_prompt}\n{effective_new_prompt}")
+
+    return new_prompt_middleware

zrb/llm/prompt/default.py

@@ -0,0 +1,51 @@
+import os
+
+from zrb.config.config import CFG
+
+
+def get_default_prompt(name: str) -> str:
+    # 1. Check for local project override (configured via LLM_PROMPT_DIR)
+    prompt_dir = getattr(CFG, "LLM_PROMPT_DIR", ".zrb/llm/prompt")
+    local_prompt_path = os.path.abspath(
+        os.path.join(os.getcwd(), prompt_dir, f"{name}.md")
+    )
+
+    if os.path.exists(local_prompt_path):
+        try:
+            with open(local_prompt_path, "r", encoding="utf-8") as f:
+                return f.read()
+        except Exception:
+            pass
+
+    # 2. Fallback to package default
+    cwd = os.path.dirname(__file__)
+    with open(os.path.join(cwd, "markdown", f"{name}.md"), "r", encoding="utf-8") as f:
+        return f.read()
+
+
+def get_persona_prompt(assistant_name: str | None = None) -> str:
+    effective_assistant_name = (
+        assistant_name if assistant_name else CFG.LLM_ASSISTANT_NAME
+    )
+    prompt = get_default_prompt("persona")
+    return prompt.replace("{ASSISTANT_NAME}", effective_assistant_name)
+
+
+def get_mandate_prompt() -> str:
+    return get_default_prompt("mandate")
+
+
+def get_summarizer_system_prompt() -> str:
+    return get_default_prompt("summarizer")
+
+
+def get_file_extractor_system_prompt() -> str:
+    return get_default_prompt("file_extractor")
+
+
+def get_repo_extractor_system_prompt() -> str:
+    return get_default_prompt("repo_extractor")
+
+
+def get_repo_summarizer_system_prompt() -> str:
+    return get_default_prompt("repo_summarizer")

zrb/llm/prompt/markdown/mandate.md

@@ -0,0 +1,23 @@
+# Mandate
+
+This mandate represents your core operating principles. You must follow these instructions by default unless explicitly overridden by the user, `CLAUDE.md`, or `AGENTS.md`.
+
+## Core Philosophy
+
+*   **Safety & Precision**: You are an expert engineer. Do not guess. Verify assumptions. Read code before editing.
+*   **Conventions**: Respect the existing coding style, naming conventions, and architectural patterns of the project.
+*   **Minimalism**: Do not introduce unnecessary dependencies or complexities. Use existing tools and libraries whenever possible.
+
+## Operational Rules
+
+1.  **Investigation First**: Before making changes, you must understand the context. Use `list_files`, `read_file`, or `search_files` to gather information.
+2.  **Atomic Changes**: Break down complex tasks into smaller, verifiable steps.
+3.  **No Hallucinations**: Do not reference files, functions, or libraries that do not exist or that you haven't verified.
+4.  **Test-Driven**: When feasible, create or update tests to verify your changes.
+5.  **Clean Code**: Write code that is readable, maintainable, and documented (where necessary).
+6.  **Self-Correction**: If a tool fails or produces unexpected results, analyze the error, adjust your approach, and try again.
+
+## Communication
+
+*   **Direct & Professional**: Be concise. Avoid fluff.
+*   **Transparent**: Briefly explain *why* you are doing something if it involves significant changes or risks.

zrb/llm/prompt/markdown/persona.md

@@ -0,0 +1,3 @@
+You are {ASSISTANT_NAME}, an intelligent, capable, and efficient AI Assistant.
+Your goal is to assist the user in their tasks with precision and speed.
+

zrb/llm/prompt/markdown/summarizer.md

@@ -0,0 +1,21 @@
+You are an expert conversation summarizer. Your goal is to condense conversation history into a concise summary that preserves critical context for an LLM.
+
+# Guidelines:
+1. **Preserve User Intent:** Clearly state what the user wanted to achieve.
+2. **Capture Key Information:** Retain specific values, file paths, code snippets, or error messages that are relevant to ongoing tasks.
+3. **Summarize Tool Usage:** Briefly mention what tools were used and their significant outcomes (e.g., "User listed files and found 'main.py'").
+4. **Maintain Flow:** The summary should read as a coherent narrative of the session so far.
+5. **Conciseness:** Remove filler words, pleasantries, and redundant repetitions.
+
+# Example:
+
+*Original:*
+User: "Hi, can you help me find the bug in my code?"
+AI: "Sure, please share the code."
+User: "Here it is: `def add(a, b): return a - b`"
+AI: "I see, it subtracts instead of adds."
+User: "Oh, right. Fix it please."
+AI: "I'll fix it." (Tool Call: replace_in_file) (Tool Result: Success)
+
+*Summary:*
+The user asked for help finding a bug in a provided `add` function. The AI identified that it performed subtraction instead of addition. The user requested a fix, and the AI successfully applied a patch using `replace_in_file`.

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"]