feat✨: add new tools for CLI plotting, codebase search, file search, grep, and directory listing

AlessandroAnnini · AlessandroAnnini · commit f510673169b1 · 2025-07-01T10:48:34.000+02:00
diff --git a/agent_loop/tools/cli_plot.py b/agent_loop/tools/cli_plot.py
@@ -0,0 +1,131 @@
+import plotext as plt
+import time
+
+tool_definition = {
+    "name": "cli_plot",
+    "description": (
+        "This is the default plotting and charting tool."
+        "Render advanced terminal charts using plotext. Use this to visualize data in the terminal, for charts, graphs, plots, etc. that have not to be saved to a file."
+        "Supports line, scatter, bar, histogram, datetime, log-scales, subplots, styling, and streaming."
+    ),
+    "input_schema": {
+        "type": "object",
+        "properties": {
+            "data": {
+                "type": "array",
+                "items": {
+                    "type": "object",
+                    "properties": {
+                        "x": {"type": "array", "items": {"type": ["number", "string"]}},
+                        "y": {"type": "array", "items": {"type": "number"}},
+                        "label": {"type": "string"},
+                        "type": {
+                            "type": "string",
+                            "enum": ["line", "scatter", "bar", "hist"],
+                        },
+                        "yaxis": {"type": "string", "enum": ["left", "right"]},
+                        "color": {"type": "string"},
+                        "marker": {"type": "string"},
+                        "fill": {"type": "boolean"},
+                    },
+                    "required": ["y", "type"],
+                },
+                "description": "List of series to plot.",
+            },
+            "style": {
+                "type": "object",
+                "properties": {
+                    "width": {"type": "integer"},
+                    "height": {"type": "integer"},
+                    "title": {"type": "string"},
+                    "xlabel": {"type": "string"},
+                    "ylabel": {"type": "array", "items": {"type": "string"}},
+                    "grid": {"type": "array", "items": {"type": "boolean"}},
+                    "xscale": {"type": "string"},
+                    "yscale": {"type": "string"},
+                },
+                "description": "Styling options for the plot.",
+            },
+            "stream": {
+                "type": "object",
+                "properties": {
+                    "frames": {"type": "integer"},
+                    "interval": {"type": "number"},
+                    "clear_terminal": {"type": "boolean"},
+                },
+                "description": "Streaming mode configuration.",
+            },
+        },
+        "required": ["data"],
+    },
+}
+
+
+def handle_call(input_data):
+    try:
+        plt.clear_figure()
+        data_series = input_data["data"]
+        style = input_data.get("style", {})
+        stream_cfg = input_data.get("stream")
+
+        # Styling
+        if style.get("width") and style.get("height"):
+            plt.plotsize(style["width"], style["height"])
+        if style.get("title"):
+            plt.title(style["title"])
+        if style.get("xlabel"):
+            plt.xlabel(style["xlabel"])
+        if style.get("ylabel"):
+            ylabels = style["ylabel"]
+            plt.ylabel(ylabels[0], *(ylabels[1:] if len(ylabels) > 1 else []))
+        if style.get("grid"):
+            plt.grid(*style["grid"])
+        if style.get("xscale"):
+            plt.xscale(style["xscale"])
+        if style.get("yscale"):
+            plt.yscale(style["yscale"])
+
+        # Plot data
+        for series in data_series:
+            func = {
+                "line": plt.plot,
+                "scatter": plt.scatter,
+                "bar": plt.bar,
+                "hist": plt.hist,
+            }[series["type"]]
+
+            args = []
+            if "x" in series:
+                args.append(series["x"])
+            args.append(series["y"])
+
+            kwargs = {}
+            for opt in ("label", "yaxis", "color", "marker"):
+                if series.get(opt):
+                    kwargs[opt] = series[opt]
+            if series.get("fill"):
+                kwargs["fillx" if series["type"] in ("line", "scatter") else "fill"] = (
+                    True
+                )
+
+            func(*args, **kwargs)
+
+        # Streaming mode
+        if stream_cfg:
+            frames = stream_cfg.get("frames", 1)
+            interval = stream_cfg.get("interval", 0.1)
+            clear_term = stream_cfg.get("clear_terminal", True)
+
+            for _ in range(frames):
+                if clear_term:
+                    plt.clear_terminal()
+                plt.show()
+                time.sleep(interval)
+                plt.clear_data()
+        else:
+            plt.show()
+
+        return {"status": "Rendered successfully."}
+
+    except Exception as e:
+        return {"error": f"Plotting failed: {e}"}
diff --git a/agent_loop/tools/codebase_search.py b/agent_loop/tools/codebase_search.py
@@ -0,0 +1,58 @@
+import subprocess
+
+tool_definition = {
+    "name": "codebase_search",
+    "description": (
+        "Find snippets of code from the codebase most relevant to the search query. "
+        "This is a semantic search tool, so the query should ask for something semantically matching what is needed. "
+        "Use target_directories to scope the search. Reuse the user's exact query unless there's a clear reason not to."
+    ),
+    "input_schema": {
+        "type": "object",
+        "properties": {
+            "query": {
+                "type": "string",
+                "description": "The search query to find relevant code. Reuse the user's exact wording.",
+            },
+            "target_directories": {
+                "type": "array",
+                "items": {"type": "string"},
+                "description": "Glob patterns for directories to search over.",
+            },
+            "explanation": {
+                "type": "string",
+                "description": "Why the agent is performing this semantic search and how it helps.",
+            },
+        },
+        "required": ["query"],
+    },
+}
+
+
+def handle_call(input_data):
+    query = input_data["query"]
+    target_dirs = input_data.get("target_directories", ["."])
+
+    try:
+        args = [
+            "semantic-code-search",
+            "--query",
+            query,
+        ]
+
+        for directory in target_dirs:
+            args.extend(["--dir", directory])
+
+        result = subprocess.run(args, capture_output=True, text=True, timeout=20)
+
+        if result.returncode != 0:
+            return {"error": result.stderr.strip(), "exit_code": result.returncode}
+
+        return {
+            "query": query,
+            "output": result.stdout.strip(),
+            "exit_code": result.returncode,
+        }
+
+    except Exception as e:
+        return {"error": f"Execution error: {e}"}
diff --git a/agent_loop/tools/file_search.py b/agent_loop/tools/file_search.py
@@ -0,0 +1,40 @@
+import subprocess
+
+tool_definition = {
+    "name": "file_search",
+    "description": (
+        "Fast file search based on fuzzy matching against file path. "
+        "Use if you know part of the file path but don't know where it's located exactly. "
+        "Response will be capped to 10 results. Make your query more specific if need to filter results further."
+    ),
+    "input_schema": {
+        "type": "object",
+        "properties": {
+            "query": {"type": "string", "description": "Fuzzy filename to search for"},
+            "explanation": {
+                "type": "string",
+                "description": "Why the agent is performing this file search and how it helps.",
+            },
+        },
+        "required": ["query", "explanation"],
+    },
+}
+
+
+def handle_call(input_data):
+    query = input_data["query"]
+
+    try:
+        result = subprocess.run(
+            ["fdfind", query], capture_output=True, text=True, timeout=10
+        )
+
+        if result.returncode not in [0, 1]:  # 0: found, 1: not found
+            return {"error": result.stderr.strip(), "exit_code": result.returncode}
+
+        matches = result.stdout.strip().splitlines()[:10]
+
+        return {"query": query, "matches": matches, "exit_code": result.returncode}
+
+    except Exception as e:
+        return {"error": f"Execution error: {e}"}
diff --git a/agent_loop/tools/grep.py b/agent_loop/tools/grep.py
@@ -0,0 +1,84 @@
+import subprocess
+import shlex
+
+tool_definition = {
+    "name": "grep_search",
+    "description": (
+        "Search for exact strings or regular expressions using `grep`. "
+        "This tool is ideal when the agent knows exactly what to look for (e.g., symbols, patterns, or keywords in files). "
+        "Use include/exclude filters to narrow down scope."
+    ),
+    "input_schema": {
+        "type": "object",
+        "properties": {
+            "query": {
+                "type": "string",
+                "description": "Regex or literal pattern to search for.",
+            },
+            "include_pattern": {
+                "type": "string",
+                "description": "Glob pattern for files to include, e.g. '*.py'",
+            },
+            "exclude_pattern": {
+                "type": "string",
+                "description": "Glob pattern for files to exclude, e.g. '*test*'",
+            },
+            "case_sensitive": {
+                "type": "boolean",
+                "description": "Case-sensitive search or not.",
+            },
+            "explanation": {
+                "type": "string",
+                "description": "Why the agent is performing this search.",
+            },
+            "directory": {
+                "type": "string",
+                "description": "Base directory to search in (default: current directory)",
+                "default": ".",
+            },
+        },
+        "required": ["query"],
+    },
+}
+
+
+def handle_call(input_data):
+    query = input_data["query"]
+    include = input_data.get("include_pattern")
+    exclude = input_data.get("exclude_pattern")
+    case_sensitive = input_data.get("case_sensitive", True)
+    directory = input_data.get("directory", ".")
+
+    # Base grep command
+    cmd = ["grep", "-r", "--line-number", "--color=never"]
+
+    if not case_sensitive:
+        cmd.append("-i")
+
+    if include:
+        cmd.extend(["--include", include])
+
+    if exclude:
+        cmd.extend(["--exclude", exclude])
+
+    cmd.extend([shlex.quote(query), shlex.quote(directory)])
+
+    try:
+        result = subprocess.run(
+            " ".join(cmd), shell=True, capture_output=True, text=True, timeout=15
+        )
+        output = result.stdout.strip()
+        if result.returncode == 0 or output:
+            return {
+                "matches": output.splitlines()[:50],
+                "exit_code": result.returncode,
+                "stderr": result.stderr,
+            }
+        else:
+            return {
+                "matches": [],
+                "exit_code": result.returncode,
+                "stderr": result.stderr,
+            }
+    except Exception as e:
+        return {"error": f"Execution error: {e}"}
diff --git a/agent_loop/tools/list_dir.py b/agent_loop/tools/list_dir.py
@@ -0,0 +1,46 @@
+import os
+
+tool_definition = {
+    "name": "list_dir",
+    "description": (
+        "List the contents of a directory. The quick tool to use for discovery, "
+        "before using more targeted tools like semantic search or file reading. "
+        "Useful to try to understand the file structure before diving deeper into specific files."
+    ),
+    "input_schema": {
+        "type": "object",
+        "properties": {
+            "relative_workspace_path": {
+                "type": "string",
+                "description": "Path to list contents of, relative to the workspace root.",
+            },
+            "explanation": {
+                "type": "string",
+                "description": "Why the agent is listing this directory and how it helps.",
+            },
+        },
+        "required": ["relative_workspace_path"],
+    },
+}
+
+
+def handle_call(input_data):
+    path = input_data["relative_workspace_path"]
+
+    try:
+        entries = os.listdir(path)
+        entries_info = []
+        for entry in sorted(entries):
+            full_path = os.path.join(path, entry)
+            if os.path.isdir(full_path):
+                entry_type = "directory"
+            elif os.path.isfile(full_path):
+                entry_type = "file"
+            else:
+                entry_type = "other"
+            entries_info.append({"name": entry, "type": entry_type})
+
+        return {"path": path, "contents": entries_info}
+
+    except Exception as e:
+        return {"error": f"Could not list directory '{path}': {e}"}
