/ src / python / txtai / agent / tool / grep.py
grep.py
 1  """
 2  Grep imports
 3  """
 4  
 5  import glob
 6  import os
 7  import re
 8  
 9  from smolagents import Tool
10  
11  
12  class GrepTool(Tool):
13      """
14      The GrepTool "greps" for file matches.
15      """
16  
17      # pylint: disable=W0231
18      def __init__(self):
19          """
20          Creates a GrepTool.
21          """
22  
23          # Tool parameters
24          self.name = "grep"
25          self.description = (
26              "Implementation of a grep tool. Searches a list of files for matches. "
27              "Returns a dictionary with the file path as the key and list of matches as the value."
28          )
29          self.inputs = {
30              "search": {"type": "string", "description": "Search grep pattern"},
31              "files": {"type": "string", "description": "File or file glob pattern to search"},
32          }
33          self.output_type = "any"
34  
35          # Validate parameters and initialize tool
36          super().__init__()
37  
38      # pylint: disable=W0221
39      def forward(self, search, files):
40          """
41          Greps a set of files for a pattern.
42  
43          Args:
44              search: search pattern
45              files: file or file glob pattern to search
46  
47          Returns:
48              {path: [matches]}
49          """
50  
51          # Compile the search pattern
52          pattern = re.compile(search)
53  
54          # Format file pattern
55          files = os.path.join(files, "*") if os.path.isdir(files) else files
56  
57          # Iterate over each found file
58          results = {}
59          for path in glob.glob(files, recursive=True):
60              if os.path.isfile(path):
61                  try:
62                      with open(path, "r", encoding="utf-8") as f:
63                          # Search for matches line-by-line
64                          matches = []
65                          for line in f:
66                              if pattern.search(line):
67                                  matches.append(line)
68  
69                          if matches:
70                              results[path] = matches
71  
72                  except (UnicodeDecodeError, FileNotFoundError):
73                      pass
74  
75          return results