zrb 1.13.1__py3-none-any.whl → 1.21.33__py3-none-any.whl

zrb/builtin/llm/tool/cli.py

@@ -1,33 +1,68 @@
-import json
 import subprocess
+import sys
 
+if sys.version_info >= (3, 12):
+    from typing import TypedDict
+else:
+    from typing_extensions import TypedDict
 
-def run_shell_command(command: str) -> str:
+
+class ShellCommandResult(TypedDict):
+    """
+    Result of shell command execution
+
+    Attributes:
+        return_code: The return code, 0 indicating no error
+        stdout: Standard output
+        stderr: Standard error
     """
-    Executes a shell command on the user's local machine and returns the output.
 
-    This tool is powerful and should be used for tasks that require interacting with the command line, such as running scripts, managing system processes, or using command-line tools.
+    return_code: int
+    stdout: str
+    stderr: str
 
-    **Security Warning:** This tool executes commands with the same permissions as the user running the assistant. Before executing any command that could modify files or system state (e.g., `git`, `npm`, `pip`, `docker`), you MUST explain what the command does and ask the user for confirmation.
+
+def run_shell_command(command: str, timeout: int = 30) -> ShellCommandResult:
+    """
+    Executes a non-interactive shell command on the user's machine.
+
+    **EFFICIENCY TIP:**
+    Combine multiple shell commands into a single call using `&&` or `;` to save steps.
+    Example: `mkdir new_dir && cd new_dir && touch file.txt`
+
+    CRITICAL: This tool runs with user-level permissions. Explain commands that modify
+        the system (e.g., `git`, `pip`) and ask for confirmation.
+    IMPORTANT: Long-running processes should be run in the background (e.g., `command &`).
+
+    Example:
+    run_shell_command(command='ls -l', timeout=30)
 
     Args:
-        command (str): The exact shell command to execute.
+        command (str): The exact shell command to be executed.
+        timeout (int): The maximum time in seconds to wait for the command to finish.
+            Defaults to 30.
 
     Returns:
-        str: A JSON string containing return code, standard output (stdout),
-            and standard error (stderr) from the command.
-            Example: {"return_code": 0, "stdout": "ok", "stderr": ""}
+        dict: return_code, stdout, and stderr.
     """
