✨ Feat: Enhance Python tool and add logging

- Add logging to the agent to track model responses and tool execution.
- Introduce a `_truncate` helper for logging long strings.
- Enhance the Python tool to:
    - Accept a list of extra allowed imports via configuration.
    - Use AST to parse code and enforce import restrictions.
    - Add `psutil` to default allowed imports and requirements.
- Update `load_default_tools` to accept settings for Python tool imports.
- Add verbose and quiet logging options to the CLI.
- Update README with examples and configuration details for Python tool imports.

Author: PythonicNinja

.env.example

@@ -5,3 +5,4 @@ GEMINI_API_KEY=your-gemini-key
 GEMINI_MODEL=gemini-2.5-flash
 AGENT_SYSTEM_PROMPT=You are a concise assistant. Use tools when needed.
 REQUEST_TIMEOUT=30
+PYTHON_TOOL_IMPORTS=os,sys,psutil

README.md

@@ -24,6 +24,27 @@ The repo ships with a lightweight `Makefile` wired to [`uvx`](https://docs.astra
    make run PROMPT="Summarize the latest message."
    ```
 
+### Examples
+
+#### Disk space:
+
+```bash
+| => python main.py "get disk space on current machine using python"
+Disk Space:
+Total: 926.35 GB
+Used: 452.60 GB
+Free: 473.76 GB
+```
+
+#### CPU usage:
+
+```bash
+| => python main.py "give current cpu usage" -v
+INFO: Running tool 'python'.
+INFO: Responding without tool use.
+Current CPU usage: 16.5%
+```
+
 ### Configuration
 
 All settings live in `.env` (loaded with `python-dotenv`):
@@ -37,6 +58,7 @@ All settings live in `.env` (loaded with `python-dotenv`):
 | `GEMINI_MODEL` | Defaults to `gemini-2.5-flash`.            |
 | `AGENT_SYSTEM_PROMPT` | Optional custom system prompt.             |
 | `REQUEST_TIMEOUT` | Request timeout in seconds (default `30`). |
+| `PYTHON_TOOL_IMPORTS` | Optional comma list of extra python-tool imports (`os,sys,psutil`). |
 
 The repository already contains `.env.example` with placeholders for these values.
 
@@ -51,6 +73,8 @@ python main.py --help
 - `--max-turns`: maximum number of tool iterations.
 - `--no-tools`: disable tool use.
 - `--list-tools`: inspect available tools.
+- `-v/--verbose`: increase logging (use `-vv` for debug-level traces about tool usage).
+- `-q/--quiet`: suppress logs (errors only).
 
 ### Tools
 
@@ -59,7 +83,9 @@ Tools live under `simple_agent/tools` and implement a tiny interface (`name`, `d
 - `time`: returns the current UTC timestamp.
 - `calculator`: evaluates small arithmetic expressions safely.
 - `file_reader`: dumps a snippet of a local text file (`path[:start-end]`).
-- `python`: runs a short Python snippet in a separate interpreter (stdout/stderr returned).
+- `python`: runs a short Python snippet in a separate interpreter (default imports include `math`, `json`, `os`, `sys`, `psutil`; extend via `PYTHON_TOOL_IMPORTS`).
+
+The python tool executes with a module allowlist. By default it includes: `collections`, `datetime`, `functools`, `itertools`, `json`, `math`, `os`, `pathlib`, `psutil`, `random`, `statistics`, `sys`, `time`. Set `PYTHON_TOOL_IMPORTS` (comma separated) to append additional modules if needed.
 
 Adding new tools only requires dropping a module next to the others and including it in `load_default_tools()`.
 

main.py

@@ -3,27 +3,60 @@
 from __future__ import annotations
 
 import argparse
+import logging
 from dataclasses import replace
 
 from simple_agent import SimpleAgent, get_backend, load_default_tools
 from simple_agent.config import get_settings
 
 
+def configure_logging(verbosity: int, quiet: bool) -> None:
+    if quiet:
+        level = logging.ERROR
+    else:
+        if verbosity >= 2:
+            level = logging.DEBUG
+        elif verbosity == 1:
+            level = logging.INFO
+        else:
+            level = logging.WARNING
+
+    logging.basicConfig(level=level, format="%(levelname)s: %(message)s", force=True)
+
+
 def build_parser() -> argparse.ArgumentParser:
     parser = argparse.ArgumentParser(description="Run a small tool-enabled agent.")
     parser.add_argument("prompt", nargs="?", help="Prompt to send to the agent. If omitted, stdin is used.")
     parser.add_argument("--backend", choices=["chatgpt", "gemini"], help="Override the backend specified in .env.")
     parser.add_argument("--max-turns", type=int, default=5, help="Maximum number of tool loops before giving up.")
     parser.add_argument("--no-tools", action="store_true", help="Disable tool usage and respond directly.")
     parser.add_argument("--list-tools", action="store_true", help="List available tools and exit.")
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        action="count",
+        default=0,
+        help="Increase log verbosity (use -vv for debug).",
+    )
+    parser.add_argument(
+        "-q",
+        "--quiet",
+        action="store_true",
+        help="Quiet mode (errors only). Overrides --verbose.",
+    )
     return parser
 
 
 def main() -> None:
     parser = build_parser()
     args = parser.parse_args()
