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

zrb/llm/agent/agent.py

@@ -0,0 +1,215 @@
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Any, Callable
+
+from zrb.llm.config.config import llm_config as default_llm_config
+from zrb.llm.config.limiter import LLMLimiter
+from zrb.llm.util.attachment import normalize_attachments
+from zrb.llm.util.prompt import expand_prompt
+
+# Context variable to propagate tool confirmation callback to sub-agents
+tool_confirmation_var: ContextVar[Callable[[Any], Any] | None] = ContextVar(
+    "tool_confirmation", default=None
+)
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent, DeferredToolRequests, DeferredToolResults, Tool
+    from pydantic_ai._agent_graph import HistoryProcessor
+    from pydantic_ai.messages import UserPromptPart
+    from pydantic_ai.models import Model
+    from pydantic_ai.output import OutputDataT, OutputSpec
+    from pydantic_ai.settings import ModelSettings
+    from pydantic_ai.tools import ToolFuncEither
+    from pydantic_ai.toolsets import AbstractToolset
+
+
+def create_agent(
+    model: "Model | str | None" = None,
+    system_prompt: str = "",
+    tools: list["Tool | ToolFuncEither"] = [],
+    toolsets: list["AbstractToolset[None]"] = [],
+    model_settings: "ModelSettings | None" = None,
+    history_processors: list["HistoryProcessor"] | None = None,
+    output_type: "OutputSpec[OutputDataT]" = str,
+    retries: int = 1,
+    yolo: bool = False,
+) -> "Agent[None, Any]":
+    from pydantic_ai import Agent, DeferredToolRequests
+    from pydantic_ai.toolsets import FunctionToolset
+
+    # Expand system prompt with references
+    effective_system_prompt = expand_prompt(system_prompt)
+
+    final_output_type = output_type
+    effective_toolsets = list(toolsets)
+    if tools:
+        effective_toolsets.append(FunctionToolset(tools=tools))
+
+    if not yolo:
+        final_output_type = output_type | DeferredToolRequests
+        effective_toolsets = [ts.approval_required() for ts in effective_toolsets]
+
+    if model is None:
+        model = default_llm_config.model
+
+    return Agent(
+        model=model,
+        output_type=final_output_type,
+        instructions=effective_system_prompt,
+        toolsets=effective_toolsets,
+        model_settings=model_settings,
+        history_processors=history_processors,
+        retries=retries,
+    )
+
+
+async def run_agent(
+    agent: "Agent[None, Any]",
+    message: str | None,
+    message_history: list[Any],
+    limiter: LLMLimiter,
+    attachments: list[Any] | None = None,
+    print_fn: Callable[[str], Any] = print,
+    event_handler: Callable[[Any], Any] | None = None,
+    tool_confirmation: Callable[[Any], Any] | None = None,
+) -> tuple[Any, list[Any]]:
+    """
+    Runs the agent with rate limiting, history management, and optional CLI confirmation loop.
+    Returns (result_output, new_message_history).
+    """
+    import asyncio
+
+    from pydantic_ai import AgentRunResultEvent, DeferredToolRequests
+
+    # Resolve tool confirmation callback (Arg > Context > None)
+    effective_tool_confirmation = tool_confirmation
+    if effective_tool_confirmation is None:
+        effective_tool_confirmation = tool_confirmation_var.get()
+
+    # Set context var for sub-agents
+    token = tool_confirmation_var.set(effective_tool_confirmation)
+
+    try:
+        # Expand user message with references
+        effective_message = expand_prompt(message) if message else message
+
+        # Prepare Prompt Content
+        prompt_content = _get_prompt_content(effective_message, attachments, print_fn)
+
+        # 1. Prune & Throttle
+        current_history = await _acquire_rate_limit(
+            limiter, prompt_content, message_history, print_fn
+        )
+        current_message = prompt_content
+        current_results = None
+
+        # 2. Execution Loop
+        while True:
+            result_output = None
+            run_history = []
+
+            async for event in agent.run_stream_events(
+                current_message,
+                message_history=current_history,
+                deferred_tool_results=current_results,
+            ):
+                await asyncio.sleep(0)
+                if isinstance(event, AgentRunResultEvent):
+                    result = event.result
+                    result_output = result.output
+                    run_history = result.all_messages()
+                if event_handler:
+                    await event_handler(event)
+
+            # Handle Deferred Calls
+            if isinstance(result_output, DeferredToolRequests):
+                current_results = await _process_deferred_requests(
+                    result_output, effective_tool_confirmation
+                )
+                if current_results is None:
+                    return result_output, run_history
+                # Prepare next iteration
+                current_message = None
+                current_history = run_history
+                continue
+            return result_output, run_history
+    finally:
+        tool_confirmation_var.reset(token)
+
+
+def _get_prompt_content(
+    message: str | None, attachments: list[Any] | None, print_fn: Callable[[str], Any]
+) -> "list[UserPromptPart] | str | None":
+    from pydantic_ai.messages import UserPromptPart
+
+    prompt_content = message
+    if attachments:
+        attachments = normalize_attachments(attachments, print_fn)
+        parts: list[UserPromptPart] = []
+        if message:
+            parts.append(UserPromptPart(content=message))
+        parts.extend(attachments)
+        prompt_content = parts
+    return prompt_content
+
+
+async def _acquire_rate_limit(
+    limiter: LLMLimiter,
+    message: str | None,
+    message_history: list[Any],
+    print_fn: Callable[[str], Any],
+) -> list[Any]:
+    """Prunes history and waits if rate limits are exceeded."""
+    if not message:
+        return message_history
+
+    # Prune
+    pruned_history = limiter.fit_context_window(message_history, message)
+
+    # Throttle
+    est_tokens = limiter.count_tokens(pruned_history) + limiter.count_tokens(message)
+    await limiter.acquire(
+        est_tokens, notifier=lambda msg: print_fn(msg) if msg else None
+    )
+
+    return pruned_history
+
+
+async def _process_deferred_requests(
+    result_output: "DeferredToolRequests",
+    effective_tool_confirmation: Callable[[Any], Any] | None,
+) -> "DeferredToolResults | None":
+    """Handles tool approvals/denials via callback or CLI fallback."""
+    import asyncio
+    import inspect
+
+    from pydantic_ai import DeferredToolResults, ToolApproved, ToolDenied
+
+    all_requests = (result_output.calls or []) + (result_output.approvals or [])
+    if not all_requests:
+        return None
+
+    current_results = DeferredToolResults()
+
+    for call in all_requests:
+        if effective_tool_confirmation:
+            res = effective_tool_confirmation(call)
+            if inspect.isawaitable(res):
+                result = await res
+            else:
+                result = res
+            current_results.approvals[call.tool_call_id] = result
+        else:
+            # CLI Fallback
+            prompt_text = f"Execute tool '{call.tool_name}' with args {call.args}?"
+            prompt_cli = f"\n[?] {prompt_text} (y/N) "
+
+            # We use asyncio.to_thread(input, ...) to avoid blocking the loop
+            user_input = await asyncio.to_thread(input, prompt_cli)
+            answer = user_input.strip().lower() in ("y", "yes")
+
+            if answer:
+                current_results.approvals[call.tool_call_id] = ToolApproved()
+            else:
+                current_results.approvals[call.tool_call_id] = ToolDenied("User denied")
+
+    return current_results

