code-lm 0.1.0__py3-none-any.whl

gemini_cli/tools/system_tools.py

@@ -0,0 +1,65 @@
+"""
+System operation tools.
+"""
+
+import os
+import subprocess
+import tempfile
+from .base import BaseTool
+
+class BashTool(BaseTool):
+    """Tool to execute bash commands."""
+    
+    name = "bash"
+    description = "Execute a bash command"
+    
+    # List of banned commands for security
+    BANNED_COMMANDS = [
+        'curl', 'wget', 'nc', 'netcat', 'telnet',
+        'lynx', 'w3m', 'links', 'ssh',
+    ]
+    
+    def execute(self, command, timeout=30000):
+        """
+        Execute a bash command.
+        
+        Args:
+            command: The command to execute
+            timeout: Timeout in milliseconds (optional)
+        """
+        try:
+            # Check for banned commands
+            for banned in self.BANNED_COMMANDS:
+                if banned in command.split():
+                    return f"Error: The command '{banned}' is not allowed for security reasons."
+            
+            # Convert timeout to seconds (with better error handling)
+            try:
+                timeout_sec = int(timeout) / 1000
+            except ValueError:
+                # If timeout can't be converted to int, use default
+                timeout_sec = 30
+            
+            # Remove the temporary directory context
+            process = subprocess.Popen(
+                command,
+                shell=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.PIPE,
+                text=True,
+            )
+            
+            try:
+                stdout, stderr = process.communicate(timeout=timeout_sec)
+                
+                if process.returncode != 0:
+                    return f"Command exited with status {process.returncode}\n\nSTDOUT:\n{stdout}\n\nSTDERR:\n{stderr}"
+                
+                return stdout
+            
+            except subprocess.TimeoutExpired:
+                process.kill()
+                return f"Error: Command timed out after {timeout_sec} seconds"
+        
+        except Exception as e:
+            return f"Error executing command: {str(e)}"

gemini_cli/tools/task_complete_tool.py

@@ -0,0 +1,53 @@
+"""
+Tool to signal task completion.
+"""
+import logging
+from .base import BaseTool
+
+log = logging.getLogger(__name__)
+
+class TaskCompleteTool(BaseTool):
+    """
+    Signals that the current task/request is fully completed.
+    This MUST be the final tool called by the assistant for a given request.
+    """
+    name = "task_complete"
+    description = "Signals task completion. MUST be called as the final step, providing a user-friendly summary."
+
+    def execute(self, summary: str) -> str:
+        """
+        Signals completion and returns the summary provided by the LLM.
+
+        Args:
+            summary: A concise, user-friendly summary of the actions taken and the final outcome.
+
+        Returns:
+            The summary string provided as input. The orchestrator uses this call as a signal to stop.
+        """
+        log.info(f"Task completion signaled by LLM.")
+        
+        # --- ADDED/MODIFIED: More Robust Cleaning --- 
+        cleaned_summary = summary
+        if isinstance(summary, str):
+            log.debug(f"Original summary from LLM (Type: {type(summary)}, Length: {len(summary)}): \"{summary}\"")
+            # Repeatedly strip common leading/trailing chars like quotes, spaces
+            strippable_chars = ' "\'\n\t'
+            while cleaned_summary.startswith(tuple(strippable_chars)) or cleaned_summary.endswith(tuple(strippable_chars)):
+                prev_summary = cleaned_summary
+                cleaned_summary = cleaned_summary.strip(strippable_chars)
+                if cleaned_summary == prev_summary: # Avoid infinite loop if strip doesn't change anything
+                    break 
+            log.debug(f"Final cleaned summary: \"{cleaned_summary}\"")
+        else:
+             log.warning(f"TaskCompleteTool received non-string summary type: {type(summary)}")
+             cleaned_summary = str(summary).strip() # Attempt to convert and strip
+        # --- END ADDED/MODIFIED SECTION ---
+        
+        log.debug(f"Processing summary (cleaned): {cleaned_summary}")
+        if not cleaned_summary or len(cleaned_summary) < 5:
+             log.warning("TaskCompleteTool called with missing or very short summary.")
+             # Provide a default confirmation if summary is bad/missing
+             return "Task marked as complete, but the provided summary was insufficient."
+        # The orchestrator loop will see this tool was called and use this summary.
+        # We just return the summary itself.
+        return cleaned_summary