+    configure_logging(args.verbose, args.quiet)
+
+    settings = get_settings()
+    if args.backend:
+        settings = replace(settings, backend=args.backend)  # type: ignore[arg-type]
 
-    tools = [] if args.no_tools else load_default_tools()
+    tools = [] if args.no_tools else load_default_tools(settings)
 
     if args.list_tools:
         if not tools:
@@ -43,10 +76,6 @@ def main() -> None:
     if not prompt:
         parser.error("A prompt is required.")
 
-    settings = get_settings()
-    if args.backend:
-        settings = replace(settings, backend=args.backend)  # type: ignore[arg-type]
-
     backend = get_backend(settings)
     agent = SimpleAgent(backend=backend, tools=tools, system_prompt=settings.system_prompt)
     try:

pyproject.toml

@@ -8,6 +8,7 @@ requires-python = ">=3.11"
 dependencies = [
     "python-dotenv>=1.0",
     "requests>=2.32",
+    "psutil>=5.9",
 ]
 
 [project.optional-dependencies]

requirements.txt

@@ -1,2 +1,3 @@
 python-dotenv>=1.0
 requests>=2.32
+psutil>=5.9

simple_agent/agent.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import json
+import logging
 import re
 from dataclasses import dataclass, field
 from typing import Dict, Iterable, List
@@ -18,7 +19,11 @@
 If a tool is necessary, respond with ONLY a JSON object that looks like:
 {{"tool": "<tool name>", "input": "<plain text request for the tool>"}}
 Do not wrap the JSON in backticks or add commentary.
-If no tool is needed, answer the user directly in natural language."""
+If no tool is needed, answer the user directly in natural language.
+<IMPORTANT>
+If no explicit tool try to use python tool to execute code in it!
+</IMPORTANT>
+"""
 
 
 @dataclass(slots=True)
@@ -30,6 +35,7 @@ class SimpleAgent:
     system_prompt: str
     tool_map: Dict[str, Tool] = field(init=False)
     _prepared_system_prompt: str = field(init=False)
+    _logger: logging.Logger = field(init=False, repr=False)
 
     def __post_init__(self) -> None:
         self.tool_map: Dict[str, Tool] = {tool.name: tool for tool in self.tools}
@@ -38,6 +44,7 @@ def __post_init__(self) -> None:
             user_prompt=self.system_prompt,
             tool_descriptions=descriptions,
         )
