Check-in missing function

Author: camelid

accelerant/chat.py

@@ -1,7 +1,6 @@
 import json
 from rich.markup import escape as rescape
 from typing import Any, List, Optional
-from llm_utils import number_group_of_lines
 import openai
 from openai import NotGiven
 from openai.types.chat import (
@@ -30,6 +29,7 @@
     GetSurroundingCodeTool,
     LLMToolRunner,
 )
+from accelerant.util import custom_number_group_of_lines
 from perfparser import LineLoc
 
 
@@ -64,7 +64,7 @@ def _build_hotspot_prompt(
         # FIXME: avoid crashing
         assert parent_sym is not None
         sline = parent_sym["range"]["start"]["line"]
-        prettyline = number_group_of_lines(
+        prettyline = custom_number_group_of_lines(
             project.get_range(filename, parent_sym["range"]),
             max(sline + 1, 1),
             with_note=lambda n: " <--- HOTSPOT" if n == lineno else "",

accelerant/util.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, TypedDict
+from typing import Callable, List, Optional, TypedDict
 
 
 class SymbolLoc(TypedDict):
@@ -47,3 +47,43 @@ def truncate_for_llm(text: str, char_limit: int):
     if len(text) > char_limit:
         return text[:char_limit] + "[...too long...]"
     return text
+
+
+def custom_number_group_of_lines(
+    group: List[str],
+    first: int,
+    strip: bool = True,
+    with_note: Optional[Callable[[int], str]] = None,
+) -> str:
+    """
+    Add line numbers for each line in the input list.
+
+    Args:
+        group (List[str]): The lines to number.
+        first (int): The number for the first line.
+        strip (bool): Whether to strip leading and trailing blank lines.
+
+    Returns:
+        A string concatenation of the numbered lines.
+    """
+    if strip:
+        while group and not group[0].strip():
+            group = group[1:]
+            first += 1
+        while group and not group[-1].strip():
+            group = group[:-1]
+
+    last = first + len(group) - 1
+    max_line_number_length = len(str(last))
+    result = "\n".join(
+        [
+            "{0:>{1}} {2}{3}".format(
+                first + i,
+                max_line_number_length,
+                line,
+                with_note(first + i) if with_note else "",
+            )
+            for i, line in enumerate(group)
+        ]
+    )
+    return result