gemini_cli/tools/test_runner.py

@@ -0,0 +1,99 @@
+"""
+Tool for running automated tests (e.g., pytest).
+"""
+
+import subprocess
+import logging
+import shlex
+from .base import BaseTool
+
+# Configure logging for this tool
+log = logging.getLogger(__name__)
+
+class TestRunnerTool(BaseTool):
+    """
+    Tool to execute automated tests using a test runner like pytest.
+    Assumes the test runner command (e.g., 'pytest') is available
+    in the environment where the CLI is run.
+    """
+    name = "test_runner"
+    description = "Runs automated tests using the project's test runner (defaults to trying 'pytest'). Use after making code changes to verify correctness."
+
+    def execute(self, test_path: str | None = None, options: str | None = None, runner_command: str = "pytest") -> str:
+        """
+        Executes automated tests.
+
+        Args:
+            test_path: Specific file or directory to test (optional, runs tests discovered by the runner if omitted).
+            options: Additional command-line options for the test runner (e.g., '-k my_test', '-v', '--cov'). Optional.
+            runner_command: The command to invoke the test runner (default: 'pytest').
+
+        Returns:
+            A string summarizing the test results, including output on failure.
+        """
+        command = [runner_command]
+
+        if options:
+            # Use shlex to safely split options string respecting quotes
+            try:
+                command.extend(shlex.split(options))
+            except ValueError as e:
+                log.warning(f"Could not parse options string '{options}': {e}. Ignoring options.")
+                # Optionally return an error message here
+                # return f"Error: Invalid options string provided: {options}"
+
+        if test_path:
+            command.append(test_path)
+
+        log.info(f"Executing test command: {' '.join(command)}")
+
+        try:
+            # Execute the command, capture output, set a timeout (e.g., 5 minutes)
+            process = subprocess.run(
+                command,
+                capture_output=True,
+                text=True,
+                check=False, # Don't raise exception on non-zero exit code, we'll check it manually
+                timeout=300 # Timeout in seconds (e.g., 5 minutes)
+            )
+
+            exit_code = process.returncode
+            stdout = process.stdout.strip()
+            stderr = process.stderr.strip()
+
+            log.info(f"Test run completed. Exit Code: {exit_code}")
+            log.debug(f"Test stdout:\n{stdout}")
+            if stderr:
+                log.debug(f"Test stderr:\n{stderr}")
+
+            # Prepare summary message
+            summary = f"Test run using '{runner_command}' completed.\n"
+            summary += f"Exit Code: {exit_code}\n"
+
+            if exit_code == 0:
+                summary += "Status: SUCCESS\n"
+                # Include stdout even on success, maybe truncated?
+                summary += f"\nOutput:\n---\n{stdout[-1000:]}\n---\n" # Show last 1000 chars
+            else:
+                summary += "Status: FAILED\n"
+                summary += f"\nStandard Output:\n---\n{stdout}\n---\n"
+                if stderr:
+                    summary += f"\nStandard Error:\n---\n{stderr}\n---\n"
+
+            # Specific exit codes for pytest might be useful
+            # (e.g., 5 means no tests collected) - can add more logic here
+            if exit_code == 5 and 'pytest' in runner_command:
+                 summary += "\nNote: Pytest exit code 5 often means no tests were found or collected."
+
+
+            return summary
+
+        except FileNotFoundError:
+            log.error(f"Test runner command '{runner_command}' not found.")
+            return f"Error: Test runner command '{runner_command}' not found. Is it installed and in PATH?"
+        except subprocess.TimeoutExpired:
+            log.error("Test run timed out.")
+            return "Error: Test run exceeded the timeout limit (5 minutes)."
+        except Exception as e:
+            log.error(f"An unexpected error occurred while running tests: {e}", exc_info=True)
+            return f"Error: An unexpected error occurred: {str(e)}"

gemini_cli/tools/tree_tool.py