zrb/llm/agent/summarizer.py

@@ -0,0 +1,20 @@
+from typing import TYPE_CHECKING
+
+from zrb.llm.agent.agent import create_agent
+from zrb.llm.prompt.default import get_summarizer_system_prompt
+
+if TYPE_CHECKING:
+    from pydantic_ai import Agent
+    from pydantic_ai.models import Model
+
+
+def create_summarizer_agent(
+    model: "str | None | Model" = None,
+    system_prompt: str | None = None,
+) -> "Agent[None, str]":
+    effective_system_prompt = system_prompt or get_summarizer_system_prompt()
+
+    return create_agent(
+        model=model,
+        system_prompt=effective_system_prompt,
+    )

zrb/llm/app/__init__.py

@@ -0,0 +1,10 @@
+from zrb.llm.app.confirmation.allow_tool import allow_tool_usage
+from zrb.llm.app.confirmation.handler import ConfirmationMiddleware, last_confirmation
+from zrb.llm.app.confirmation.replace_confirmation import replace_confirmation
+
+__all__ = [
+    "allow_tool_usage",
+    "ConfirmationMiddleware",
+    "last_confirmation",
+    "replace_confirmation",
+]

zrb/llm/app/completion.py

@@ -0,0 +1,281 @@
+import os
+from datetime import datetime
+from typing import Iterable
+
+from prompt_toolkit.completion import (
+    CompleteEvent,
+    Completer,
+    Completion,
+    PathCompleter,
+)
+from prompt_toolkit.document import Document
+
+from zrb.llm.history_manager.any_history_manager import AnyHistoryManager
+from zrb.util.match import fuzzy_match
+
+
+class InputCompleter(Completer):
+    def __init__(
+        self,
+        history_manager: AnyHistoryManager,
+        attach_commands: list[str] = [],
+        exit_commands: list[str] = [],
+        info_commands: list[str] = [],
+        save_commands: list[str] = [],
+        load_commands: list[str] = [],
+        redirect_output_commands: list[str] = [],
+        summarize_commands: list[str] = [],
+        exec_commands: list[str] = [],
+    ):
+        self._history_manager = history_manager
+        self._attach_commands = attach_commands
+        self._exit_commands = exit_commands
+        self._info_commands = info_commands
+        self._save_commands = save_commands
+        self._load_commands = load_commands
+        self._redirect_output_commands = redirect_output_commands
+        self._summarize_commands = summarize_commands
+        self._exec_commands = exec_commands
+        # expanduser=True allows ~/path
+        self._path_completer = PathCompleter(expanduser=True)
+        # Cache for file listing to improve performance
+        self._file_cache: list[str] | None = None
+        self._file_cache_time = 0
+        self._cmd_history = self._get_cmd_history()
+
+    def get_completions(
+        self, document: Document, complete_event: CompleteEvent
+    ) -> Iterable[Completion]:
+        text_before_cursor = document.text_before_cursor.lstrip()
+        word = document.get_word_before_cursor(WORD=True)
+
+        all_commands = (
+            self._exit_commands
+            + self._attach_commands
+            + self._summarize_commands
+            + self._info_commands
+            + self._save_commands
+            + self._load_commands
+            + self._redirect_output_commands
+            + self._exec_commands
+        )
+        command_prefixes = {cmd[0] for cmd in all_commands if cmd}
+
+        # 1. Command and Argument Completion
+        if text_before_cursor and text_before_cursor[0] in command_prefixes:
+            parts = text_before_cursor.split()
+            # Check if we are typing the command itself or arguments
+            is_typing_command = len(parts) == 1 and not text_before_cursor.endswith(" ")
+            is_typing_arg = (len(parts) == 1 and text_before_cursor.endswith(" ")) or (
+                len(parts) >= 2
+            )
+
+            if is_typing_command:
+                lower_word = word.lower()
+                prefix = text_before_cursor[0]
+                for cmd in all_commands:
+                    if cmd.startswith(prefix) and cmd.lower().startswith(lower_word):
+                        yield Completion(cmd, start_position=-len(word))
+                return
+
+            if is_typing_arg:
+                cmd = parts[0]
+                arg_prefix = text_before_cursor[len(cmd) :].lstrip()
+
+                # Exec Command: Suggest History
+                if self._is_command(cmd, self._exec_commands):
+                    # Filter history
+                    matches = [h for h in self._cmd_history if h.startswith(arg_prefix)]
+                    # Sort matches by length (shorter first) as heuristic? Or just recent?
+                    # Since _cmd_history is set (unique), we lose order.
+                    # But Python 3.7+ dicts preserve insertion order, so if we used dict keys, we kept order.
+                    # Let's assume _get_cmd_history returns recent last.
+                    # We reverse to show most recent first.
+                    for h in reversed(matches):
+                        yield Completion(h, start_position=-len(arg_prefix))
+                    return
+
+                # Check if we are typing the second part (argument) strictly
+                # (Re-evaluating logic for other commands which only take 1 arg usually)
+                if not (
+                    (len(parts) == 1 and text_before_cursor.endswith(" "))
+                    or (len(parts) == 2 and not text_before_cursor.endswith(" "))
+                ):
+                    return
+
+                arg_prefix = parts[1] if len(parts) == 2 else ""
+
+                # Save Command: Suggest Timestamp
+                if self._is_command(cmd, self._save_commands):
+                    ts = datetime.now().strftime("%Y-%m-%d-%H-%M")
+                    if ts.startswith(arg_prefix):
+                        yield Completion(ts, start_position=-len(arg_prefix))
+                    return
+
+                # Redirect Command: Suggest Timestamp.txt
+                if self._is_command(cmd, self._redirect_output_commands):
+                    ts = datetime.now().strftime("%Y-%m-%d-%H-%M.txt")
+                    if ts.startswith(arg_prefix):
+                        yield Completion(ts, start_position=-len(arg_prefix))
+                    return
+
+                # Load Command: Search History
+                if self._is_command(cmd, self._load_commands):
+                    results = self._history_manager.search(arg_prefix)
+                    for res in results[:10]:
+                        yield Completion(res, start_position=-len(arg_prefix))
+                    return
+
+                # Attach Command: Suggest Files
+                if self._is_command(cmd, self._attach_commands):
+                    yield from self._get_file_completions(
+                        arg_prefix, complete_event, only_files=True
+                    )
+                    return
+
+                # Other commands (Exit, Info, Summarize) need no completion
+                return
+
+        # 2. File Completion (@)
+        if word.startswith("@"):
+            path_part = word[1:]
+            yield from self._get_file_completions(
+                path_part, complete_event, only_files=False
+            )
+
+    def _get_cmd_history(self) -> list[str]:
+        history_files = [
+            os.path.expanduser("~/.bash_history"),
+            os.path.expanduser("~/.zsh_history"),
+        ]
+        unique_cmds = {}  # Use dict to preserve order (insertion order)
+
+        for hist_file in history_files:
+            if not os.path.exists(hist_file):
+                continue
+            try:
+                with open(hist_file, "r", errors="ignore") as f:
+                    for line in f:
+                        line = line.strip()
+                        if not line:
+                            continue
+                        # Handle zsh timestamp format: : 1612345678:0;command
+                        if line.startswith(": ") and ";" in line:
+                            parts = line.split(";", 1)
+                            if len(parts) == 2:
+                                line = parts[1]
+
+                        if line:
+                            # Remove existing to update position to end (most recent)
+                            if line in unique_cmds:
+                                del unique_cmds[line]
+                            unique_cmds[line] = None
+            except Exception:
+                pass
+
+        return list(unique_cmds.keys())
+
+    def _is_command(self, cmd: str, cmd_list: list[str]) -> bool:
+        return cmd.lower() in [c.lower() for c in cmd_list]
+
+    def _get_file_completions(
+        self, text: str, complete_event: CompleteEvent, only_files: bool = False
+    ) -> Iterable[Completion]:
+        # Logic:
+        # - If text indicates path traversal (/, ., ~), use PathCompleter
+        # - Else, check file count. If < 5000, use Fuzzy. Else use PathCompleter.
+
+        if self._is_path_navigation(text):
+            yield from self._get_path_completions(text, complete_event, only_files)
+            return
+
+        # Count files (cached strategy could be added here if needed)
+        files = self._get_recursive_files(limit=5000)
+        if len(files) < 5000:
+            # Fuzzy Match
+            yield from self._get_fuzzy_completions(text, files, only_files)
+        else:
+            # Fallback to PathCompleter for large repos
+            yield from self._get_path_completions(text, complete_event, only_files)
+
+    def _is_path_navigation(self, text: str) -> bool:
+        return (
+            text.startswith("/")
+            or text.startswith(".")
+            or text.startswith("~")
+            or os.sep in text
+        )
+
+    def _get_path_completions(
+        self, text: str, complete_event: CompleteEvent, only_files: bool
+    ) -> Iterable[Completion]:
+        # PathCompleter needs a document where text represents the path
+        fake_document = Document(text=text, cursor_position=len(text))
+        for c in self._path_completer.get_completions(fake_document, complete_event):
+            if only_files:
+                # Check if the completed path is a directory
+                # Note: 'text' is the prefix. c.text is the completion suffix.
+                # We need to reconstruct full path to check isdir
+                # This is tricky with PathCompleter's internal logic.
+                # A simple heuristic: if it ends with path separator, it's a dir.
+                if c.text.endswith(os.sep):
+                    continue
+            yield c
+
+    def _get_fuzzy_completions(
+        self, text: str, files: list[str], only_files: bool
+    ) -> Iterable[Completion]:
+        matches = []
+        for f in files:
+            if only_files and f.endswith(os.sep):
+                continue
+            is_match, score = fuzzy_match(f, text)
+            if is_match:
+                matches.append((score, f))
+
+        # Sort by score (lower is better)
+        matches.sort(key=lambda x: x[0])
+
+        # Return top 20
+        for _, f in matches[:20]:
+            yield Completion(f, start_position=-len(text))
+
+    def _get_recursive_files(self, root: str = ".", limit: int = 5000) -> list[str]:
+        # Simple walker with exclusions
+        paths = []
+        # Check if current dir is hidden
+        cwd_is_hidden = os.path.basename(os.path.abspath(root)).startswith(".")
+
+        try:
+            for dirpath, dirnames, filenames in os.walk(root):
+                # Exclude hidden directories unless root is hidden
+                if not cwd_is_hidden:
+                    dirnames[:] = [d for d in dirnames if not d.startswith(".")]
+
+                # Exclude common ignores
+                dirnames[:] = [
+                    d
+                    for d in dirnames
+                    if d not in ("node_modules", "__pycache__", "venv", ".venv")
+                ]
+
+                rel_dir = os.path.relpath(dirpath, root)
+                if rel_dir == ".":
+                    rel_dir = ""
+
+                # Add directories
+                for d in dirnames:
+                    paths.append(os.path.join(rel_dir, d) + os.sep)
+                    if len(paths) >= limit:
+                        return paths
+
+                # Add files
+                for f in filenames:
+                    if not cwd_is_hidden and f.startswith("."):
+                        continue
+                    paths.append(os.path.join(rel_dir, f))
+                    if len(paths) >= limit:
+                        return paths
+        except Exception:
+            pass
+        return paths

