Merge pull request #186 from pjcreath/profile-repomap

Add profiling information to repo map generation

Author: dwash96

.gitignore

@@ -33,3 +33,5 @@
 aider/__version__.py
 aider/_version.py
 *.pyc
+.aider*
+env/

aider/io.py

@@ -343,6 +343,8 @@ def __init__(
         self.bell_on_next_input = False
         self.notifications = notifications
         self.verbose = verbose
+        self.profile_start_time = None
+        self.profile_last_time = None
 
         # Variables used to interface with base_coder
         self.coder = None
@@ -1373,6 +1375,25 @@ def tool_output(self, *messages, log_only=False, bold=False):
 
         self.stream_print(*messages, style=style)
 
+    def profile(self, *messages, start=False):
+        if not self.verbose:
+            return
+
+        now = time.time()
+        message_str = " ".join(map(str, messages))
+
+        # Treat uninitialized as an implicit start.
+        if start or self.profile_start_time is None:
+            self.profile_start_time = now
+            self.stream_print(f"PROFILE: {message_str}")
+        else:
+            total_elapsed = now - self.profile_start_time
+            last_elapsed = now - self.profile_last_time
+            output_message = f"PROFILE: [+{last_elapsed:6.2f}s] {message_str} (total {total_elapsed:.2f}s)"
+            self.stream_print(output_message)
+
+        self.profile_last_time = now
+
     def assistant_output(self, message, pretty=None):
         if not message:
             self.tool_warning("Empty response received from LLM. Check your provider account?")

aider/repomap.py

@@ -198,6 +198,9 @@ def __init__(
         self._mentioned_ident_similarity = 0.8
 
         if self.verbose:
+            self.io.tool_output(
+                f"RepoMap loaded entries from tags cache: {len(self.TAGS_CACHE)}"
+            )
             self.io.tool_output(
                 f"RepoMap initialized with map_mul_no_files: {self.map_mul_no_files}"
             )
@@ -696,6 +699,8 @@ def get_ranked_tags(
                 if tag.specific_kind == "import":
                     file_imports[rel_fname].add(tag.name)
 
+        self.io.profile("process files")
+
         if self.use_enhanced_map and len(file_imports) > 0:
             import_ast_mode = True
 
@@ -791,6 +796,8 @@ def get_ranked_tags(
                     weight = num_refs * use_mul * 2 ** (-1 * path_distance)
                     G.add_edge(referencer, definer, weight=weight, key=ident, ident=ident)
 
+        self.io.profile("build graph")
+
         if not references:
             pass
 
@@ -808,6 +815,8 @@ def get_ranked_tags(
             except ZeroDivisionError:
                 return []
 
+        self.io.profile("pagerank")
+
         # distribute the rank from each source node, across all of its out edges
         ranked_definitions = defaultdict(float)
         for src in G.nodes:
@@ -822,6 +831,8 @@ def get_ranked_tags(
                 ident = data["ident"]
                 ranked_definitions[(dst, ident)] += data["rank"]
 
+        self.io.profile("distribute rank")
+
         ranked_tags = []
         ranked_definitions = sorted(
             ranked_definitions.items(), reverse=True, key=lambda x: (x[1], x[0])
@@ -905,9 +916,11 @@ def get_ranked_tags_map(
 
             # Check if the result is in the cache
             if use_cache and cache_key in self.map_cache:
+                self.io.tool_output("DEBUG: get_ranked_tags_map cache hit")
                 return self.map_cache[cache_key]
 
         # If not in cache or force_refresh is True, generate the map
+        self.io.tool_output("DEBUG: get_ranked_tags_map cache miss, generating map")
         start_time = time.time()
         result = self.get_ranked_tags_map_uncached(
             chat_fnames, other_fnames, max_map_tokens, mentioned_fnames, mentioned_idents
@@ -929,6 +942,7 @@ def get_ranked_tags_map_uncached(
         mentioned_fnames=None,
         mentioned_idents=None,
     ):
+        self.io.profile("get_ranked_tags_map_uncached", start=True)
         if not other_fnames:
             other_fnames = list()
         if not max_map_tokens:
@@ -943,6 +957,7 @@ def get_ranked_tags_map_uncached(
         ranked_tags = self.get_ranked_tags(
             chat_fnames, other_fnames, mentioned_fnames, mentioned_idents, True
         )
+        self.io.profile("get_ranked_tags")
 
         other_rel_fnames = sorted(set(self.get_rel_fname(fname) for fname in other_fnames))
         special_fnames = filter_important_files(other_rel_fnames)
@@ -992,6 +1007,7 @@ def get_ranked_tags_map_uncached(
 
             middle = int((lower_bound + upper_bound) // 2)
 
+        self.io.profile("binary search and to_tree")
         return best_tree
 
     tree_cache = dict()