Modify context generation slightly (xrefs)

data_prep/project_context/context_introspector.py

@@ -162,8 +162,34 @@ def _get_xrefs_to_function(self) -> list[str]:
       logging.warning(
           'Could not retrieve xrefs for project: %s '
           'function_signature: %s', project, func_sig)
+
+    # At times, xrefs can be noisy (multiple 100s of loc).
+    # Truncate them by default
+    xrefs = self._truncate_xrefs(xrefs)
     return xrefs
 
+  def _truncate_xrefs(self, xrefs: list[str]) -> list[str]:
+    """Truncates xrefs to 10 lines before and after the
+    function name is referenced."""
+    truncated = []
+    for xref in xrefs:
+      lines = xref.split('\n')
+      line_index = -1
+      start = 0
+      end = len(lines) - 1
+      for index, line in enumerate(lines):
+        if self._benchmark.function_name in line:
+          line_index = index
+          break
+      # If name was not found, then just return the entire function.
+      # If it was, truncate it.
+      if line_index != -1:
+        start = start if line_index <= 10 else line_index - 10
+        end = end if line_index >= end - 10 else line_index + 10
+      truncated.append('\n'.join(lines[start:end]))
+
+    return truncated
+
   def get_context_info(self) -> dict:
     """Retrieves contextual information and stores them in a dictionary."""
     xrefs = self._get_xrefs_to_function()

llm_toolkit/prompt_builder.py

@@ -145,7 +145,7 @@ def format_context(self, context_info: dict) -> str:
         headers='\n'.join(context_info['files']),
         must_insert=context_info['decl'],
         func_source=context_info['func_source'],
-        xrefs='\n'.join(context_info['xrefs']),
+        xrefs=context_info['xrefs'],
     )
 
   def _select_examples(self, examples: list[list],

prompts/template_xml/context.txt

@@ -23,8 +23,11 @@ Here is the source code of the function being tested:
 {% endif %}
 {% if xrefs %}
 
-Here is the source code for functions which reference the function being tested:
+Here is the source code for locations which reference the function being tested:
+{% for xref in xrefs %}
+Reference location {{ loop.index }}:
 <code>
-{{ xrefs }}
+{{ xref }}
 </code>
+{% endfor %}
 {% endif %}