+        self._logger = logging.getLogger(self.__class__.__name__)
 
     def run(self, user_input: str, max_turns: int = 5) -> str:
         history: List[dict[str, str]] = [
@@ -47,15 +54,18 @@ def run(self, user_input: str, max_turns: int = 5) -> str:
 
         for _ in range(max_turns):
             response = self.backend.generate(history)
+            self._logger.debug("Model response: %s", _truncate(response))
             tool_request = self._maybe_extract_tool_request(response)
             if not tool_request:
+                self._logger.info("Responding without tool use.")
                 return response.strip()
 
             tool_name = tool_request.get("tool")
             tool_input = tool_request.get("input", "")
 
             tool = self.tool_map.get(tool_name or "")
             if not tool:
+                self._logger.warning("Model requested unknown tool '%s'.", tool_name)
                 history.append(
                     {
                         "role": "user",
@@ -64,7 +74,11 @@ def run(self, user_input: str, max_turns: int = 5) -> str:
                 )
                 continue
 
+            self._logger.info("Running tool '%s'.", tool_name)
+            if tool_input:
+                self._logger.debug("Tool '%s' input: %s", tool_name, _truncate(tool_input))
             tool_output = tool.run(tool_input)
+            self._logger.debug("Tool '%s' output: %s", tool_name, _truncate(tool_output))
 
             history.append({"role": "assistant", "content": json.dumps(tool_request)})
             history.append(
@@ -93,3 +107,8 @@ def _maybe_extract_tool_request(text: str) -> dict | None:
                 return data
 
         return None
+
+
+def _truncate(value: str, limit: int = 500) -> str:
+    value = value.strip()
+    return value if len(value) <= limit else f"{value[:limit]}…"

simple_agent/config.py

@@ -26,6 +26,7 @@ class Settings:
     gemini_api_key: str | None
     gemini_model: str
     request_timeout: float
+    python_tool_imports: tuple[str, ...]
 
     @staticmethod
     def _get_env(key: str, default: str | None = None) -> str | None:
@@ -51,6 +52,7 @@ def from_env(cls) -> "Settings":
             gemini_api_key=cls._get_env("GEMINI_API_KEY"),
             gemini_model=cls._get_env("GEMINI_MODEL", "gemini-1.5-flash"),
             request_timeout=float(cls._get_env("REQUEST_TIMEOUT", "30")),
+            python_tool_imports=_parse_list(cls._get_env("PYTHON_TOOL_IMPORTS")),
         )
 
 
@@ -59,3 +61,9 @@ def get_settings() -> Settings:
     """Convenience accessor with caching to avoid redundant parsing."""
 
     return Settings.from_env()
+
+
+def _parse_list(value: str | None) -> tuple[str, ...]:
+    if not value:
+        return ()
+    return tuple(item.strip() for item in value.split(",") if item.strip())

simple_agent/tools/__init__.py

@@ -2,19 +2,31 @@
 
 from __future__ import annotations
 
-from typing import List
+from typing import List, TYPE_CHECKING
 
 from .base import Tool
 from .file_read_tool import FileReadTool
 from .math_tool import MathTool
-from .time_tool import TimeTool
 from .python_tool import PythonSandboxTool
+from .time_tool import TimeTool
 
+if TYPE_CHECKING:  # pragma: no cover
+    from ..config import Settings
 
-def load_default_tools() -> List[Tool]:
+
+def load_default_tools(settings: "Settings | None" = None) -> List[Tool]:
     """Return the default toolset used by the CLI."""
 
-    return [TimeTool(), MathTool(), FileReadTool(), PythonSandboxTool()]
+    allowed_imports = None
+    if settings and settings.python_tool_imports:
+        allowed_imports = set(settings.python_tool_imports)
+
+    return [
+        TimeTool(),
+        MathTool(),
+        FileReadTool(),
+        PythonSandboxTool(extra_allowed_imports=allowed_imports),
+    ]
 
 
 __all__ = ["Tool", "load_default_tools"]

simple_agent/tools/python_tool.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import subprocess
+import ast
 from textwrap import dedent
 
 from .base import SimpleTool
@@ -11,18 +12,44 @@
 class PythonSandboxTool(SimpleTool):
     """Executes small Python snippets in a separate interpreter."""
 
-    def __init__(self, *, timeout: int = 5) -> None:
+    def __init__(
+        self,
+        *,
+        timeout: int = 5,
+        extra_allowed_imports: set[str] | None = None,
+    ) -> None:
         super().__init__(
             name="python",
             description="Run Python code in a sandboxed interpreter. Provide raw code; stdout is returned.",
         )
         self.timeout = timeout
+        default_allowed = {
+            "math",
+            "statistics",
+            "datetime",
+            "time",
+            "random",
+            "json",
+            "collections",
+            "itertools",
+            "functools",
+            "os",
+            "sys",
+            "pathlib",
+            "psutil",
+        }
+        self.allowed_imports = default_allowed | (extra_allowed_imports or set())
 
     def run(self, query: str) -> str:
         code = query.strip()
         if not code:
             return "Provide Python code to run."
 
+        disallowed = _find_disallowed_imports(code, self.allowed_imports)
+        if disallowed:
+            allowed = ", ".join(sorted(self.allowed_imports))
+            return f"Imports not permitted: {', '.join(sorted(disallowed))}. Allowed modules: {allowed}."
+
         wrapped = dedent(
             f"""
             import sys
@@ -62,3 +89,26 @@ def run(self, query: str) -> str:
         if stderr:
             return f"[stderr]\n{stderr}"
         return stdout or "(no output)"
+
+
+def _find_disallowed_imports(code: str, allowed: set[str]) -> set[str]:
+    try:
+        tree = ast.parse(code)
+    except SyntaxError:
+        return set()
+
+    blocked: set[str] = set()
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Import):
+            for alias in node.names:
+                root = alias.name.split(".")[0]
+                if root not in allowed:
+                    blocked.add(root)
+        elif isinstance(node, ast.ImportFrom):
+            if node.module is None or node.level:
+                blocked.add("<relative>")
+                continue
+            root = node.module.split(".")[0]
+            if root not in allowed:
+                blocked.add(root)
+    return blocked