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

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))

zrb/llm/tool/file.py

@@ -0,0 +1,419 @@
+import fnmatch
+import os
+import re
+from typing import Any
+
+DEFAULT_EXCLUDED_PATTERNS = [
+    "__pycache__",
+    "*.pyc",
+    "*.pyo",
+    "*.pyd",
+    ".Python",
+    "build",
+    "dist",
+    ".env",
+    ".venv",
+    "env",
+    "venv",
+    ".idea",
+    ".vscode",
+    ".git",
+    "node_modules",
+    ".pytest_cache",
+    ".coverage",
+    "htmlcov",
+]
+
+
+def list_files(
+    path: str = ".",
+    include_hidden: bool = False,
+    depth: int = 3,
+    excluded_patterns: list[str] | None = None,
+) -> dict[str, list[str]]:
+    """
+    Recursively explores and lists files within a directory tree up to a defined depth.
+
+    **WHEN TO USE:**
+    - To discover the project structure or find specific files when the path is unknown.
+    - To verify the existence of files in a directory.
+
+    **EFFICIENCY TIP:**
+    - Do NOT use this tool if you already know the file path. Use `read_file` directly.
+    - Keep `depth` low (default 3) to avoid overwhelming output.
+
+    **ARGS:**
+    - `path`: The root directory to start the search from.
+    - `include_hidden`: If True, includes hidden files and directories (starting with `.`).
+    - `depth`: Maximum levels of directories to descend.
+    - `excluded_patterns`: List of glob patterns to ignore.
+    """
+    all_files: list[str] = []
+    abs_path = os.path.abspath(os.path.expanduser(path))
+    if not os.path.exists(abs_path):
+        raise FileNotFoundError(f"Path does not exist: {path}")
+
+    patterns_to_exclude = (
+        excluded_patterns
+        if excluded_patterns is not None
+        else DEFAULT_EXCLUDED_PATTERNS
+    )
+    if depth <= 0:
+        depth = 1
+
+    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 d.startswith("."))
+            and not _is_excluded(d, patterns_to_exclude)
+        ]
+
+        for filename in files:
+            if (include_hidden or not filename.startswith(".")) and not _is_excluded(
+                filename, patterns_to_exclude
+            ):
+                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)}
+
+
+def read_file(
+    path: str, start_line: int | None = None, end_line: int | None = None
+) -> str:
+    """
+    Reads content from a file, optionally specifying a line range.
+
+    **EFFICIENCY TIP:**
+    - Prefer reading the **entire file** at once for full context (imports, class definitions).
+    - Only use `start_line` and `end_line` for extremely large files (e.g., logs).
+
+    **ARGS:**
+    - `path`: Path to the file to read.
+    - `start_line`: The 1-based line number to start reading from.
+    - `end_line`: The 1-based line number to stop reading at (inclusive).
+    """
+    abs_path = os.path.abspath(os.path.expanduser(path))
+    if not os.path.exists(abs_path):
+        return f"Error: File not found: {path}"
+
+    try:
+        with open(abs_path, "r", encoding="utf-8") as f:
+            lines = f.readlines()
+
+        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 = "".join(selected_lines)
+
+        if start_line is not None or end_line is not None:
+            return f"File: {path} (Lines {start_idx + 1}-{end_idx} of {total_lines})\n{content_result}"
+        return content_result
+
+    except Exception as e:
+        return f"Error reading file {path}: {e}"
+
+
+def read_files(paths: list[str]) -> dict[str, str]:
+    """
+    Reads content from multiple files simultaneously.
+
+    **USAGE:**
+    - Use this when you need context from several related files (e.g., a class definition and its tests).
+
+    **ARGS:**
+    - `paths`: List of file paths to read.
+    """
+    results = {}
+    for path in paths:
+        results[path] = read_file(path)
+    return results
+
+
+def write_file(path: str, content: str, mode: str = "w") -> str:
+    """
+    Writes or appends content to a file.
+
+    **CRITICAL - PREVENT ERRORS:**
+    1. **ESCAPING:** Do NOT double-escape quotes in your JSON tool call.
+    2. **SIZE LIMIT:** DO NOT write more than 4000 characters in a single call.
+    3. **CHUNKING:** For large files, use `mode="w"` for the first chunk and `mode="a"` for the rest.
+
+    **ARGS:**
+    - `path`: Target file path.
+    - `content`: Text content to write.
+    - `mode`: File opening mode ("w" to overwrite, "a" to append).
+    """
+    abs_path = os.path.abspath(os.path.expanduser(path))
+    try:
+        os.makedirs(os.path.dirname(abs_path), exist_ok=True)
+        with open(abs_path, mode, encoding="utf-8") as f:
+            f.write(content)
+        return f"Successfully wrote to {path}"
+    except Exception as e:
+        return f"Error writing to file {path}: {e}"
+
+
+def write_files(files: list[dict[str, str]]) -> dict[str, str]:
+    """
+    Performs batch write operations to multiple files.
+
+    **ARGS:**
+    - `files`: A list of dictionaries, each containing:
+        - `path` (str): Target file path.
+        - `content` (str): Text to write.
+        - `mode` (str, optional): "w" (overwrite, default) or "a" (append).
+    """
+    results = {}
+    for file_info in files:
+        path = file_info.get("path")
+        content = file_info.get("content")
+        mode = file_info.get("mode", "w")
+        if not path or content is None:
+            results[str(path)] = "Error: Missing path or content"
+            continue
+        results[path] = write_file(path, content, mode)
+    return results
+
+
+def replace_in_file(path: str, old_text: str, new_text: str, count: int = -1) -> str:
+    """
+    Replaces exact text sequences within a file.
+
+    **CRITICAL INSTRUCTIONS:**
+    1. **PRECISION:** `old_text` must match the file content EXACTLY.
+    2. **READ FIRST:** Always `read_file` before replacing.
+    3. **MINIMAL CONTEXT:** Include 2-3 lines of context in `old_text` to ensure uniqueness.
+
+    **ARGS:**
+    - `path`: Path to the file to modify.
+    - `old_text`: The exact literal text to be replaced.
+    - `new_text`: The replacement text.
+    - `count`: Number of occurrences to replace (default -1 for all).
+    """
+    abs_path = os.path.abspath(os.path.expanduser(path))
+    if not os.path.exists(abs_path):
+        return f"Error: File not found: {path}"
+
+    try:
+        with open(abs_path, "r", encoding="utf-8") as f:
+            content = f.read()
+
+        if old_text not in content:
+            return f"Error: '{old_text}' not found in {path}"
+
+        new_content = content.replace(old_text, new_text, count)
+
+        if content == new_content:
+            return f"No changes made to {path}"
+
+        with open(abs_path, "w", encoding="utf-8") as f:
+            f.write(new_content)
+        return f"Successfully updated {path}"
+    except Exception as e:
+        return f"Error replacing text in {path}: {e}"
+
+
+def search_files(
+    path: str,
+    regex: str,
+    file_pattern: str | None = None,
+    include_hidden: bool = True,
+) -> dict[str, Any]:
+    """
+    Searches for a regular expression pattern within files.
+
+    **WHEN TO USE:**
+    - To find usages of a function, variable, or string across the project.
+
+    **ARGS:**
+    - `path`: Root directory to search.
+    - `regex`: A standard Python regular expression.
+    - `file_pattern`: Optional glob (e.g., "*.py") to restrict the search.
+    - `include_hidden`: Whether to search in hidden files/dirs.
+    """
+    try:
+        pattern = re.compile(regex)
+    except re.error as e:
+        return {"error": f"Invalid regex pattern: {e}"}
+
+    search_results = {"summary": "", "results": []}
+    match_count = 0
+    searched_file_count = 0
+    file_match_count = 0
+
+    abs_path = os.path.abspath(os.path.expanduser(path))
+    if not os.path.exists(abs_path):
+        return {"error": f"Path not found: {path}"}
+
+    try:
+        for root, dirs, files in os.walk(abs_path):
+            # Skip hidden directories
+            dirs[:] = [d for d in dirs if include_hidden or not d.startswith(".")]
+            for filename in files:
+                # Skip hidden files
+                if not include_hidden and filename.startswith("."):
+                    continue
+                # Apply file pattern filter if provided
+                if file_pattern and not fnmatch.fnmatch(filename, file_pattern):
+                    continue
+
+                file_path = os.path.join(root, filename)
+                rel_file_path = os.path.relpath(file_path, os.getcwd())
+                searched_file_count += 1
+
+                try:
+                    matches = _get_file_matches(file_path, pattern)
+                    if matches:
+                        file_match_count += 1
+                        match_count += len(matches)
+                        search_results["results"].append(
+                            {"file": rel_file_path, "matches": matches}
+                        )
+                except Exception:
+                    # Ignore read errors for binary files etc
+                    pass
+
+        if match_count == 0:
+            search_results["summary"] = (
+                f"No matches found for pattern '{regex}' in path '{path}' "
+                f"(searched {searched_file_count} files)."
+            )
+        else:
+            search_results["summary"] = (
+                f"Found {match_count} matches in {file_match_count} files "
+                f"(searched {searched_file_count} files)."
+            )
+        return search_results
+
+    except Exception as e:
+        return {"error": f"Error searching files: {e}"}
+
+
+def _get_file_matches(
+    file_path: str,
+    pattern: re.Pattern,
+    context_lines: int = 2,
+) -> list[dict[str, any]]:
+    """Search for regex matches in a file with context."""
+    with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
+        lines = f.readlines()
+    matches = []
+    for line_idx, line in enumerate(lines):
+        if pattern.search(line):
+            line_num = line_idx + 1
+            context_start = max(0, line_idx - context_lines)
+            context_end = min(len(lines), line_idx + context_lines + 1)
+            match_data = {
+                "line_number": line_num,
+                "line_content": line.rstrip(),
+                "context_before": [
+                    lines[j].rstrip() for j in range(context_start, line_idx)
+                ],
+                "context_after": [
+                    lines[j].rstrip() for j in range(line_idx + 1, context_end)
+                ],
+            }
+            matches.append(match_data)
+    return matches
+
+
+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 analyze_file(path: str, query: str) -> str:
+    """
+    Delegates deep analysis of a specific file to a specialized sub-agent.
+
+    **WHEN TO USE:**
+    - For complex questions about a file's logic, structure, or potential bugs.
+    - When you need a summary or specific details that require "understanding" the code.
+
+    **NOTE:** For simple data retrieval, use `read_file`.
+
+    **ARGS:**
+    - `path`: Path to the file to analyze.
+    - `query`: The specific analytical question or instruction.
+    """
+    # Lazy imports to avoid circular dependencies
+    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_file_extractor_system_prompt
+
+    abs_path = os.path.abspath(os.path.expanduser(path))
+    if not os.path.exists(abs_path):
+        return f"Error: File not found: {path}"
+
+    # Read content
+    content = read_file(abs_path)
+    if content.startswith("Error:"):
+        return content
+
+    # Check token limit and truncate if necessary
+    token_threshold = CFG.LLM_FILE_ANALYSIS_TOKEN_THRESHOLD
+    # Simple character-based approximation (1 token ~ 4 chars)
+    char_limit = token_threshold * 4
+
+    clipped_content = content
+    if len(content) > char_limit:
+        clipped_content = content[:char_limit] + "\n...[TRUNCATED]..."
+
+    system_prompt = get_file_extractor_system_prompt()
+
+    # Create the sub-agent
+    agent = create_agent(
+        model=llm_config.model,
+        system_prompt=system_prompt,
+        tools=[
+            read_file,
+            search_files,
+        ],
+    )
+
+    # Construct the user message
+    user_message = f"""
+    Instruction: {query}
+    File Path: {abs_path}
+    File Content:
+    ```
+    {clipped_content}
+    ```
+    """
+
+    # Run the agent
+    # We pass empty history as this is a fresh sub-task
+    # We use print as the print_fn (which streams to stdout)
+    result, _ = await run_agent(
+        agent=agent,
+        message=user_message,
+        message_history=[],
+        limiter=llm_limiter,
+    )
+
+    return str(result)