-    result = subprocess.run(
-        command,
-        shell=True,
-        capture_output=True,
-        text=True,
-    )
-    return json.dumps(
-        {
-            "return_code": result.returncode,
-            "stdout": result.stdout,
-            "stderr": result.stderr,
+    try:
+        result = subprocess.run(
+            command,
+            shell=True,
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+        )
+        return {
+            "return_code": int(result.returncode),
+            "stdout": str(result.stdout or ""),
+            "stderr": str(result.stderr or ""),
+        }
+    except subprocess.TimeoutExpired as e:
+        stdout = e.stdout.decode() if isinstance(e.stdout, bytes) else (e.stdout or "")
+        stderr = e.stderr.decode() if isinstance(e.stderr, bytes) else (e.stderr or "")
+        return {
+            "return_code": 124,
+            "stdout": str(stdout),
+            "stderr": f"{stderr}\nError: Command timed out after {timeout} seconds".strip(),
         }
-    )

zrb/builtin/llm/tool/code.py

@@ -6,6 +6,7 @@ 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.util.cli.style import stylize_faint
 
 _DEFAULT_EXTENSIONS = [
     "py",
@@ -49,36 +50,41 @@ _DEFAULT_EXTENSIONS = [
 async def analyze_repo(
     ctx: AnyContext,
     path: str,
-    goal: str,
+    query: str,
     extensions: list[str] = _DEFAULT_EXTENSIONS,
     exclude_patterns: list[str] = DEFAULT_EXCLUDED_PATTERNS,
     extraction_token_threshold: int | None = None,
     summarization_token_threshold: int | None = None,
 ) -> str:
     """
-    Performs a deep, goal-oriented analysis of a code repository or directory.
+    Analyzes a code repository or directory to answer a specific query.
 
-    This powerful tool recursively reads all relevant files in a directory, extracts key information, and then summarizes that information in relation to a specific goal. It uses intelligent sub-agents for extraction and summarization, making it ideal for complex tasks that require a holistic understanding of a codebase.
+    CRITICAL: The query must contain ALL necessary context, instructions, and information.
+        The sub-agent performing the analysis does NOT share your current conversation
+        history, memory, or global context.
+        The quality of analysis depends entirely on the query. Vague queries yield poor
+        results.
 
-    Use this tool for:
-    - Understanding a large or unfamiliar codebase.
-    - Generating high-level summaries of a project's architecture.
-    - Performing a preliminary code review.
-    - Creating documentation or diagrams (e.g., "Generate a Mermaid C4 diagram for this service").
-    - Answering broad questions like "How does the authentication in this project work?".
+    IMPORTANT: This tool can be slow and expensive on large repositories. Use judiciously.
+
+    Example:
+    analyze_repo(
+        path='src/my_project',
+        query='Summarize the main functionalities by analyzing Python files.',
+        extensions=['py']
+    )
 
     Args:
-        path (str): The path to the directory or repository to analyze.
-        goal (str): A clear and specific description of what you want to achieve. A good goal is critical for getting a useful result. For example: "Understand the database schema by analyzing all the .sql files" or "Create a summary of all the API endpoints defined in the 'api' directory".
-        extensions (list[str], optional): A list of file extensions to include in the analysis. Defaults to a comprehensive list of common code and configuration files.
-        exclude_patterns (list[str], optional): A list of glob patterns for files and directories to exclude from the analysis. Defaults to common patterns like '.git', 'node_modules', and '.venv'.
-        extraction_token_threshold (int, optional): The maximum token threshold for the extraction sub-agent.
-        summarization_token_threshold (int, optional): The maximum token threshold for the summarization sub-agent.
+        ctx (AnyContext): The execution context.
+        path (str): Path to the directory or repository.
+        query (str): Clear and specific analysis question or goal.
+        extensions (list[str], optional): File extensions to include.
+        exclude_patterns (list[str], optional): Glob patterns to exclude.
+        extraction_token_threshold (int, optional): Token limit for extraction sub-agent.
+        summarization_token_threshold (int, optional): Token limit for summarization sub-agent.
 
     Returns:
-        str: A detailed, markdown-formatted analysis and summary of the repository, tailored to the specified goal.
-    Raises:
-        Exception: If an error occurs during the analysis.
+        str: Detailed, markdown-formatted analysis and summary.
     """
     if extraction_token_threshold is None:
         extraction_token_threshold = CFG.LLM_REPO_ANALYSIS_EXTRACTION_TOKEN_THRESHOLD
@@ -88,23 +94,26 @@ async def analyze_repo(
         )
     abs_path = os.path.abspath(os.path.expanduser(path))
     file_metadatas = _get_file_metadatas(abs_path, extensions, exclude_patterns)
-    ctx.print("Extraction")
+    ctx.print(stylize_faint("  📝 Extraction"), plain=True)
     extracted_infos = await _extract_info(
         ctx,
         file_metadatas=file_metadatas,
-        goal=goal,
+        query=query,
         token_limit=extraction_token_threshold,
     )
+    if len(extracted_infos) == 0:
+        raise RuntimeError(
+            "No info can be extracted, adjust extensions or exclude_patterns."
+        )
     if len(extracted_infos) == 1:
         return extracted_infos[0]
-    ctx.print("Summarization")
     summarized_infos = extracted_infos
     while len(summarized_infos) > 1:
-        ctx.print("Summarization")
+        ctx.print(stylize_faint("  📝 Summarization"), plain=True)
         summarized_infos = await _summarize_info(
             ctx,
             extracted_infos=summarized_infos,
-            goal=goal,
+            query=query,
             token_limit=summarization_token_threshold,
         )
     return summarized_infos[0]
@@ -122,11 +131,11 @@ def _get_file_metadatas(
             if not any(file.endswith(f".{ext}") for ext in extensions):
                 continue
             file_path = os.path.join(root, file)
-            if is_excluded(file_path, exclude_patterns):
-                continue
             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:
-                    rel_path = os.path.relpath(file_path, dir_path)
                     metadata_list.append({"path": rel_path, "content": f.read()})
             except Exception as e:
                 print(f"Error reading file {file_path}: {e}")
@@ -137,13 +146,15 @@ def _get_file_metadatas(
 async def _extract_info(
     ctx: AnyContext,
     file_metadatas: list[dict[str, str]],
-    goal: str,
+    query: str,
     token_limit: int,
 ) -> list[str]:
     extract = create_sub_agent_tool(
         tool_name="extract",
         tool_description="extract",
         system_prompt=CFG.LLM_REPO_EXTRACTOR_SYSTEM_PROMPT,
+        auto_summarize=False,
+        remember_history=False,
     )
     extracted_infos = []
     content_buffer = []
@@ -155,7 +166,7 @@ async def _extract_info(
         file_str = json.dumps(file_obj)
         if current_token_count + llm_rate_limitter.count_token(file_str) > token_limit:
             if content_buffer:
-                prompt = _create_extract_info_prompt(goal, content_buffer)
+                prompt = json.dumps(_create_extract_info_prompt(query, content_buffer))
                 extracted_info = await extract(
                     ctx, llm_rate_limitter.clip_prompt(prompt, token_limit)
                 )
@@ -165,10 +176,9 @@ async def _extract_info(
         else:
             content_buffer.append(file_obj)
             current_token_count += llm_rate_limitter.count_token(file_str)
-
     # Process any remaining content in the buffer
     if content_buffer:
-        prompt = _create_extract_info_prompt(goal, content_buffer)
+        prompt = json.dumps(_create_extract_info_prompt(query, content_buffer))
         extracted_info = await extract(
             ctx, llm_rate_limitter.clip_prompt(prompt, token_limit)
         )
@@ -176,25 +186,25 @@ async def _extract_info(
     return extracted_infos
 
 
-def _create_extract_info_prompt(goal: str, content_buffer: list[dict]) -> str:
-    return json.dumps(
-        {
-            "main_assistant_goal": goal,
-            "files": content_buffer,
-        }
-    )
+def _create_extract_info_prompt(query: str, content_buffer: list[dict]) -> dict:
+    return {
+        "main_assistant_query": query,
+        "files": content_buffer,
+    }
 
 
 async def _summarize_info(
     ctx: AnyContext,
     extracted_infos: list[str],
-    goal: str,
+    query: str,
     token_limit: int,
 ) -> list[str]:
     summarize = create_sub_agent_tool(
         tool_name="extract",
         tool_description="extract",
         system_prompt=CFG.LLM_REPO_SUMMARIZER_SYSTEM_PROMPT,
+        auto_summarize=False,
+        remember_history=False,
     )
     summarized_infos = []
     content_buffer = ""
@@ -202,7 +212,9 @@ async def _summarize_info(
         new_prompt = content_buffer + extracted_info
         if llm_rate_limitter.count_token(new_prompt) > token_limit:
             if content_buffer:
-                prompt = _create_summarize_info_prompt(goal, content_buffer)
+                prompt = json.dumps(
+                    _create_summarize_info_prompt(query, content_buffer)
+                )
                 summarized_info = await summarize(
                     ctx, llm_rate_limitter.clip_prompt(prompt, token_limit)
                 )
@@ -213,7 +225,7 @@ async def _summarize_info(
 
     # Process any remaining content in the buffer
     if content_buffer:
-        prompt = _create_summarize_info_prompt(goal, content_buffer)
+        prompt = json.dumps(_create_summarize_info_prompt(query, content_buffer))
         summarized_info = await summarize(
             ctx, llm_rate_limitter.clip_prompt(prompt, token_limit)
         )
@@ -221,10 +233,8 @@ async def _summarize_info(
     return summarized_infos
 
 
-def _create_summarize_info_prompt(goal: str, content_buffer: str) -> str:
-    return json.dumps(
-        {
-            "main_assistant_goal": goal,
-            "extracted_info": content_buffer,
-        }
-    )
+def _create_summarize_info_prompt(query: str, content_buffer: str) -> dict:
+    return {
+        "main_assistant_query": query,
+        "extracted_info": content_buffer,
+    }