zrb/llm/app/confirmation/allow_tool.py

@@ -0,0 +1,66 @@
+import json
+import re
+from typing import Any, Awaitable, Callable, Dict, Optional
+
+from zrb.llm.app.confirmation.handler import ConfirmationMiddleware, UIProtocol
+
+
+def allow_tool_usage(
+    tool_name: str, kwargs: Optional[Dict[str, str]] = None
+) -> ConfirmationMiddleware:
+    """
+    Creates a confirmation middleware that automatically approves a tool execution
+    if it matches the specified tool_name and argument constraints.
+
+    :param tool_name: The name of the tool to allow.
+    :param kwargs: A dictionary of regex patterns for arguments.
+                   If None or empty, the tool is allowed regardless of arguments.
+                   If provided, arguments in the tool call must match the regex patterns
+                   specified in kwargs (only for arguments present in both).
+    :return: A ConfirmationMiddleware function.
+    """
+    from pydantic_ai import ToolApproved
+
+    async def middleware(
+        ui: UIProtocol,
+        call: Any,
+        response: str,
+        next_handler: Callable[[UIProtocol, Any, str], Awaitable[Any]],
+    ) -> Any:
+        # Check if tool name matches
+        if call.tool_name != tool_name:
+            return await next_handler(ui, call, response)
+
+        # If kwargs is empty or None, approve
+        if not kwargs:
+            ui.append_to_output(f"\n✅ Auto-approved tool: {tool_name}")
+            return ToolApproved()
+
+        # Parse arguments
+        try:
+            args = call.args
+            if isinstance(args, str):
+                args = json.loads(args)
+
+            if not isinstance(args, dict):
+                # If args is not a dict (e.g. primitive), and kwargs is not empty,
+                # we assume it doesn't match complex constraints (or we can't check keys).
+                # So we delegate to the next handler.
+                return await next_handler(ui, call, response)
+
+        except (json.JSONDecodeError, ValueError):
+            return await next_handler(ui, call, response)
+
+        # Check constraints
+        # "all parameter in the call parameter has to match the ones in kwargs (if that parameter defined in the kwargs)"
+        for arg_name, arg_value in args.items():
+            if arg_name in kwargs:
+                pattern = kwargs[arg_name]
+                # Convert arg_value to string for regex matching
+                if not re.search(pattern, str(arg_value)):
+                    return await next_handler(ui, call, response)
+
+        ui.append_to_output(f"\n✅ Auto-approved tool: {tool_name} with matching args")
+        return ToolApproved()
+
+    return middleware