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

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)

zrb/llm/tool/note.py

@@ -0,0 +1,70 @@
+import os
+from typing import Callable, List
+
+from zrb.llm.note.manager import NoteManager
+
+
+def create_note_tools(note_manager: NoteManager) -> List[Callable]:
+    async def read_long_term_note() -> str:
+        """
+        Retrieves your GLOBAL 🧠 Long-Term Memory.
+        This contains established preferences, personal facts, and context spanning multiple projects.
+        ALWAYS check this at the start of a session.
+        """
+        return note_manager.read("~")
+
+    read_long_term_note.__name__ = "read_long_term_note"
+
+    async def write_long_term_note(content: str) -> str:
+        """
+        Updates your GLOBAL 🧠 Long-Term Memory with CRITICAL information.
+        Use this to persist user preferences, personal facts, and cross-project rules.
+
+        **WARNING:** This COMPLETELY OVERWRITES the existing Long-Term Note.
+
+        **ARGS:**
+        - `content`: The full text to store in the global memory.
+        """
+        note_manager.write("~", content)
+        return "Global long-term note saved."
+
+    write_long_term_note.__name__ = "write_long_term_note"
+
+    async def read_contextual_note(path: str | None = None) -> str:
+        """
+        Retrieves LOCAL 📝 Contextual Notes for a specific project or directory.
+        Use this to recall architectural decisions or project-specific guidelines.
+
+        **ARGS:**
+        - `path`: Target file/dir path. Defaults to current working directory.
+        """
+        if path is None:
+            path = os.getcwd()
+        return note_manager.read(path)
+
+    read_contextual_note.__name__ = "read_contextual_note"
+
+    async def write_contextual_note(content: str, path: str | None = None) -> str:
+        """
+        Persists LOCAL 📝 Contextual Notes for a specific project or directory.
+        Use this to save architectural patterns or progress markers for the current task.
+
+        **WARNING:** This COMPLETELY OVERWRITES the contextual note for the specified path.
+
+        **ARGS:**
+        - `content`: The full text to store in the local memory.
+        - `path`: Target file/dir path. Defaults to current working directory.
+        """
+        if path is None:
+            path = os.getcwd()
+        note_manager.write(path, content)
+        return f"Contextual note saved for: {path}"
+
+    write_contextual_note.__name__ = "write_contextual_note"
+
+    return [
+        read_long_term_note,
+        write_long_term_note,
+        read_contextual_note,
+        write_contextual_note,
+    ]

zrb/{builtin/llm → llm}/tool/rag.py

@@ -5,6 +5,7 @@ import os
 import sys
 from collections.abc import Callable
 from textwrap import dedent
+from typing import Any
 
 import ulid
 
