Repository: ryanhoangt/locify

Issue:
Add max token count control
It would be helpful to have a way to control max tokens allowed for the map generated.

Discussion:

Patch:
diff --git a/locify/indexing/full_map/strategy.py b/locify/indexing/full_map/strategy.py
index 6449cea..663957f 100644
--- a/locify/indexing/full_map/strategy.py
+++ b/locify/indexing/full_map/strategy.py
@@ -2,7 +2,9 @@ from collections import defaultdict
 from pathlib import Path
 
 from grep_ast import TreeContext
+from tqdm import tqdm
 
+from locify.indexing.prompts import Prompts
 from locify.tree_sitter.parser import ParsedTag, TagKind, TreeSitterParser
 from locify.utils.file import GitRepoUtils, read_text
 from locify.utils.llm import get_token_count_from_text
@@ -10,21 +12,60 @@ from locify.utils.path import PathUtils
 
 
 class FullMapStrategy:
-    def __init__(self, model_name='gpt-4o', root='./') -> None:
+    def __init__(
+        self,
+        model_name='gpt-4o',
+        root='./',
+        max_map_token=1024 * 3,
+        content_prefix=Prompts.repo_content_prefix,
+    ) -> None:
         if not Path(root).is_absolute():
             root = str(Path(root).resolve())
 
         self.root = root
         self.model_name = model_name
+        self.max_map_token = max_map_token
+        self.content_prefix = content_prefix
 
         self.git_utils = GitRepoUtils(root)
         self.path_utils = PathUtils(root)
         self.ts_parser = TreeSitterParser()
 
+    def tags_to_tree(self, parsed_tags: list[ParsedTag]) -> str:
+        num_tags = len(parsed_tags)
+        lower_bound, upper_bound = 0, num_tags
+        best_tree_token_count = 0
+
+        # Assume each tag has 16 tokens on average, this trick is to deal with repos with lots of tags and using the exact middle as the starting point is too slow
+        mid = min(self.max_map_token // 16, num_tags) // 2
+        while lower_bound < upper_bound:
+            tree_repr = self.tag_list_to_tree(parsed_tags[:mid])
+            token_count = get_token_count_from_text(self.model_name, tree_repr)
+
+            if (
+                token_count <= self.max_map_token
+                and token_count > best_tree_token_count
+            ):
+                best_tree_token_count = token_count
+
+            if token_count > self.max_map_token:
+                upper_bound = mid
+            else:
+                lower_bound = mid + 1
+
+            mid = (lower_bound + upper_bound) // 2
+
+        result_tags = parsed_tags[:mid]
+        # Sort by relative file path and tag's line number
+        result_tags.sort(key=lambda tag: (tag.rel_path, tag.start_line))
+        return self.tag_list_to_tree(result_tags)
+
     def get_map(self, depth: int | None = None, rel_dir_path: str | None = None) -> str:
+        # t0 = time.time()
         ranked_tags = self.get_ranked_tags(rel_dir_path=rel_dir_path, depth=depth)
-        tree_repr = self.tag_list_to_tree(ranked_tags)
-        return tree_repr
+        tree_repr = self.tags_to_tree(ranked_tags)
+        # print(f'Getting map took {time.time() - t0:.2f}s')
+        return self.content_prefix + tree_repr
 
     def get_map_with_token_count(
         self, depth: int | None = None, rel_dir_path: str | None = None
@@ -48,7 +89,7 @@ class FullMapStrategy:
             set
         )  # (relative file, symbol identifier) -> set of its tags
 
-        for abs_file in all_abs_files:
+        for abs_file in tqdm(all_abs_files, desc='Parsing tags', unit='file'):
             rel_file = self.path_utils.get_relative_path_str(abs_file)
             parsed_tags = self.ts_parser.get_tags_from_file(abs_file, rel_file)
 
diff --git a/locify/indexing/prompts.py b/locify/indexing/prompts.py
new file mode 100644
index 0000000..8eff25a
--- /dev/null
+++ b/locify/indexing/prompts.py
@@ -0,0 +1,2 @@
+class Prompts:
+    repo_content_prefix = """Below are a tree of classes and functions for files in the workspace. It's chosen based on some heuristics and is not always extensive. So, please only use it for quick code navigation and not as a definitive source of truth:\n\n"""
diff --git a/locify/indexing/repo_map/strategy.py b/locify/indexing/repo_map/strategy.py
index 78d3041..115992e 100644
--- a/locify/indexing/repo_map/strategy.py
+++ b/locify/indexing/repo_map/strategy.py
@@ -2,14 +2,22 @@ import math
 from collections import Counter, defaultdict
 
 import networkx as nx
+from tqdm import tqdm
 
 from locify.indexing.full_map.strategy import FullMapStrategy
+from locify.indexing.prompts import Prompts
 from locify.tree_sitter.parser import ParsedTag, TagKind
 
 
 class RepoMapStrategy(FullMapStrategy):
-    def __init__(self, model_name='gpt-4o', root='./') -> None:
-        super().__init__(model_name, root)
+    def __init__(
+        self,
+        model_name='gpt-4o',
+        root='./',
+        max_map_token=1024 * 3,
+        content_prefix=Prompts.repo_content_prefix,
+    ) -> None:
+        super().__init__(model_name, root, max_map_token, content_prefix)
 
     def get_ranked_tags(
         self,
@@ -43,7 +51,7 @@ class RepoMapStrategy(FullMapStrategy):
         personalization_dict = {}
         personalization_val = 100 / num_files
 
-        for abs_file in all_abs_files:
+        for abs_file in tqdm(all_abs_files, desc='Parsing tags', unit='file'):
             rel_file = self.path_utils.get_relative_path_str(abs_file)
 
             if rel_file in mentioned_rel_files:
diff --git a/pyproject.toml b/pyproject.toml
index 70ccb39..39ae259 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,7 +1,7 @@
 [tool.poetry]
 name = "locify"
 version = "0.0.2"
-description = "A library for navigating over large codebases."
+description = "A library for LLM-based agents to navigate large codebases efficiently."
 authors = ["Hoang Tran <descience.thh10@gmail.com>"]
 license = "MIT"
 readme = "README.md"


Test patch:
diff --git a/tests/integration/test_repo_map_strategy.py b/tests/integration/test_repo_map_strategy.py
index 77a242a..a123dac 100644
--- a/tests/integration/test_repo_map_strategy.py
+++ b/tests/integration/test_repo_map_strategy.py
@@ -106,11 +106,15 @@ def test_repo_map_strategy_get_ranked_tags_with_refs(setup_git_repo_with_refs):
     # Ensure the token count is a reasonable positive integer
     assert token_count > 0
 
-    expected_map = """test_file1.py:
+    repo_content_prefix = strategy.content_prefix
+    expected_map = (
+        repo_content_prefix
+        + """test_file1.py:
   1│def foo():
 ...⋮...
   4│def bar():
 ...⋮...
 test_file2.py:
 ...⋮..."""
+    )
     assert repo_map == expected_map