zrb 1.8.10__py3-none-any.whl → 1.21.29__py3-none-any.whl

zrb/builtin/llm/tool/file.py

@@ -5,32 +5,12 @@ import re
 from typing import Any, Optional
 
 from zrb.builtin.llm.tool.sub_agent import create_sub_agent_tool
+from zrb.config.config import CFG
+from zrb.config.llm_rate_limitter import llm_rate_limitter
 from zrb.context.any_context import AnyContext
-from zrb.llm_rate_limitter import llm_rate_limitter
+from zrb.task.llm.file_tool_model import FileReplacement, FileToRead, FileToWrite
 from zrb.util.file import read_file, read_file_with_line_numbers, write_file
 
-_EXTRACT_INFO_FROM_FILE_SYSTEM_PROMPT = """
-You are an extraction info agent.
-Your goal is to help to extract relevant information to help the main assistant.
-You write your output is in markdown format containing path and relevant information.
-Extract only information that relevant to main assistant's goal.
-
-Extracted Information format (Use this as reference, extract relevant information only):
-# imports
-- <imported-package>
-- ...
-# variables
-- <variable-type> <variable-name>: <the-purpose-of-the-variable>
-- ...
-# functions
-- <function-name>:
-  - parameters: <parameters>
-  - logic/description: <what-the-function-do-and-how-it-works>
-...
-...
-""".strip()
-
-
 DEFAULT_EXCLUDED_PATTERNS = [
     # Common Python artifacts
     "__pycache__",
@@ -102,21 +82,25 @@ DEFAULT_EXCLUDED_PATTERNS = [
 
 def list_files(
     path: str = ".",
-    recursive: bool = True,
     include_hidden: bool = False,
+    depth: int = 3,
     excluded_patterns: Optional[list[str]] = None,
-) -> str:
-    """List files/directories in a path, excluding specified patterns.
+) -> dict[str, list[str]]:
+    """
+    Lists files recursively up to a specified depth.
+
+    Example:
+    list_files(path='src', include_hidden=False, depth=2)
+
     Args:
-        path (str): Path to list. Pass exactly as provided, including '~'. Defaults to ".".
-        recursive (bool): List recursively. Defaults to True.
-        include_hidden (bool): Include hidden files/dirs. Defaults to False.
-        excluded_patterns (Optional[List[str]]): List of glob patterns to exclude.
-            Defaults to a comprehensive list of common temporary/artifact patterns.
+        path (str): Directory path. Defaults to current directory.
+        include_hidden (bool): Include hidden files. Defaults to False.
+        depth (int): Maximum depth to traverse. Defaults to 3.
+            Minimum depth is 1 (current directory only).
+        excluded_patterns (list[str]): Glob patterns to exclude.
+
     Returns:
-        str: JSON string: {"files": ["file1.txt", ...]} or {"error": "..."}
-    Raises:
-        Exception: If an error occurs.
+        dict: {'files': [relative_paths]}
     """
     all_files: list[str] = []
     abs_path = os.path.abspath(os.path.expanduser(path))
@@ -130,50 +114,30 @@ def list_files(
         if excluded_patterns is not None
         else DEFAULT_EXCLUDED_PATTERNS
     )
+    if depth <= 0:
+        depth = 1
     try:
-        if recursive:
-            for root, dirs, files in os.walk(abs_path, topdown=True):
-                # Filter directories in-place
-                dirs[:] = [
-                    d
-                    for d in dirs
-                    if (include_hidden or not _is_hidden(d))
-                    and not is_excluded(d, patterns_to_exclude)
-                ]
-                # Process files
-                for filename in files:
-                    if (include_hidden or not _is_hidden(filename)) and not is_excluded(
-                        filename, patterns_to_exclude
-                    ):
-                        full_path = os.path.join(root, filename)
-                        # Check rel path for patterns like '**/node_modules/*'
-                        rel_full_path = os.path.relpath(full_path, abs_path)
-                        is_rel_path_excluded = is_excluded(
-                            rel_full_path, patterns_to_exclude
-                        )
-                        if not is_rel_path_excluded:
-                            all_files.append(full_path)
-        else:
-            # Non-recursive listing (top-level only)
-            for item in os.listdir(abs_path):
-                full_path = os.path.join(abs_path, item)
-                # Include both files and directories if not recursive
-                if (include_hidden or not _is_hidden(item)) and not is_excluded(
-                    item, patterns_to_exclude
+        initial_depth = abs_path.rstrip(os.sep).count(os.sep)
+        for root, dirs, files in os.walk(abs_path, topdown=True):
+            current_depth = root.rstrip(os.sep).count(os.sep) - initial_depth
+            if current_depth >= depth - 1:
+                del dirs[:]
+            dirs[:] = [
+                d
+                for d in dirs
+                if (include_hidden or not _is_hidden(d))
+                and not is_excluded(d, patterns_to_exclude)
+            ]
+            for filename in files:
+                if (include_hidden or not _is_hidden(filename)) and not is_excluded(
+                    filename, patterns_to_exclude
                 ):
-                    all_files.append(full_path)
-        # Return paths relative to the original path requested
-        try:
-            rel_files = [os.path.relpath(f, abs_path) for f in all_files]
-            return json.dumps({"files": sorted(rel_files)})
-        except (
-            ValueError
-        ) as e:  # Handle case where path is '.' and abs_path is CWD root
-            if "path is on mount '" in str(e) and "' which is not on mount '" in str(e):
-                # If paths are on different mounts, just use absolute paths
-                rel_files = all_files
-                return json.dumps({"files": sorted(rel_files)})
-            raise
+                    full_path = os.path.join(root, filename)
+                    rel_full_path = os.path.relpath(full_path, abs_path)
+                    if not is_excluded(rel_full_path, patterns_to_exclude):
+                        all_files.append(rel_full_path)
+        return {"files": sorted(all_files)}
+
     except (OSError, IOError) as e:
         raise OSError(f"Error listing files in {path}: {e}")
     except Exception as e:
@@ -210,97 +174,146 @@ def is_excluded(name: str, patterns: list[str]) -> bool:
 
 
 def read_from_file(
-    path: str,
-    start_line: Optional[int] = None,
-    end_line: Optional[int] = None,
-) -> str:
-    """Read file content (or specific lines) at a path, including line numbers.
+    file: FileToRead | list[FileToRead],
+) -> dict[str, Any]:
+    """
+    Reads content from one or more files, optionally specifying line ranges.
+
+    Examples:
+    ```
+    # Read entire content of a single file
+    read_from_file(file={'path': 'path/to/file.txt'})
+
+    # Read specific lines from a file
+    # The content will be returned with line numbers in the format: "LINE_NUMBER | line content"
+    read_from_file(file={'path': 'path/to/large_file.log', 'start_line': 100, 'end_line': 150})
+
+    # Read multiple files
+    read_from_file(file=[
+        {'path': 'path/to/file1.txt'},
+        {'path': 'path/to/file2.txt', 'start_line': 1, 'end_line': 5}
+    ])
+    ```
+
     Args:
-        path (str): Path to read. Pass exactly as provided, including '~'.
-        start_line (Optional[int]): Starting line number (1-based).
-            Defaults to None (start of file).
-        end_line (Optional[int]): Ending line number (1-based, inclusive).
-            Defaults to None (end of file).
+        file (FileToRead | list[FileToRead]): A single file configuration or a list of them.
+
     Returns:
-        str: JSON: {"path": "...", "content": "...", "start_line": N, ...} or {"error": "..."}
-        The content includes line numbers.
-    Raises:
-        Exception: If an error occurs.
+        dict: Content and metadata for a single file, or a dict of results for multiple files.
+        The `content` field in the returned dictionary will have line numbers in the format: "LINE_NUMBER | line content"
     """
-    abs_path = os.path.abspath(os.path.expanduser(path))
-    # Check if file exists
-    if not os.path.exists(abs_path):
-        raise FileNotFoundError(f"File not found: {path}")
-    try:
-        content = read_file_with_line_numbers(abs_path)
-        lines = content.splitlines()
-        total_lines = len(lines)
-        # Adjust line indices (convert from 1-based to 0-based)
-        start_idx = (start_line - 1) if start_line is not None else 0
-        end_idx = end_line if end_line is not None else total_lines
-        # Validate indices
-        if start_idx < 0:
-            start_idx = 0
-        if end_idx > total_lines:
-            end_idx = total_lines
-        if start_idx > end_idx:
-            start_idx = end_idx
-        # Select the lines for the result
-        selected_lines = lines[start_idx:end_idx]
-        content_result = "\n".join(selected_lines)
-        return json.dumps(
-            {
+    is_list = isinstance(file, list)
+    files = file if is_list else [file]
+
+    results = {}
+    for file_config in files:
+        path = file_config["path"]
+        start_line = file_config.get("start_line", None)
+        end_line = file_config.get("end_line", None)
+        try:
+            abs_path = os.path.abspath(os.path.expanduser(path))
+            if not os.path.exists(abs_path):
+                raise FileNotFoundError(f"File not found: {path}")
+
+            content = read_file_with_line_numbers(abs_path)
+            lines = content.splitlines()
+            total_lines = len(lines)
+
+            start_idx = (start_line - 1) if start_line is not None else 0
+            end_idx = end_line if end_line is not None else total_lines
+
+            if start_idx < 0:
+                start_idx = 0
+            if end_idx > total_lines:
+                end_idx = total_lines
+            if start_idx > end_idx:
+                start_idx = end_idx
+
+            selected_lines = lines[start_idx:end_idx]
+            content_result = "\n".join(selected_lines)
+
+            results[path] = {
                 "path": path,
                 "content": content_result,
-                "start_line": start_idx + 1,  # Convert back to 1-based for output
-                "end_line": end_idx,  # end_idx is already exclusive upper bound
+                "start_line": start_idx + 1,
+                "end_line": end_idx,
                 "total_lines": total_lines,
             }
-        )
-    except (OSError, IOError) as e:
-        raise OSError(f"Error reading file {path}: {e}")
-    except Exception as e:
-        raise RuntimeError(f"Unexpected error reading file {path}: {e}")
+        except Exception as e:
+            if not is_list:
+                if isinstance(e, (OSError, IOError)):
+                    raise OSError(f"Error reading file {path}: {e}") from e
+                raise RuntimeError(f"Unexpected error reading file {path}: {e}") from e
+            results[path] = f"Error reading file: {e}"
+
+    if is_list:
+        return results
+
+    return results[files[0]["path"]]
 
 
 def write_to_file(
-    path: str,
-    content: str,
-    line_count: int,
-) -> str:
-    """Write full content to a file. Creates/overwrites file.
+    file: FileToWrite | list[FileToWrite],
+) -> str | dict[str, Any]:
+    """
+    Writes content to one or more files, with options for overwrite, append, or exclusive
+    creation.
+
+    **CRITICAL - PREVENT JSON ERRORS:**
+    1. **ESCAPING:** Do NOT double-escape quotes.
+       - CORRECT: "content": "He said \"Hello\""
+       - WRONG:   "content": "He said \\"Hello\\""  <-- This breaks JSON parsing!
+    2. **SIZE LIMIT:** Content MUST NOT exceed 4000 characters.
+       - Exceeding this causes truncation and EOF errors.
+       - Split larger content into multiple sequential calls (first 'w', then 'a').
+
+    Examples:
+    ```
+    # Overwrite 'file.txt' with initial content
+    write_to_file(file={'path': 'path/to/file.txt', 'content': 'Initial content.'})
+
+    # Append a second chunk to 'file.txt' (note the newline at the beginning of the content)
+    write_to_file(file={'path': 'path/to/file.txt', 'content': '\nSecond chunk.', 'mode': 'a'})
+
+    # Write to multiple files
+    write_to_file(file=[
+        {'path': 'path/to/file1.txt', 'content': 'Content for file 1'},
+        {'path': 'path/to/file2.txt', 'content': 'Content for file 2', 'mode': 'w'}
+    ])
+    ```
+
     Args:
-        path (str): Path to write. Pass exactly as provided, including '~'.
-        content (str): Full file content.
-            MUST be complete, no truncation/omissions. Exclude line numbers.
-        line_count (int): Number of lines in the provided content.
+        file (FileToWrite | list[FileToWrite]): A single file configuration or a list of them.
+
     Returns:
-        str: JSON: {"success": true, "path": "f.txt", "warning": "..."} or {"error": "..."}
-    Raises:
-        Exception: If an error occurs.
+        Success message for single file, or dict with success/errors for multiple files.
     """
-    actual_lines = len(content.splitlines())
-    warning = None
-    if actual_lines != line_count:
-        warning = (
-            f"Provided line_count ({line_count}) does not match actual "
-            f"content lines ({actual_lines}) for file {path}"
-        )
-    try:
-        abs_path = os.path.abspath(os.path.expanduser(path))
-        # Ensure directory exists
-        directory = os.path.dirname(abs_path)
-        if directory and not os.path.exists(directory):
-            os.makedirs(directory, exist_ok=True)
-        write_file(abs_path, content)
-        result_data = {"success": True, "path": path}
-        if warning:
-            result_data["warning"] = warning
-        return json.dumps(result_data)
-    except (OSError, IOError) as e:
-        raise OSError(f"Error writing file {path}: {e}")
-    except Exception as e:
-        raise RuntimeError(f"Unexpected error writing file {path}: {e}")
+    # Normalize to list
+    files = file if isinstance(file, list) else [file]
+
+    success = []
+    errors = {}
+    for file_config in files:
+        path = file_config["path"]
+        content = file_config["content"]
+        mode = file_config.get("mode", "w")
+        try:
+            abs_path = os.path.abspath(os.path.expanduser(path))
+            # The underlying utility creates the directory, so we don't need to do it here.
+            write_file(abs_path, content, mode=mode)
+            success.append(path)
+        except Exception as e:
+            errors[path] = f"Error writing file: {e}"
+
+    # Return appropriate response based on input type
+    if isinstance(file, list):
+        return {"success": success, "errors": errors}
+    else:
+        if errors:
+            raise RuntimeError(
+                f"Error writing file {file['path']}: {errors[file['path']]}"
+            )
+        return f"Successfully wrote to file: {file['path']} in mode '{file.get('mode', 'w')}'"
 
 
 def search_files(
@@ -308,18 +321,21 @@ def search_files(
     regex: str,
     file_pattern: Optional[str] = None,
     include_hidden: bool = True,
-) -> str:
-    """Search files in a directory using regex, showing context.
+) -> dict[str, Any]:
+    """
+    Searches for a regex pattern in files within a directory.
+
+    Example:
+    search_files(path='src', regex='class \\w+', file_pattern='*.py', include_hidden=False)
+
     Args:
-        path (str): Path to search. Pass exactly as provided, including '~'.
-        regex (str): Python regex pattern to search for.
-        file_pattern (Optional[str]): Glob pattern to filter files
-            (e.g., '*.py'). Defaults to None.
-        include_hidden (bool): Include hidden files/dirs. Defaults to True.
+        path (str): Directory to search.
+        regex (str): Regex pattern.
+        file_pattern (str): Glob pattern filter.
+        include_hidden (bool): Include hidden files.
+
     Returns:
-        str: JSON: {"summary": "...", "results": [{"file":"f.py", ...}]} or {"error": "..."}
-    Raises:
-        Exception: If error occurs or regex is invalid.
+        dict: Summary and list of matches.
     """
     try:
         pattern = re.compile(regex)
@@ -366,9 +382,7 @@ def search_files(
                 f"Found {match_count} matches in {file_match_count} files "
                 f"(searched {searched_file_count} files)."
             )
-        return json.dumps(
-            search_results
-        )  # No need for pretty printing for LLM consumption
+        return search_results
     except (OSError, IOError) as e:
         raise OSError(f"Error searching files in {path}: {e}")
     except Exception as e:
@@ -376,7 +390,9 @@ def search_files(
 
 
 def _get_file_matches(
-    file_path: str, pattern: re.Pattern, context_lines: int = 2
+    file_path: str,
+    pattern: re.Pattern,
+    context_lines: int = 2,
 ) -> list[dict[str, Any]]:
     """Search for regex matches in a file with context."""
     try:
@@ -406,95 +422,139 @@ def _get_file_matches(
         raise RuntimeError(f"Unexpected error processing {file_path}: {e}")
 
 
-def apply_diff(
-    path: str,
-    start_line: int,
-    end_line: int,
-    search_content: str,
-    replace_content: str,
-) -> str:
-    """Apply a precise search/replace to a file based on line numbers and content.
+def replace_in_file(
+    file: FileReplacement | list[FileReplacement],
+) -> str | dict[str, Any]:
+    """
+    Replaces exact text in files.
+
+    **CRITICAL INSTRUCTIONS:**
+    1. **READ FIRST:** Use `read_file` to get exact content. Do not guess.
+    2. **EXACT MATCH:** `old_text` must match file content EXACTLY (whitespace, newlines).
+    3. **ESCAPING:** Do NOT double-escape quotes in `new_text`. Use `\"`, not `\\"`.
+    4. **SIZE LIMIT:** `new_text` MUST NOT exceed 4000 chars to avoid truncation/EOF errors.
+    5. **MINIMAL CONTEXT:** Keep `old_text` small (target lines + 2-3 context lines).
+    6. **DEFAULT:** Replaces **ALL** occurrences. Set `count=1` for first occurrence only.
+
+    Examples:
+    ```
+    # Replace ALL occurrences
+    replace_in_file(file=[
+        {'path': 'file.txt', 'old_text': 'foo', 'new_text': 'bar'},
+        {'path': 'file.txt', 'old_text': 'baz', 'new_text': 'qux'}
+    ])
+
+    # Replace ONLY the first occurrence
+    replace_in_file(
+        file={'path': 'file.txt', 'old_text': 'foo', 'new_text': 'bar', 'count': 1}
+    )
+
+    # Replace code block (include context for safety)
+    replace_in_file(
+        file={
+            'path': 'app.py',
+            'old_text': '    def old_fn():\n        pass',
+            'new_text': '    def new_fn():\n        pass'
+        }
+    )
+    ```
+
     Args:
-        path (str): Path to modify. Pass exactly as provided, including '~'.
-        start_line (int): The 1-based starting line number of the content to replace.
-        end_line (int): The 1-based ending line number (inclusive) of the content to replace.
-        search_content (str): The exact content expected to be found in the specified
-            line range. Must exactly match file content including whitespace/indentation,
-            excluding line numbers.
-        replace_content (str): The new content to replace the search_content with.
-            Excluding line numbers.
+        file: Single replacement config or list of them.
+
     Returns:
-        str: JSON: {"success": true, "path": "f.py"} or {"success": false, "error": "..."}
-    Raises:
-        Exception: If an error occurs.
+        Success message or error dict.
     """
-    abs_path = os.path.abspath(os.path.expanduser(path))
-    if not os.path.exists(abs_path):
-        raise FileNotFoundError(f"File not found: {path}")
-    try:
-        content = read_file(abs_path)
-        lines = content.splitlines()
-        if start_line < 1 or end_line > len(lines) or start_line > end_line:
-            raise ValueError(
-                f"Invalid line range {start_line}-{end_line} for file with {len(lines)} lines"
-            )
-        original_content = "\n".join(lines[start_line - 1 : end_line])
-        if original_content != search_content:
-            error_message = (
-                f"Search content does not match file content at "
-                f"lines {start_line}-{end_line}.\n"
-                f"Expected ({len(search_content.splitlines())} lines):\n"
-                f"---\n{search_content}\n---\n"
-                f"Actual ({len(lines[start_line-1:end_line])} lines):\n"
-                f"---\n{original_content}\n---"
-            )
-            return json.dumps({"success": False, "path": path, "error": error_message})
-        new_lines = (
-            lines[: start_line - 1] + replace_content.splitlines() + lines[end_line:]
-        )
-        new_content = "\n".join(new_lines)
-        if content.endswith("\n"):
-            new_content += "\n"
-        write_file(abs_path, new_content)
-        return json.dumps({"success": True, "path": path})
-    except ValueError as e:
-        raise ValueError(f"Error parsing diff: {e}")
-    except (OSError, IOError) as e:
-        raise OSError(f"Error applying diff to {path}: {e}")
-    except Exception as e:
-        raise RuntimeError(f"Unexpected error applying diff to {path}: {e}")
+    # Normalize to list
+    file_replacements = file if isinstance(file, list) else [file]
+    # Group replacements by file path to minimize file I/O
+    replacements_by_path: dict[str, list[FileReplacement]] = {}
+    for r in file_replacements:
+        path = r["path"]
+        if path not in replacements_by_path:
+            replacements_by_path[path] = []
+        replacements_by_path[path].append(r)
+    success = []
+    errors = {}
+    for path, replacements in replacements_by_path.items():
+        try:
+            abs_path = os.path.abspath(os.path.expanduser(path))
+            if not os.path.exists(abs_path):
+                raise FileNotFoundError(f"File not found: {path}")
+            content = read_file(abs_path)
+            original_content = content
+            # Apply all replacements for this file
+            for replacement in replacements:
+                old_text = replacement["old_text"]
+                new_text = replacement["new_text"]
+                count = replacement.get("count", -1)
+                if old_text not in content:
+                    raise ValueError(f"old_text not found in file: {path}")
+                # Replace occurrences
+                content = content.replace(old_text, new_text, count)
+            # Only write if content actually changed
+            if content != original_content:
+                write_file(abs_path, content)
+                success.append(path)
+            else:
+                success.append(f"{path} (no changes needed)")
+        except Exception as e:
+            errors[path] = f"Error applying replacement to {path}: {e}"
+    # Return appropriate response based on input type
+    if isinstance(file, list):
+        return {"success": success, "errors": errors}
+    path = file["path"]
+    if errors:
+        error_message = errors[path]
+        raise RuntimeError(f"Error applying replacement to {path}: {error_message}")
+    return f"Successfully applied replacement(s) to {path}"
 
 
 async def analyze_file(
-    ctx: AnyContext, path: str, query: str, token_limit: int = 30000
-) -> str:
-    """Analyze file using LLM capability to reduce context usage.
-    Use this tool for:
-    - summarization
-    - outline/structure extraction
-    - code review
-    - other tasks
+    ctx: AnyContext, path: str, query: str, token_threshold: int | None = None
+) -> dict[str, Any]:
+    """
+    Analyzes a file using a sub-agent for complex questions.
+
+    Example:
+    analyze_file(path='src/main.py', query='Summarize the main function.')
+
     Args:
-        path (str): File path to be analyze. Pass exactly as provided, including '~'.
-        query(str): Instruction to analyze the file
-        token_limit(Optional[int]): Max token length to be taken from file
+        ctx (AnyContext): The execution context.
+        path (str): The path to the file to analyze.
+        query (str): A specific analysis query with clear guidelines and
+            necessary information.
+        token_threshold (int | None): Max tokens.
+
     Returns:
-        str: The analysis result
-    Raises:
-        Exception: If an error occurs.
+        Analysis results.
     """
+    if token_threshold is None:
+        token_threshold = CFG.LLM_FILE_ANALYSIS_TOKEN_THRESHOLD
     abs_path = os.path.abspath(os.path.expanduser(path))
     if not os.path.exists(abs_path):
         raise FileNotFoundError(f"File not found: {path}")
     file_content = read_file(abs_path)
     _analyze_file = create_sub_agent_tool(
         tool_name="analyze_file",
-        tool_description="analyze file with LLM capability",
-        system_prompt=_EXTRACT_INFO_FROM_FILE_SYSTEM_PROMPT,
+        tool_description=(
+            "Analyze file content using LLM sub-agent "
+            "for complex questions about code structure, documentation "
+            "quality, or file-specific analysis. Use for questions that "
+            "require understanding beyond simple text reading."
+        ),
+        system_prompt=CFG.LLM_FILE_EXTRACTOR_SYSTEM_PROMPT,
         tools=[read_from_file, search_files],
+        auto_summarize=False,
+        remember_history=False,
     )
     payload = json.dumps(
-        {"instruction": query, "file_path": abs_path, "file_content": file_content}
+        {
+            "instruction": query,
+            "file_path": abs_path,
+            "file_content": llm_rate_limitter.clip_prompt(
+                file_content, token_threshold
+            ),
+        }
     )
-    clipped_payload = llm_rate_limitter.clip_prompt(payload, token_limit)
-    return await _analyze_file(ctx, clipped_payload)
+    return await _analyze_file(ctx, payload)