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

zrb/llm/task/llm_task.py

@@ -0,0 +1,245 @@
+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.env.any_env import AnyEnv
+from zrb.input.any_input import AnyInput
+from zrb.llm.agent.agent import create_agent, run_agent
+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.config.limiter import llm_limiter as default_llm_limitter
+from zrb.llm.history_manager.any_history_manager import AnyHistoryManager
+from zrb.llm.history_manager.file_history_manager import FileHistoryManager
+from zrb.llm.history_processor.summarizer import (
+    summarize_history,
+)
+from zrb.llm.prompt.compose import PromptManager
+from zrb.llm.util.attachment import get_attachments
+from zrb.llm.util.stream_response import (
+    create_event_handler,
+    create_faint_printer,
+)
+from zrb.task.any_task import AnyTask
+from zrb.task.base_task import BaseTask
+from zrb.util.attr import get_attr, get_bool_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
+
+
+class LLMTask(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,
+        summarize_command: list[str] = [],
+        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 = (
+            default_llm_limitter if llm_limitter is None else 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._summarize_command = summarize_command
+
+    @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)
+
+    async def _exec_action(self, ctx: AnyContext) -> Any:
+        conversation_name = self._get_conversation_name(ctx)
+        message_history = self._history_manager.load(conversation_name)
+        user_message = get_attr(ctx, self._message, "", self._render_message)
+        user_attachments = get_attachments(ctx, self._attachment)
+
+        if (
+            isinstance(user_message, str)
+            and user_message.strip() in self._summarize_command
+        ):
+            ctx.print("Compressing conversation history...", plain=True)
+            new_history = await summarize_history(message_history)
+            self._history_manager.update(conversation_name, new_history)
+            self._history_manager.save(conversation_name)
+            return "Conversation history compressed."
+
+        yolo = get_bool_attr(ctx, self._yolo, False)
+        system_prompt = self._get_system_prompt(ctx)
+        ctx.log_debug(f"SYSTEM PROMPT: {system_prompt}")
+        agent = create_agent(
+            model=self._get_model(ctx),
+            system_prompt=self._get_system_prompt(ctx),
+            tools=self._tools,
+            toolsets=self._toolsets,
+            model_settings=self._get_model_settings(ctx),
+            history_processors=self._history_processors,
+            yolo=yolo,
+        )
+
+        print_event = create_faint_printer(ctx)
+        handle_event = create_event_handler(
+            print_event,
+            show_tool_call_detail=CFG.LLM_SHOW_TOOL_CALL_PREPARATION,
+            show_tool_result=CFG.LLM_SHOW_TOOL_CALL_RESULT,
+        )
+
+        output, new_history = await run_agent(
+            agent=agent,
+            message=user_message,
+            message_history=message_history,
+            limiter=self._llm_limitter,
+            attachments=user_attachments,
+            print_fn=ctx.print,
+            event_handler=handle_event,
+            tool_confirmation=self._tool_confirmation,
+        )
+
+        self._history_manager.update(conversation_name, new_history)
+        self._history_manager.save(conversation_name)
+        ctx.log_debug(f"All messages: {new_history}")
+
+        return output
+
+    def _get_system_prompt(self, ctx: AnyContext) -> str:
+        if self._prompt_manager is None:
+            return str(
+                get_attr(ctx, self._system_prompt, "", self._render_system_prompt)
+            )
+        compose_prompt = self._prompt_manager.compose_prompt()
+        return compose_prompt(ctx)
+
+    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_settings(self, ctx: AnyContext) -> "ModelSettings | None":
+        model_settings = self._model_settings
+        rendered_model_settings = get_attr(ctx, model_settings, None)
+        if rendered_model_settings is not None:
+            return rendered_model_settings
+        return self._llm_config.model_settings
+
+    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

zrb/llm/tool/__init__.py

@@ -0,0 +1,39 @@
+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.rag import create_rag_from_directory
+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
+
+__all__ = [
+    "run_shell_command",
+    "analyze_code",
+    "list_files",
+    "read_file",
+    "read_files",
+    "write_file",
+    "write_files",
+    "replace_in_file",
+    "search_files",
+    "analyze_file",
+    "create_note_tools",
+    "create_rag_from_directory",
+    "create_activate_skill_tool",
+    "create_sub_agent_tool",
+    "open_web_page",
+    "search_internet",
+    "create_list_zrb_task_tool",
+    "create_run_zrb_task_tool",
+]

zrb/llm/tool/bash.py

@@ -0,0 +1,75 @@
+import asyncio
+import re
+
+from zrb.util.cli.style import stylize_faint
+
+
+async def run_shell_command(command: str, timeout: int = 30) -> str:
+    """
+    Executes a shell command on the host system and returns its combined stdout and stderr.
+    This is a powerful tool for running builds, tests, or system utilities.
+
+    **CRITICAL SAFETY:**
+    - DO NOT run destructive commands (e.g., `rm -rf /`) without absolute certainty.
+    - Prefer specialized tools (like `read_file` or `write_file`) for file operations.
+
+    **USAGE GUIDELINES:**
+    - Use non-interactive commands.
+    - If a command is expected to produce massive output, use `timeout` or pipe to a file.
+    - The output is streamed to the console in real-time.
+
+    Args:
+        command (str): The full shell command to execute.
+        timeout (int): Maximum wait time in seconds before terminating the process. Defaults to 30.
+    """
+    ANSI_ESCAPE = re.compile(
+        r"(?:\x1B\[[0-?]*[ -/]*[@-~])|"  # CSI (Control Sequence Introducer)
+        r"(?:\x1B\][^\a\x1b]*[\a\x1b])|"  # OSC (Operating System Command)
+        r"(?:\x1B[0-9=>])"  # Simple 2-byte (DECSC, DECRC, etc.)
+    )
+    try:
+        process = await asyncio.create_subprocess_shell(
+            command, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE
+        )
+
+        stdout_lines = []
+        stderr_lines = []
+
+        async def read_stream(stream, lines_list, prefix=""):
+            while True:
+                line = await stream.readline()
+                if not line:
+                    break
+                decoded = line.decode()
+                if decoded:
+                    shown = ANSI_ESCAPE.sub("", decoded)
+                    shown = stylize_faint(shown)
+                    print(f"{prefix}  {shown}", end="")  # Stream to console
+                    lines_list.append(decoded)
+
+        # Wait for the process to complete or timeout
+        try:
+            await asyncio.wait_for(
+                asyncio.gather(
+                    read_stream(process.stdout, stdout_lines, ""),
+                    read_stream(process.stderr, stderr_lines, "[stderr] "),
+                    process.wait(),
+                ),
+                timeout=timeout,
+            )
+        except asyncio.TimeoutError:
+            if process.returncode is None:
+                try:
+                    process.terminate()
+                    await process.wait()
+                except ProcessLookupError:
+                    pass
+            return f"Error: Command timed out after {timeout} seconds."
+
+        output = "\n".join(stdout_lines)
+        error = "\n".join(stderr_lines)
+
+        return f"Exit Code: {process.returncode}\nStdout:\n{output}\nStderr:\n{error}"
+
+    except Exception as e:
+        return f"Error executing command: {e}"

zrb/llm/tool/code.py

@@ -0,0 +1,266 @@
+import fnmatch
+import json
+import os
+
+from zrb.config.config import CFG
+from zrb.llm.agent.agent import create_agent, run_agent
+from zrb.llm.config.config import llm_config
+from zrb.llm.config.limiter import llm_limiter
+from zrb.llm.prompt.default import (
+    get_repo_extractor_system_prompt,
+    get_repo_summarizer_system_prompt,
+)
+from zrb.llm.tool.file import DEFAULT_EXCLUDED_PATTERNS
+
+_DEFAULT_EXTENSIONS = [
+    "py",
+    "go",
+    "java",
+    "ts",
+    "js",
+    "rs",
+    "rb",
+    "php",
+    "sh",
+    "bash",
+    "c",
+    "cpp",
+    "h",
+    "hpp",
+    "cs",
+    "swift",
+    "kt",
+    "scala",
+    "m",
+    "pl",
+    "lua",
+    "sql",
+    "html",
+    "css",
+    "scss",
+    "less",
+    "json",
+    "yaml",
+    "yml",
+    "toml",
+    "ini",
+    "xml",
+    "md",
+    "rst",
+    "txt",
+]
+
+
+async def analyze_code(
+    path: str,
+    query: str,
+    extensions: list[str] | None = None,
+    exclude_patterns: list[str] | None = None,
+) -> str:
+    """
+    Performs a deep, semantic analysis of an entire codebase or directory to answer complex architectural or logic-based questions.
+    This tool uses a 'map-reduce' strategy to handle large repositories that exceed single-prompt limits.
+
+    **WHEN TO USE:**
+    - To understand system-wide flows, architectural patterns, or cross-file dependencies.
+    - When you need a summary of how a feature is implemented across multiple files.
+    - To identify potential refactoring opportunities or technical debt.
+
+    **LIMITATIONS:**
+    - It extracts and summarizes information; it does not read every single byte if the repo is massive.
+    - For precise line-by-line reading of a known file, use `read_file` instead.
+
+    Args:
+        path (str): Path to the directory or repository.
+        query (str): A clear, specific question or analysis goal (e.g., "How is authentication handled?").
+        extensions (list[str], optional): File extensions to include (e.g., ["py", "ts"]).
+        exclude_patterns (list[str], optional): Glob patterns to ignore (e.g., ["tests/*"]).
+
+    Returns:
+        str: A comprehensive analytical report.
+    """
+    if extensions is None:
+        extensions = _DEFAULT_EXTENSIONS
+    if exclude_patterns is None:
+        exclude_patterns = DEFAULT_EXCLUDED_PATTERNS
+
+    abs_path = os.path.abspath(os.path.expanduser(path))
+    if not os.path.exists(abs_path):
+        return f"Error: Path not found: {path}"
+
+    # 1. Gather files
+    file_metadatas = _get_file_metadatas(abs_path, extensions, exclude_patterns)
+    if not file_metadatas:
+        return "No files found matching the criteria."
+
+    print(f"  📝 Extraction ({len(file_metadatas)} files)")
+
+    # 2. Extract Info
+    extraction_token_threshold = CFG.LLM_REPO_ANALYSIS_EXTRACTION_TOKEN_THRESHOLD
+    extracted_infos = await _extract_info(
+        file_metadatas=file_metadatas,
+        query=query,
+        token_limit=extraction_token_threshold,
+    )
+
+    if not extracted_infos:
+        return "No information could be extracted from the files."
+
+    if len(extracted_infos) == 1:
+        return extracted_infos[0]
+
+    # 3. Summarize Info (Reduce)
+    summarization_token_threshold = CFG.LLM_REPO_ANALYSIS_SUMMARIZATION_TOKEN_THRESHOLD
+    summarized_infos = extracted_infos
+
+    while len(summarized_infos) > 1:
+        print(f"  📝 Summarization ({len(summarized_infos)} chunks)")
+        summarized_infos = await _summarize_info(
+            extracted_infos=summarized_infos,
+            query=query,
+            token_limit=summarization_token_threshold,
+        )
+
+    return summarized_infos[0]
+
+
+def _get_file_metadatas(
+    dir_path: str,
+    extensions: list[str],
+    exclude_patterns: list[str],
+) -> list[dict[str, str]]:
+    metadata_list = []
+    for root, _, files in os.walk(dir_path):
+        files.sort()
+        for file in files:
+            if not any(file.endswith(f".{ext}") for ext in extensions):
+                continue
+            file_path = os.path.join(root, file)
+            try:
+                rel_path = os.path.relpath(file_path, dir_path)
+                if _is_excluded(rel_path, exclude_patterns):
+                    continue
+                with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+                    metadata_list.append({"path": rel_path, "content": f.read()})
+            except Exception as e:
+                print(f"Error reading file {file_path}: {e}")
+    metadata_list.sort(key=lambda m: m["path"])
+    return metadata_list
+
+
+def _is_excluded(name: str, patterns: list[str]) -> bool:
+    for pattern in patterns:
+        if fnmatch.fnmatch(name, pattern):
+            return True
+        parts = name.split(os.path.sep)
+        for part in parts:
+            if fnmatch.fnmatch(part, pattern):
+                return True
+    return False
+
+
+async def _extract_info(
+    file_metadatas: list[dict[str, str]],
+    query: str,
+    token_limit: int,
+) -> list[str]:
+    agent = create_agent(
+        model=llm_config.model,
+        system_prompt=get_repo_extractor_system_prompt(),
+    )
+
+    extracted_infos = []
+    content_buffer = []
+    current_token_count = 0
+
+    # We estimate token count of the prompt template overhead
+    base_overhead = 100
+
+    for metadata in file_metadatas:
+        path = metadata.get("path", "")
+        content = metadata.get("content", "")
+        file_obj = {"path": path, "content": content}
+        file_str = json.dumps(file_obj)
+        file_tokens = llm_limiter.count_tokens(file_str)
+
+        if current_token_count + file_tokens + base_overhead > token_limit:
+            if content_buffer:
+                await _run_extraction(agent, query, content_buffer, extracted_infos)
+
+            content_buffer = [file_obj]
+            current_token_count = file_tokens
+        else:
+            content_buffer.append(file_obj)
+            current_token_count += file_tokens
+
+    # Process remaining buffer
+    if content_buffer:
+        await _run_extraction(agent, query, content_buffer, extracted_infos)
+
+    return extracted_infos
+
+
+async def _run_extraction(agent, query, content_buffer, extracted_infos):
+    prompt_data = {
+        "main_assistant_query": query,
+        "files": content_buffer,
+    }
+    # We serialize to JSON for the prompt
+    message = json.dumps(prompt_data)
+
+    result, _ = await run_agent(
+        agent=agent,
+        message=message,
+        message_history=[],  # Stateless
+        limiter=llm_limiter,
+    )
+    extracted_infos.append(str(result))
+
+
+async def _summarize_info(
+    extracted_infos: list[str],
+    query: str,
+    token_limit: int,
+) -> list[str]:
+    agent = create_agent(
+        model=llm_config.model,
+        system_prompt=get_repo_summarizer_system_prompt(),
+    )
+
+    summarized_infos = []
+    content_buffer = ""
+    # Overhead for prompt structure
+    base_overhead = 100
+
+    for info in extracted_infos:
+        # Check if adding this info exceeds limit
+        if (
+            llm_limiter.count_tokens(content_buffer + info) + base_overhead
+            > token_limit
+        ):
+            if content_buffer:
+                await _run_summarization(agent, query, content_buffer, summarized_infos)
+            content_buffer = info
+        else:
+            content_buffer += info + "\n"
+
+    if content_buffer:
+        await _run_summarization(agent, query, content_buffer, summarized_infos)
+
+    return summarized_infos
+
+
+async def _run_summarization(agent, query, content_buffer, summarized_infos):
+    prompt_data = {
+        "main_assistant_query": query,
+        "extracted_info": content_buffer,
+    }
+    message = json.dumps(prompt_data)
+
+    result, _ = await run_agent(
+        agent=agent,
+        message=message,
+        message_history=[],  # Stateless
+        limiter=llm_limiter,
+    )
+    summarized_infos.append(str(result))