@@ -44,49 +45,40 @@ def create_rag_from_directory(
     openai_embedding_model: str | None = None,
 ):
     """
-    Creates a powerful RAG tool for querying a local knowledge base.
-
-    This factory function generates a tool that can perform semantic searches
-    over a directory of documents. It automatically indexes the documents into
-    a vector database, keeping it updated as files change. The generated tool
-    is ideal for answering questions based on a specific set of documents,
-    such as project documentation, research papers, or internal wikis.
-
-    The created tool will:
-    1.  Monitor a specified directory for file changes.
-    2.  Automatically update a vector database (ChromaDB) with the latest
-        content.
-    3.  Accept a user query, embed it, and perform a similarity search against
-        the document vectors.
-    4.  Return the most relevant document chunks that match the query.
+    Create a powerful RAG (Retrieval-Augmented Generation) tool for querying a local
+    knowledge base.
+
+    This factory function generates a tool that performs semantic search over a directory of
+    documents. It automatically indexes the documents into a vector database (ChromaDB) and
+    keeps it updated as files change.
+
+    The generated tool is ideal for answering questions based on a specific set of documents,
+    such as project documentation or internal wikis.
 
     Args:
-        tool_name (str): The name for the generated RAG tool (e.g.,
-            "search_project_docs").
-        tool_description (str): A clear description of what the generated tool
-            does and when to use it (e.g., "Searches the project's technical
-            documentation to answer questions.").
-        document_dir_path (str, optional): The path to the directory
-            containing the documents to be indexed.
-        vector_db_path (str, optional): The path to store the ChromaDB vector
+        tool_name (str): The name for the generated RAG tool (e.g., "search_project_docs").
+        tool_description (str): A clear description of what the tool does and when to use it.
+            This is what the LLM will see.
+        document_dir_path (str, optional): The path to the directory containing the documents
+            to be indexed.
+        vector_db_path (str, optional): The path where the ChromaDB vector database will be
+            stored.
+        vector_db_collection (str, optional): The name of the collection within the vector
             database.
-        vector_db_collection (str, optional): The name of the collection within
-            the vector database.
         chunk_size (int, optional): The size of text chunks for embedding.
         overlap (int, optional): The overlap between text chunks.
-        max_result_count (int, optional): The maximum number of search results
-            to return.
-        file_reader (list[RAGFileReader], optional): Custom file readers for
+        max_result_count (int, optional): The maximum number of search results to return.
+        file_reader (list[RAGFileReader], optional): A list of custom file readers for
             specific file types.
-        openai_api_key (str, optional): OpenAI API key for embeddings.
-        openai_base_url (str, optional): OpenAI base URL for embeddings.
+        openai_api_key (str, optional): Your OpenAI API key for generating embeddings.
+        openai_base_url (str, optional): An optional base URL for the OpenAI API.
         openai_embedding_model (str, optional): The embedding model to use.
 
     Returns:
-        Callable: An asynchronous function that serves as the RAG tool.
+        An asynchronous function that serves as the RAG tool.
     """
 
-    async def retrieve(query: str) -> str:
+    async def retrieve(query: str) -> dict[str, Any]:
         # Docstring will be set dynamically below
         from chromadb import PersistentClient
         from chromadb.config import Settings
@@ -201,16 +193,20 @@ def create_rag_from_directory(
             query_embeddings=query_vector,
             n_results=max_result_count_val,
         )
-        return json.dumps(results)
+        return dict(results)
 
     retrieve.__name__ = tool_name
     retrieve.__doc__ = dedent(
         f"""
         {tool_description}
-        Args:
-            query (str): The user query to search for in documents.
-        Returns:
-            str: JSON string with search results: {{"ids": [...], "documents": [...], ...}}
+        This tool performs a semantic search across a curated knowledge base of documents.
+        It is highly effective for answering questions that require specific project knowledge not found in general training data.
+
+        **ARGS:**
+        - `query` (str): The semantic search query or question.
+
+        **RETURNS:**
+        - A dictionary containing matching document chunks ("documents") and their metadata.
         """
     ).strip()
     return retrieve

zrb/llm/tool/search/brave.py

@@ -0,0 +1,53 @@
+from typing import Any
+
+import requests
+
+from zrb.config.config import CFG
+
+
+def search_internet(
+    query: str,
+    page: int = 1,
+    safe_search: str | None = None,
+    language: str | None = None,
+) -> dict[str, Any]:
+    """
+    Performs a live internet search using Brave Search to retrieve up-to-date information, news, or documentation.
+
+    **WHEN TO USE:**
+    - To find the latest information on rapidly changing topics (e.g., library updates, current events).
+    - To search for documentation or examples not present in the local codebase.
+    - To verify facts or find external resources.
+
+    **ARGS:**
+    - `query`: The search string or question.
+    - `page`: Result page number (default 1).
+    """
+    if safe_search is None:
+        safe_search = CFG.BRAVE_API_SAFE
+    if language is None:
+        language = CFG.BRAVE_API_LANG
+
+    user_agent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36"
+
+    response = requests.get(
+        "https://api.search.brave.com/res/v1/web/search",
+        headers={
+            "User-Agent": user_agent,
+            "Accept": "application/json",
+            "x-subscription-token": CFG.BRAVE_API_KEY,
+        },
+        params={
+            "q": query,
+            "count": "10",
+            "offset": (page - 1) * 10,
+            "safesearch": safe_search,
+            "search_lang": language,
+            "summary": "true",
+        },
+    )
+    if response.status_code != 200:
+        raise Exception(
+            f"Error: Unable to retrieve search results (status code: {response.status_code})"
+        )
+    return response.json()