@@ -0,0 +1,92 @@
+"""
+Tool for displaying directory structure using the 'tree' command.
+"""
+import subprocess
+import logging
+from google.generativeai.types import FunctionDeclaration, Tool
+
+from .base import BaseTool
+
+log = logging.getLogger(__name__)
+
+DEFAULT_TREE_DEPTH = 3
+MAX_TREE_DEPTH = 10
+
+class TreeTool(BaseTool):
+    name: str = "tree"
+    description: str = (
+        f"""Displays the directory structure as a tree. Shows directories and files.
+        Use this to understand the hierarchy and layout of the current working directory or a subdirectory.
+        Defaults to a depth of {DEFAULT_TREE_DEPTH}. Use the 'depth' argument to specify a different level.
+        Optionally specify a 'path' to view a subdirectory instead of the current directory."""
+    )
+    args_schema: dict = {
+         "path": {
+            "type": "string",
+            "description": "Optional path to a specific directory relative to the workspace root. If omitted, uses the current directory.",
+        },
+        "depth": {
+            "type": "integer",
+            "description": f"Optional maximum display depth of the directory tree (Default: {DEFAULT_TREE_DEPTH}, Max: {MAX_TREE_DEPTH}).",
+        },
+    }
+    # Optional args: path, depth
+    required_args: list[str] = []
+
+    def execute(self, path: str | None = None, depth: int | None = None) -> str:
+        """Executes the tree command."""
+        
+        if depth is None:
+            depth_limit = DEFAULT_TREE_DEPTH
+        else:
+            # Clamp depth to be within reasonable limits
+            depth_limit = max(1, min(depth, MAX_TREE_DEPTH))
+            
+        command = ['tree', f'-L {depth_limit}']
+        
+        # Add path if specified
+        target_path = "." # Default to current directory
+        if path:
+            # Basic path validation/sanitization might be needed depending on security context
+            target_path = path
+            command.append(target_path)
+
+        log.info(f"Executing tree command: {' '.join(command)}")
+        try:
+            # Adding '-a' might be useful to show hidden files, but could be verbose.
+            # Adding '-F' appends / to dirs, * to executables, etc.
+            # Using shell=True is generally discouraged, but might be needed if tree isn't directly in PATH
+            # or if handling complex paths. Sticking to list format for now.
+            process = subprocess.run(
+                command, 
+                capture_output=True, 
+                text=True, 
+                check=False, # Don't raise exception on non-zero exit code
+                timeout=15 # Add a timeout
+            )
+
+            if process.returncode == 0:
+                log.info(f"Tree command successful for path '{target_path}' with depth {depth_limit}.")
+                # Limit output size? Tree can be huge.
+                output = process.stdout.strip()
+                if len(output.splitlines()) > 200: # Limit lines as a proxy for size
+                     log.warning(f"Tree output for '{target_path}' exceeded 200 lines. Truncating.")
+                     output = "\n".join(output.splitlines()[:200]) + "\n... (output truncated)"
+                return output
+            elif process.returncode == 127 or "command not found" in process.stderr.lower():
+                 log.error(f"\'tree\' command not found. It might not be installed.")
+                 return "Error: 'tree' command not found. Please ensure it is installed and in the system's PATH."
+            else:
+                log.error(f"Tree command failed with return code {process.returncode}. Path: '{target_path}', Depth: {depth_limit}. Stderr: {process.stderr.strip()}")
+                error_detail = process.stderr.strip() if process.stderr else "(No stderr)"
+                return f"Error executing tree command (Code: {process.returncode}): {error_detail}"
+
+        except FileNotFoundError:
+             log.error(f"\'tree\' command not found (FileNotFoundError). It might not be installed.")
+             return "Error: 'tree' command not found. Please ensure it is installed and in the system's PATH."
+        except subprocess.TimeoutExpired:
+             log.error(f"Tree command timed out for path '{target_path}' after 15 seconds.")
+             return f"Error: Tree command timed out for path '{target_path}'. The directory might be too large or complex."
+        except Exception as e:
+            log.exception(f"An unexpected error occurred while executing tree command for path '{target_path}': {e}")
+            return f"An unexpected error occurred while executing tree: {str(e)}" 

gemini_cli/utils.py

@@ -0,0 +1,20 @@
+"""
+Utility functions for the Gemini CLI tool.
+"""
+
+import tiktoken
+import json
+
+def count_tokens(text):
+    """
+    Count the number of tokens in a text string.
+    
+    This is a rough estimate for Gemini 2.5 Pro, using GPT-4 tokenizer as a proxy.
+    For production, you'd want to use model-specific token counting.
+    """
+    try:
+        encoding = tiktoken.encoding_for_model("gpt-4")
+        return len(encoding.encode(text))
+    except Exception:
+        # Fallback method: roughly 4 chars per token
+        return len(text) // 4