mirror of
https://github.com/hanzoai/aci.git
synced 2026-10-30 18:49:51 +00:00
1118 lines
45 KiB
Python
1118 lines
45 KiB
Python
# Check if llama-index is installed
|
|
try:
|
|
import llama_index
|
|
import llama_index.core
|
|
import llama_index.retrievers.bm25
|
|
except ImportError:
|
|
raise ImportError("llama-index is required for these tools. Install with: pip install hanzo-aci[llama]")
|
|
|
|
import collections
|
|
import json
|
|
import os
|
|
import pickle
|
|
import re
|
|
from collections import defaultdict
|
|
from copy import deepcopy
|
|
from typing import List, Optional
|
|
|
|
import networkx as nx
|
|
|
|
from .repo.chunk_index.code_retriever import (
|
|
build_code_retriever_from_repo as build_code_retriever,
|
|
)
|
|
from .repo.dependency_graph import (
|
|
EDGE_TYPE_CONTAINS,
|
|
NODE_TYPE_CLASS,
|
|
NODE_TYPE_DIRECTORY,
|
|
NODE_TYPE_FILE,
|
|
NODE_TYPE_FUNCTION,
|
|
VALID_EDGE_TYPES,
|
|
VALID_NODE_TYPES,
|
|
RepoDependencySearcher,
|
|
RepoEntitySearcher,
|
|
build_graph,
|
|
traverse_tree_structure,
|
|
)
|
|
from .repo.dependency_graph.index_retriever import (
|
|
build_module_retriever_from_graph as build_module_retriever,
|
|
)
|
|
from .repo.dependency_graph.index_retriever import (
|
|
build_retriever_from_persist_dir as load_retriever,
|
|
)
|
|
from .repo.dependency_graph.index_retriever import (
|
|
fuzzy_retrieve_from_graph_nodes as fuzzy_retrieve,
|
|
)
|
|
from .repo.dependency_graph.traverse_graph import (
|
|
is_test_file,
|
|
traverse_json_structure,
|
|
)
|
|
from .results import (
|
|
QueryInfo,
|
|
QueryResult,
|
|
)
|
|
from .utils import (
|
|
find_matching_files_from_list,
|
|
merge_intervals,
|
|
)
|
|
|
|
REPO_PATH: str | None = None
|
|
GRAPH_INDEX_DIR: str | None = None
|
|
BM25_INDEX_DIR: str | None = None
|
|
ALL_FILE: list | None = None
|
|
ALL_CLASS: list | None = None
|
|
ALL_FUNC: list | None = None
|
|
|
|
DP_GRAPH_ENTITY_SEARCHER: RepoEntitySearcher | None = None
|
|
DP_GRAPH_DEPENDENCY_SEARCHER: RepoDependencySearcher | None = None
|
|
DP_GRAPH: nx.MultiDiGraph | None = None
|
|
|
|
|
|
def parse_repo_index():
|
|
global REPO_PATH, GRAPH_INDEX_DIR, BM25_INDEX_DIR
|
|
global DP_GRAPH_ENTITY_SEARCHER, DP_GRAPH_DEPENDENCY_SEARCHER, DP_GRAPH
|
|
|
|
if not REPO_PATH:
|
|
REPO_PATH = os.environ.get('REPO_PATH', '')
|
|
# print('='*5, REPO_PATH, '='*5)
|
|
|
|
if not GRAPH_INDEX_DIR or BM25_INDEX_DIR:
|
|
BASE_INDEX_DIR = os.path.join(REPO_PATH, '_index_data')
|
|
GRAPH_INDEX_DIR = os.path.join(BASE_INDEX_DIR, 'graph_index_v2.3')
|
|
BM25_INDEX_DIR = os.path.join(BASE_INDEX_DIR, 'bm25_index')
|
|
# print('='*5, BASE_INDEX_DIR, '='*5)
|
|
|
|
if not DP_GRAPH:
|
|
# setup graph traverser
|
|
graph_index_file = os.path.join(GRAPH_INDEX_DIR, 'code_graph.pkl')
|
|
if not os.path.exists(graph_index_file):
|
|
try:
|
|
os.makedirs(GRAPH_INDEX_DIR, exist_ok=True)
|
|
G = build_graph(REPO_PATH, global_import=True)
|
|
with open(graph_index_file, 'wb') as f:
|
|
pickle.dump(G, f)
|
|
# logger.debug(f'Parsed repo `{REPO_PATH}`')
|
|
print(f'Parsed repo `{REPO_PATH}`')
|
|
except Exception as e:
|
|
# logger.error(f'Error processing repo `{REPO_PATH}`: {e}')
|
|
print(f'Error processing repo `{REPO_PATH}`: {e}')
|
|
else:
|
|
G = pickle.load(open(graph_index_file, 'rb'))
|
|
|
|
DP_GRAPH_ENTITY_SEARCHER = RepoEntitySearcher(G)
|
|
DP_GRAPH_DEPENDENCY_SEARCHER = RepoDependencySearcher(G)
|
|
DP_GRAPH = G
|
|
|
|
global ALL_FILE, ALL_CLASS, ALL_FUNC
|
|
ALL_FILE = DP_GRAPH_ENTITY_SEARCHER.get_all_nodes_by_type(NODE_TYPE_FILE)
|
|
ALL_CLASS = DP_GRAPH_ENTITY_SEARCHER.get_all_nodes_by_type(NODE_TYPE_CLASS)
|
|
ALL_FUNC = DP_GRAPH_ENTITY_SEARCHER.get_all_nodes_by_type(NODE_TYPE_FUNCTION)
|
|
|
|
if not os.path.exists(os.path.join(BM25_INDEX_DIR, 'corpus.jsonl')):
|
|
absolute_repo_dir = os.path.abspath(REPO_PATH)
|
|
build_code_retriever(
|
|
absolute_repo_dir,
|
|
persist_path=BM25_INDEX_DIR,
|
|
similarity_top_k=10,
|
|
)
|
|
|
|
|
|
def get_current_repo_modules():
|
|
return ALL_FILE, ALL_CLASS, ALL_FUNC
|
|
|
|
|
|
def get_graph_entity_searcher() -> RepoEntitySearcher:
|
|
return DP_GRAPH_ENTITY_SEARCHER
|
|
|
|
|
|
def get_graph_dependency_searcher() -> RepoDependencySearcher:
|
|
return DP_GRAPH_DEPENDENCY_SEARCHER
|
|
|
|
|
|
def get_graph():
|
|
assert DP_GRAPH is not None
|
|
return DP_GRAPH
|
|
|
|
|
|
def get_module_name_by_line_num(file_path: str, line_num: int):
|
|
# TODO:
|
|
# if the given line isn't in a function of a class and the class is large,
|
|
# find the nearest two member functions and return
|
|
|
|
entity_searcher = get_graph_entity_searcher()
|
|
dp_searcher = get_graph_dependency_searcher()
|
|
|
|
cur_module = None
|
|
if entity_searcher.has_node(file_path):
|
|
module_nids, _ = dp_searcher.get_neighbors(
|
|
file_path, etype_filter=[EDGE_TYPE_CONTAINS]
|
|
)
|
|
module_ndatas = entity_searcher.get_node_data(module_nids)
|
|
for module in module_ndatas:
|
|
if module['start_line'] <= line_num <= module['end_line']:
|
|
cur_module = module # ['node_id']
|
|
break
|
|
if cur_module and cur_module['type'] == NODE_TYPE_CLASS:
|
|
func_nids, _ = dp_searcher.get_neighbors(
|
|
cur_module['node_id'], etype_filter=[EDGE_TYPE_CONTAINS]
|
|
)
|
|
func_ndatas = entity_searcher.get_node_data(
|
|
func_nids, return_code_content=True
|
|
)
|
|
for func in func_ndatas:
|
|
if func['start_line'] <= line_num <= func['end_line']:
|
|
cur_module = func # ['node_id']
|
|
break
|
|
|
|
if cur_module: # and cur_module['type'] in [NODE_TYPE_CLASS, NODE_TYPE_FUNCTION]
|
|
return cur_module
|
|
# module_ndata = entity_searcher.get_node_data([cur_module['node_id']], return_code_content=True)
|
|
# return module_ndata[0]
|
|
return None
|
|
|
|
|
|
def get_code_block_by_line_nums(query_info, context_window=20):
|
|
# file_path: str, line_nums: List[int]
|
|
searcher = get_graph_entity_searcher()
|
|
|
|
file_path = query_info.file_path_or_pattern
|
|
line_nums = query_info.line_nums
|
|
cur_query_results = []
|
|
|
|
file_data = searcher.get_node_data([file_path], return_code_content=False)[0]
|
|
line_intervals = []
|
|
res_modules = []
|
|
# res_code_blocks = None
|
|
for line in line_nums:
|
|
# First check if within a certain module
|
|
module_data = get_module_name_by_line_num(file_path, line)
|
|
|
|
# if not within any modules, then search for the nearby lines.
|
|
if not module_data:
|
|
min_line_num = max(1, line - context_window)
|
|
max_line_num = min(file_data['end_line'], line + context_window)
|
|
line_intervals.append((min_line_num, max_line_num))
|
|
|
|
elif module_data['node_id'] not in res_modules:
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='preview',
|
|
nid=module_data['node_id'],
|
|
ntype=module_data['type'],
|
|
start_line=module_data['start_line'],
|
|
end_line=module_data['end_line'],
|
|
retrieve_src=f'Retrieved code context including {query_info.term}.',
|
|
)
|
|
cur_query_results.append(query_result)
|
|
res_modules.append(module_data['node_id'])
|
|
|
|
if line_intervals:
|
|
line_intervals = merge_intervals(line_intervals)
|
|
for interval in line_intervals:
|
|
start_line, end_line = interval
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='code_snippet',
|
|
nid=file_path,
|
|
file_path=file_path,
|
|
start_line=start_line,
|
|
end_line=end_line,
|
|
retrieve_src=f'Retrieved code context including {query_info.term}.',
|
|
)
|
|
cur_query_results.append(query_result)
|
|
# res_code_blocks = line_wrap_content('\n'.join(file_content), line_intervals)
|
|
|
|
# return res_code_blocks, res_modules
|
|
return cur_query_results
|
|
|
|
|
|
def parse_node_id(nid: str):
|
|
nfile = nid.split(':')[0]
|
|
nname = nid.split(':')[-1]
|
|
return nfile, nname
|
|
|
|
|
|
def search_entity_in_global_dict(
|
|
term: str, include_files: Optional[List[str]] = None, prefix_term=None
|
|
):
|
|
searcher = get_graph_entity_searcher()
|
|
|
|
# TODO: hard code cases like "class Migration" and "function testing"
|
|
if term.startswith(('class ', 'Class')):
|
|
term = term[len('class ') :].strip()
|
|
elif term.startswith(('function ', 'Function ')):
|
|
term = term[len('function ') :].strip()
|
|
elif term.startswith(('method ', 'Method ')):
|
|
term = term[len('method ') :].strip()
|
|
elif term.startswith('def '):
|
|
term = term[len('def ') :].strip()
|
|
|
|
# TODO: lower case if not find
|
|
# TODO: filename xxx.py as key (also lowercase if not find)
|
|
# global_name_dict = None
|
|
if term in searcher.global_name_dict:
|
|
global_name_dict = searcher.global_name_dict
|
|
nids = global_name_dict[term]
|
|
elif term.lower() in searcher.global_name_dict_lowercase:
|
|
term = term.lower()
|
|
global_name_dict = searcher.global_name_dict_lowercase
|
|
nids = global_name_dict[term]
|
|
else:
|
|
return None
|
|
|
|
node_datas = searcher.get_node_data(nids, return_code_content=False)
|
|
found_entities_filter_dict = collections.defaultdict(list)
|
|
for ndata in node_datas:
|
|
nfile, _ = parse_node_id(ndata['node_id'])
|
|
if not include_files or nfile in include_files:
|
|
prefix_terms = []
|
|
# candidite_prefixes = ndata['node_id'].lower().replace('.py', '').replace('/', '.').split('.')
|
|
candidite_prefixes = re.split(
|
|
r'[./:]', ndata['node_id'].lower().replace('.py', '')
|
|
)[:-1]
|
|
if prefix_term:
|
|
prefix_terms = prefix_term.lower().split('.')
|
|
if not prefix_term or all(
|
|
[prefix in candidite_prefixes for prefix in prefix_terms]
|
|
):
|
|
found_entities_filter_dict[ndata['type']].append(ndata['node_id'])
|
|
|
|
return found_entities_filter_dict
|
|
|
|
|
|
def search_entity(query_info, include_files: List[str] = None):
|
|
term = query_info.term
|
|
searcher = get_graph_entity_searcher()
|
|
# cur_result = ''
|
|
continue_search = True
|
|
|
|
cur_query_results = []
|
|
|
|
# first: exact match in graph
|
|
if searcher.has_node(term):
|
|
continue_search = False
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='complete',
|
|
nid=term,
|
|
retrieve_src=f'Exact match found for entity name `{term}`.',
|
|
)
|
|
cur_query_results.append(query_result)
|
|
|
|
# TODO: __init__ not exsit
|
|
elif term.endswith('.__init__'):
|
|
nid = term[: -(len('.__init__'))]
|
|
if searcher.has_node(nid):
|
|
continue_search = False
|
|
node_data = searcher.get_node_data([nid], return_code_content=True)[0]
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='preview',
|
|
nid=nid,
|
|
ntype=node_data['type'],
|
|
start_line=node_data['start_line'],
|
|
end_line=node_data['end_line'],
|
|
retrieve_src=f'Exact match found for entity name `{nid}`.',
|
|
)
|
|
cur_query_results.append(query_result)
|
|
|
|
# second: search in global name dict
|
|
if continue_search:
|
|
found_entities_dict = search_entity_in_global_dict(term, include_files)
|
|
if not found_entities_dict:
|
|
found_entities_dict = search_entity_in_global_dict(term)
|
|
|
|
use_sub_term = False
|
|
used_term = term
|
|
if not found_entities_dict and '.' in term:
|
|
# for cases: class_name.method_name
|
|
try:
|
|
prefix_term = '.'.join(term.split('.')[:-1]).split()[
|
|
-1
|
|
] # incase of 'class '/ 'function '
|
|
except IndexError:
|
|
prefix_term = None
|
|
split_term = term.split('.')[-1].strip()
|
|
used_term = split_term
|
|
found_entities_dict = search_entity_in_global_dict(
|
|
split_term, include_files, prefix_term
|
|
)
|
|
if not found_entities_dict:
|
|
found_entities_dict = search_entity_in_global_dict(
|
|
split_term, prefix_term
|
|
)
|
|
if not found_entities_dict:
|
|
use_sub_term = True
|
|
found_entities_dict = search_entity_in_global_dict(split_term)
|
|
|
|
# TODO: split the term and find in global dict
|
|
|
|
if found_entities_dict:
|
|
for ntype, nids in found_entities_dict.items():
|
|
if not nids:
|
|
continue
|
|
# if not continue_search: break
|
|
|
|
# procee class and function in the same way
|
|
if ntype in [NODE_TYPE_FUNCTION, NODE_TYPE_CLASS, NODE_TYPE_FILE]:
|
|
if len(nids) <= 3:
|
|
node_datas = searcher.get_node_data(
|
|
nids, return_code_content=True
|
|
)
|
|
for ndata in node_datas:
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='preview',
|
|
nid=ndata['node_id'],
|
|
ntype=ndata['type'],
|
|
start_line=ndata['start_line'],
|
|
end_line=ndata['end_line'],
|
|
retrieve_src=f'Match found for entity name `{used_term}`.',
|
|
)
|
|
cur_query_results.append(query_result)
|
|
# continue_search = False
|
|
else:
|
|
node_datas = searcher.get_node_data(
|
|
nids, return_code_content=False
|
|
)
|
|
for ndata in node_datas:
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='fold',
|
|
nid=ndata['node_id'],
|
|
ntype=ndata['type'],
|
|
retrieve_src=f'Match found for entity name `{used_term}`.',
|
|
)
|
|
cur_query_results.append(query_result)
|
|
if not use_sub_term:
|
|
continue_search = False
|
|
else:
|
|
continue_search = True
|
|
|
|
# third: bm25 search (entity + content)
|
|
if continue_search:
|
|
module_nids = []
|
|
|
|
# append the file name to keyword?
|
|
# # if not any(symbol in file_path_or_pattern for symbol in ['*','?', '[', ']']):
|
|
# term_with_file = f'{file_path_or_pattern}:{term}'
|
|
# module_nids = bm25_module_retrieve(query=term_with_file, include_files=include_files)
|
|
|
|
# search entity by keyword
|
|
module_nids = bm25_module_retrieve(query=term, include_files=include_files)
|
|
if not module_nids:
|
|
module_nids = bm25_module_retrieve(query=term)
|
|
|
|
if not module_nids:
|
|
# result += f"No entity found using BM25 search. Try to use fuzzy search...\n"
|
|
module_nids = fuzzy_retrieve(term, graph=get_graph(), similarity_top_k=3)
|
|
|
|
module_datas = searcher.get_node_data(module_nids, return_code_content=True)
|
|
showed_module_num = 0
|
|
for module in module_datas[:5]:
|
|
if module['type'] in [NODE_TYPE_FILE, NODE_TYPE_DIRECTORY]:
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='fold',
|
|
nid=module['node_id'],
|
|
ntype=module['type'],
|
|
retrieve_src='Retrieved entity using keyword search (bm25).',
|
|
)
|
|
cur_query_results.append(query_result)
|
|
elif showed_module_num < 3:
|
|
showed_module_num += 1
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='preview',
|
|
nid=module['node_id'],
|
|
ntype=module['type'],
|
|
start_line=module['start_line'],
|
|
end_line=module['end_line'],
|
|
retrieve_src='Retrieved entity using keyword search (bm25).',
|
|
)
|
|
cur_query_results.append(query_result)
|
|
|
|
return (cur_query_results, continue_search)
|
|
|
|
|
|
def merge_query_results(query_results):
|
|
priority = ['complete', 'code_snippet', 'preview', 'fold']
|
|
merged_results = {}
|
|
all_query_results: List[QueryResult] = []
|
|
|
|
for qr in query_results:
|
|
if qr.format_mode == 'code_snippet':
|
|
all_query_results.append(qr)
|
|
|
|
elif qr.nid and qr.nid in merged_results:
|
|
# Merge query_info_list
|
|
if qr.query_info_list[0] not in merged_results[qr.nid].query_info_list:
|
|
merged_results[qr.nid].query_info_list.extend(qr.query_info_list)
|
|
|
|
# Select the format_mode with the highest priority
|
|
existing_format_mode = merged_results[qr.nid].format_mode
|
|
if priority.index(qr.format_mode) < priority.index(existing_format_mode):
|
|
merged_results[qr.nid].format_mode = qr.format_mode
|
|
merged_results[qr.nid].start_line = qr.start_line
|
|
merged_results[qr.nid].end_line = qr.end_line
|
|
merged_results[qr.nid].retrieve_src = qr.retrieve_src
|
|
|
|
elif qr.nid:
|
|
merged_results[qr.nid] = qr
|
|
|
|
all_query_results += list(merged_results.values())
|
|
return all_query_results
|
|
|
|
|
|
def rank_and_aggr_query_results(query_results, fixed_query_info_list):
|
|
query_info_list_dict = {}
|
|
|
|
for qr in query_results:
|
|
# Convert the query_info_list to a tuple so it can be used as a dictionary key
|
|
key = tuple(qr.query_info_list)
|
|
|
|
if key in query_info_list_dict:
|
|
query_info_list_dict[key].append(qr)
|
|
else:
|
|
query_info_list_dict[key] = [qr]
|
|
|
|
# for the key: sort by query
|
|
def sorting_key(key):
|
|
# Find the first matching element index from fixed_query_info_list in the key (tuple of query_info_list)
|
|
for i, fixed_query in enumerate(fixed_query_info_list):
|
|
if fixed_query in key:
|
|
return i
|
|
# If no match is found, assign a large index to push it to the end
|
|
return len(fixed_query_info_list)
|
|
|
|
sorted_keys = sorted(query_info_list_dict.keys(), key=sorting_key)
|
|
sorted_query_info_list_dict = {
|
|
key: query_info_list_dict[key] for key in sorted_keys
|
|
}
|
|
|
|
# for the value: sort by format priority
|
|
priority = {
|
|
'complete': 1,
|
|
'code_snippet': 2,
|
|
'preview': 3,
|
|
'fold': 4,
|
|
} # Lower value indicates higher priority
|
|
# TODO: merge the same node in 'code_snippet' and 'preview'
|
|
|
|
organized_dict = {}
|
|
for key, values in sorted_query_info_list_dict.items():
|
|
nested_dict = {priority_key: [] for priority_key in priority.keys()}
|
|
for qr in values:
|
|
# Place the qr in the nested dictionary based on its format_mode
|
|
if qr.format_mode in nested_dict:
|
|
nested_dict[qr.format_mode].append(qr)
|
|
|
|
# Only add keys with non-empty lists to keep the result clean
|
|
organized_dict[key] = {k: v for k, v in nested_dict.items() if v}
|
|
|
|
return organized_dict
|
|
|
|
|
|
def search_code_snippets(
|
|
search_terms: Optional[List[str]] = None,
|
|
line_nums: Optional[List] = None,
|
|
file_path_or_pattern: Optional[str] = '**/*.py',
|
|
) -> str:
|
|
"""Searches the codebase to retrieve relevant code snippets based on given queries(terms or line numbers).
|
|
This function supports retrieving the complete content of a code entity,
|
|
searching for code entities such as classes or functions by keywords, or locating specific lines within a file.
|
|
It also supports filtering searches based on a file path or file pattern.
|
|
Note:
|
|
1. If `search_terms` are provided, it searches for code snippets based on each term:
|
|
- If a term is formatted as 'file_path:QualifiedName' (e.g., 'src/helpers/math_helpers.py:MathUtils.calculate_sum') ,
|
|
or just 'file_path', the corresponding complete code is retrieved or file content is retrieved.
|
|
- If a term matches a file, class, or function name, matched entities are retrieved.
|
|
- If there is no match with any module name, it attempts to find code snippets that likely contain the term.
|
|
2. If `line_nums` is provided, it searches for code snippets at the specified lines within the file defined by
|
|
`file_path_or_pattern`.
|
|
Args:
|
|
search_terms (Optional[List[str]]): A list of names, keywords, or code snippets to search for within the codebase.
|
|
Terms can be formatted as 'file_path:QualifiedName' to search for a specific module or entity within a file
|
|
(e.g., 'src/helpers/math_helpers.py:MathUtils.calculate_sum') or as 'file_path' to retrieve the complete content
|
|
of a file. This can also include potential function names, class names, or general code fragments.
|
|
line_nums (Optional[List[int]]): Specific line numbers to locate code snippets within a specified file.
|
|
When provided, `file_path_or_pattern` must specify a valid file path.
|
|
file_path_or_pattern (Optional[str]): A glob pattern or specific file path used to filter search results
|
|
to particular files or directories. Defaults to '**/*.py', meaning all Python files are searched by default.
|
|
If `line_nums` are provided, this must specify a specific file path.
|
|
Returns:
|
|
str: The search results, which may include code snippets, matching entities, or complete file content.
|
|
Example Usage:
|
|
# Search for the full content of a specific file
|
|
result = search_code_snippets(search_terms=['src/my_file.py'])
|
|
# Search for a specific function
|
|
result = search_code_snippets(search_terms=['src/my_file.py:MyClass.func_name'])
|
|
# Search for specific lines (10 and 15) within a file
|
|
result = search_code_snippets(line_nums=[10, 15], file_path_or_pattern='src/example.py')
|
|
# Combined search for a module name and within a specific file pattern
|
|
result = search_code_snippets(search_terms=["MyClass"], file_path_or_pattern="src/**/*.py")
|
|
"""
|
|
parse_repo_index()
|
|
files, _, _ = get_current_repo_modules()
|
|
all_file_paths = [file['name'] for file in files]
|
|
|
|
result = ''
|
|
# exclude_files = find_matching_files_from_list(all_file_paths, "**/test*/**")
|
|
if file_path_or_pattern:
|
|
include_files = find_matching_files_from_list(
|
|
all_file_paths, file_path_or_pattern
|
|
)
|
|
if not include_files:
|
|
include_files = all_file_paths
|
|
result += f"No files found for file pattern '{file_path_or_pattern}'. Will search all files.\n...\n"
|
|
else:
|
|
include_files = all_file_paths
|
|
|
|
query_info_list = []
|
|
all_query_results = []
|
|
|
|
if search_terms:
|
|
# search all terms together
|
|
filter_terms = []
|
|
for term in search_terms:
|
|
if is_test_file(term):
|
|
result += f'No results for test files: `{term}`. Please do not search for any test files.\n\n'
|
|
else:
|
|
filter_terms.append(term)
|
|
|
|
joint_terms = deepcopy(filter_terms)
|
|
if len(filter_terms) > 1:
|
|
filter_terms.append(' '.join(filter_terms))
|
|
|
|
for i, term in enumerate(filter_terms):
|
|
term = term.strip().strip('.')
|
|
if not term:
|
|
continue
|
|
|
|
query_info = QueryInfo(term=term)
|
|
query_info_list.append(query_info)
|
|
|
|
cur_query_results = []
|
|
|
|
# search entity
|
|
query_results, continue_search = search_entity(
|
|
query_info=query_info, include_files=include_files
|
|
)
|
|
cur_query_results.extend(query_results)
|
|
|
|
# search content
|
|
if continue_search:
|
|
query_results = bm25_content_retrieve(
|
|
query_info=query_info, include_files=include_files
|
|
)
|
|
cur_query_results.extend(query_results)
|
|
|
|
elif i != (len(filter_terms) - 1):
|
|
joint_terms[i] = ''
|
|
filter_terms[-1] = ' '.join([t for t in joint_terms if t.strip()])
|
|
if filter_terms[-1] in filter_terms[:-1]:
|
|
filter_terms[-1] = ''
|
|
|
|
all_query_results.extend(cur_query_results)
|
|
|
|
if file_path_or_pattern in all_file_paths and line_nums:
|
|
if isinstance(line_nums, int):
|
|
line_nums = [line_nums]
|
|
file_path = file_path_or_pattern
|
|
term = file_path + ':line ' + ', '.join([str(line) for line in line_nums])
|
|
# result += f"Search `line(s) {line_nums}` in file `{file_path}` ...\n"
|
|
query_info = QueryInfo(
|
|
term=term, line_nums=line_nums, file_path_or_pattern=file_path
|
|
)
|
|
|
|
# Search for codes based on file name and line number
|
|
query_results = get_code_block_by_line_nums(query_info)
|
|
all_query_results.extend(query_results)
|
|
|
|
merged_results = merge_query_results(all_query_results)
|
|
ranked_query_to_results = rank_and_aggr_query_results(
|
|
merged_results, query_info_list
|
|
)
|
|
|
|
# format output
|
|
# format_mode: 'complete', 'preview', 'code_snippet', 'fold': 4
|
|
searcher = get_graph_entity_searcher()
|
|
|
|
for query_infos, format_to_results in ranked_query_to_results.items():
|
|
term_desc = ', '.join([f'"{query.term}"' for query in query_infos])
|
|
result += f'##Searching for term {term_desc}...\n'
|
|
result += '### Search Result:\n'
|
|
cur_result = ''
|
|
for format_mode, query_results in format_to_results.items():
|
|
if format_mode == 'fold':
|
|
cur_retrieve_src = ''
|
|
for qr in query_results:
|
|
if not cur_retrieve_src:
|
|
cur_retrieve_src = qr.retrieve_src
|
|
|
|
if cur_retrieve_src != qr.retrieve_src:
|
|
cur_result += 'Source: ' + cur_retrieve_src + '\n\n'
|
|
cur_retrieve_src = qr.retrieve_src
|
|
|
|
cur_result += qr.format_output(searcher)
|
|
|
|
cur_result += 'Source: ' + cur_retrieve_src + '\n'
|
|
if len(query_results) > 1:
|
|
cur_result += 'Hint: Use more detailed query to get the full content of some if needed.\n'
|
|
else:
|
|
cur_result += f'Hint: Search `{query_results[0].nid}` for the full content if needed.\n'
|
|
cur_result += '\n'
|
|
|
|
elif format_mode == 'complete':
|
|
for qr in query_results:
|
|
cur_result += qr.format_output(searcher)
|
|
cur_result += '\n'
|
|
|
|
elif format_mode == 'preview':
|
|
# Remove the small modules, leaving only the large ones
|
|
filtered_results = []
|
|
grouped_by_file = defaultdict(list)
|
|
for qr in query_results:
|
|
if (qr.end_line - qr.start_line) < 100:
|
|
grouped_by_file[qr.file_path].append(qr)
|
|
else:
|
|
filtered_results.append(qr)
|
|
|
|
for file_path, results in grouped_by_file.items():
|
|
# Sort by start_line and then by end_line in descending order
|
|
sorted_results = sorted(
|
|
results, key=lambda qr: (qr.start_line, -qr.end_line)
|
|
)
|
|
|
|
max_end_line = -1
|
|
for qr in sorted_results:
|
|
# If the current QueryResult's range is not completely covered by the largest range seen so far, keep it
|
|
if qr.end_line > max_end_line:
|
|
filtered_results.append(qr)
|
|
max_end_line = max(max_end_line, qr.end_line)
|
|
|
|
# filtered_results = query_results
|
|
for qr in filtered_results:
|
|
cur_result += qr.format_output(searcher)
|
|
cur_result += '\n'
|
|
|
|
elif format_mode == 'code_snippet':
|
|
for qr in query_results:
|
|
cur_result += qr.format_output(searcher)
|
|
cur_result += '\n'
|
|
|
|
cur_result += '\n\n'
|
|
|
|
if cur_result.strip():
|
|
result += cur_result
|
|
else:
|
|
result += 'No locations found.\n\n'
|
|
|
|
return result.strip()
|
|
|
|
|
|
def get_entity_contents(entity_names: List[str]):
|
|
"""
|
|
Retrieves the complete implementations of specified entities from the codebase.
|
|
This function allows searching for function implementations, class definitions, or entire file contents
|
|
based on the provided entity names. Entity names can be specified in the following formats:
|
|
- `file_path:QualifiedName` (e.g., `'src/helpers/math_helpers.py:MathUtils.calculate_sum'`)
|
|
to retrieve a specific function or class definition.
|
|
- `file_path` (e.g., `'src/my_file.py'`) to retrieve the entire content of a file.
|
|
Args:
|
|
entity_names (List[str]):
|
|
A list of entity names to query. Each entity name can represent a function, class, or file.
|
|
Returns:
|
|
Dict[str, str]:
|
|
A dictionary mapping each entity name to its corresponding implementation or file content.
|
|
Example:
|
|
>>> get_entity_contents(['src/my_file.py:MyClass.func_name'])
|
|
>>> get_entity_contents(['src/my_file.py'])
|
|
"""
|
|
|
|
parse_repo_index()
|
|
searcher = get_graph_entity_searcher()
|
|
|
|
result = ''
|
|
for name in entity_names:
|
|
name = name.strip().strip('.')
|
|
if not name:
|
|
continue
|
|
|
|
result += f'##Searching for entity `{name}`...\n'
|
|
result += '### Search Result:\n'
|
|
query_info = QueryInfo(term=name)
|
|
|
|
if searcher.has_node(name):
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='complete',
|
|
nid=name,
|
|
retrieve_src=f'Exact match found for entity name `{name}`.',
|
|
)
|
|
result += query_result.format_output(searcher)
|
|
result += '\n\n'
|
|
else:
|
|
result += 'Invalid name. \nHint: Valid entity name should be formatted as "file_path:QualifiedName" or just "file_path".'
|
|
result += '\n\n'
|
|
return result.strip()
|
|
|
|
|
|
def bm25_module_retrieve(
|
|
query: str,
|
|
include_files: Optional[List[str]] = None,
|
|
# file_pattern: Optional[str] = None,
|
|
search_scope: str = 'all',
|
|
similarity_top_k: int = 10,
|
|
# sort_by_type = False
|
|
):
|
|
retriever = build_module_retriever(
|
|
entity_searcher=get_graph_entity_searcher(),
|
|
search_scope=search_scope,
|
|
similarity_top_k=similarity_top_k,
|
|
)
|
|
try:
|
|
retrieved_nodes = retriever.retrieve(query)
|
|
except IndexError as e:
|
|
# logger.warning(f'{e}. Probably because the query `{query}` is too short.')
|
|
print(f'{e}. Probably because the query `{query}` is too short.')
|
|
return []
|
|
|
|
filter_nodes = []
|
|
all_nodes = []
|
|
for node in retrieved_nodes:
|
|
if node.score <= 0:
|
|
continue
|
|
if not include_files or node.text.split(':')[0] in include_files:
|
|
filter_nodes.append(node.text)
|
|
all_nodes.append(node.text)
|
|
|
|
if filter_nodes:
|
|
return filter_nodes
|
|
else:
|
|
return all_nodes
|
|
|
|
|
|
def bm25_content_retrieve(
|
|
query_info: QueryInfo,
|
|
# query: str,
|
|
include_files: Optional[List[str]] = None,
|
|
# file_pattern: Optional[str] = None,
|
|
similarity_top_k: int = 10,
|
|
) -> str:
|
|
"""Retrieves code snippets from the codebase using the BM25 algorithm based on the provided query, class names, and function names. This function helps in finding relevant code sections that match specific criteria, aiding in code analysis and understanding.
|
|
Args:
|
|
query (Optional[str]): A textual query to search for relevant code snippets. Defaults to an empty string if not provided.
|
|
class_names (list[str]): A list of class names to include in the search query. If None, class names are not included.
|
|
function_names (list[str]): A list of function names to include in the search query. If None, function names are not included.
|
|
file_pattern (Optional[str]): A glob pattern to filter search results to specific file types or directories. If None, the search includes all files.
|
|
similarity_top_k (int): The number of top similar documents to retrieve based on the BM25 ranking. Defaults to 15.
|
|
Returns:
|
|
str: A formatted string containing the search results, including file paths and the retrieved code snippets (the partial code of a module or the skeleton of the specific module).
|
|
"""
|
|
|
|
# instance = get_current_issue_data()
|
|
# instance_id = get_current_issue_id()
|
|
query = query_info.term
|
|
|
|
# persist_path = os.path.join(BM25_INDEX_DIR, instance_id)
|
|
if os.path.exists(os.path.join(BM25_INDEX_DIR, 'corpus.jsonl')):
|
|
# TODO: if similairy_top_k > cache's setting, then regenerate
|
|
retriever = load_retriever(BM25_INDEX_DIR)
|
|
else:
|
|
absolute_repo_dir = os.path.abspath(REPO_PATH)
|
|
retriever = build_code_retriever(
|
|
absolute_repo_dir,
|
|
persist_path=BM25_INDEX_DIR,
|
|
similarity_top_k=similarity_top_k,
|
|
)
|
|
|
|
# similarity: {score}
|
|
cur_query_results = []
|
|
retrieved_nodes = retriever.retrieve(query)
|
|
for node in retrieved_nodes:
|
|
file = node.metadata['file_path']
|
|
# print(node.metadata)
|
|
if not include_files or file in include_files:
|
|
# drop the import code
|
|
# if len(node.metadata['span_ids']) == 1 and node.metadata['span_ids'][0] == 'imports':
|
|
# continue
|
|
if all(
|
|
[
|
|
span_id in ['docstring', 'imports', 'comments']
|
|
for span_id in node.metadata['span_ids']
|
|
]
|
|
):
|
|
# TODO: drop ?
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='code_snippet',
|
|
nid=node.metadata['file_path'],
|
|
file_path=node.metadata['file_path'],
|
|
start_line=node.metadata['start_line'],
|
|
end_line=node.metadata['end_line'],
|
|
retrieve_src='Retrieved code content using keyword search (bm25).',
|
|
)
|
|
cur_query_results.append(query_result)
|
|
|
|
elif any(
|
|
[
|
|
span_id in ['docstring', 'imports', 'comments']
|
|
for span_id in node.metadata['span_ids']
|
|
]
|
|
):
|
|
nids = []
|
|
for span_id in node.metadata['span_ids']:
|
|
nid = f'{file}:{span_id}'
|
|
searcher = get_graph_entity_searcher()
|
|
if searcher.has_node(nid):
|
|
nids.append(nid)
|
|
# TODO: warning if not find
|
|
|
|
node_datas = searcher.get_node_data(nids, return_code_content=True)
|
|
sorted_ndatas = sorted(node_datas, key=lambda x: x['start_line'])
|
|
sorted_nids = [ndata['node_id'] for ndata in sorted_ndatas]
|
|
|
|
message = ''
|
|
if sorted_nids:
|
|
if sorted_ndatas[0]['start_line'] < node.metadata['start_line']:
|
|
nid = sorted_ndatas[0]['node_id']
|
|
ntype = sorted_ndatas[0]['type']
|
|
# The code for {ntype} {nid} is incomplete; search {nid} for the full content if needed.
|
|
message += f'The code for {ntype} `{nid}` is incomplete; search `{nid}` for the full content if needed.\n'
|
|
if sorted_ndatas[-1]['end_line'] > node.metadata['end_line']:
|
|
nid = sorted_ndatas[-1]['node_id']
|
|
ntype = sorted_ndatas[-1]['type']
|
|
message += f'The code for {ntype} `{nid}` is incomplete; search `{nid}` for the full content if needed.\n'
|
|
if message.strip():
|
|
message = 'Hint: \n' + message
|
|
|
|
nids_str = ', '.join([f'`{nid}`' for nid in sorted_nids])
|
|
desc = f'Found {nids_str}.'
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='code_snippet',
|
|
nid=node.metadata['file_path'],
|
|
file_path=node.metadata['file_path'],
|
|
start_line=node.metadata['start_line'],
|
|
end_line=node.metadata['end_line'],
|
|
desc=desc,
|
|
message=message,
|
|
retrieve_src='Retrieved code content using keyword search (bm25).',
|
|
)
|
|
|
|
cur_query_results.append(query_result)
|
|
else:
|
|
for span_id in node.metadata['span_ids']:
|
|
nid = f'{file}:{span_id}'
|
|
searcher = get_graph_entity_searcher()
|
|
if searcher.has_node(nid):
|
|
ndata = searcher.get_node_data([nid], return_code_content=True)[
|
|
0
|
|
]
|
|
query_result = QueryResult(
|
|
query_info=query_info,
|
|
format_mode='preview',
|
|
nid=ndata['node_id'],
|
|
ntype=ndata['type'],
|
|
start_line=ndata['start_line'],
|
|
end_line=ndata['end_line'],
|
|
retrieve_src='Retrieved code content using keyword search (bm25).',
|
|
)
|
|
cur_query_results.append(query_result)
|
|
else:
|
|
continue
|
|
|
|
cur_query_results = cur_query_results[:5]
|
|
return cur_query_results
|
|
|
|
|
|
def _validate_graph_explorer_inputs(
|
|
start_entities: List[str],
|
|
direction: str = 'downstream',
|
|
traversal_depth: int = 1,
|
|
node_type_filter: Optional[List[str]] = None,
|
|
edge_type_filter: Optional[List[str]] = None,
|
|
):
|
|
"""evaluate input arguments"""
|
|
|
|
# assert len(invalid_entities) == 0, (
|
|
# f"Invalid value for `start_entities`: entities {invalid_entities} are not in the repository graph."
|
|
# )
|
|
assert direction in ['downstream', 'upstream', 'both'], (
|
|
"Invalid value for `direction`: Expected one of 'downstream', 'upstream', and 'both'. "
|
|
f"Received: '{direction}'."
|
|
)
|
|
assert traversal_depth == -1 or traversal_depth >= 0, (
|
|
'Invalid value for `traversal_depth`: It must be either -1 or a non-negative integer (>= 0). '
|
|
f'Received: {traversal_depth}.'
|
|
)
|
|
if isinstance(node_type_filter, list):
|
|
invalid_ntypes = []
|
|
for ntype in invalid_ntypes:
|
|
if ntype not in VALID_NODE_TYPES:
|
|
invalid_ntypes.append(ntype)
|
|
assert (
|
|
len(invalid_ntypes) == 0
|
|
), f'Invalid node types {invalid_ntypes} in node_type_filter. Expected node type in {VALID_NODE_TYPES}'
|
|
if isinstance(edge_type_filter, list):
|
|
invalid_etypes = []
|
|
for etype in edge_type_filter:
|
|
if etype not in VALID_EDGE_TYPES:
|
|
invalid_etypes.append(etype)
|
|
assert (
|
|
len(invalid_etypes) == 0
|
|
), f'Invalid edge types {invalid_etypes} in edge_type_filter. Expected edge type in {VALID_EDGE_TYPES}'
|
|
|
|
graph = get_graph()
|
|
entity_searcher = get_graph_entity_searcher()
|
|
|
|
hints = ''
|
|
valid_entities = []
|
|
for i, root in enumerate(start_entities):
|
|
# process node name
|
|
if root != '/':
|
|
root = root.strip('/')
|
|
if root.endswith('.__init__'):
|
|
root = root[: -(len('.__init__'))]
|
|
|
|
# validate node name
|
|
if root not in graph:
|
|
# search with bm25
|
|
module_nids = bm25_module_retrieve(query=root)
|
|
module_datas = entity_searcher.get_node_data(
|
|
module_nids, return_code_content=False
|
|
)
|
|
if len(module_datas) > 0:
|
|
hints += f'The entity name `{root}` is invalid. Based on your input, here are some candidate entities you might be referring to:\n'
|
|
for module in module_datas[:5]:
|
|
ntype = module['type']
|
|
nid = module['node_id']
|
|
hints += f'{ntype}: `{nid}`\n'
|
|
hints += 'Source: Retrieved entity using keyword search (bm25).\n\n'
|
|
else:
|
|
hints += f'The entity name `{root}` is invalid. There are no possible candidate entities in record.\n'
|
|
elif is_test_file(root):
|
|
hints += f'No results for the test entity: `{root}`. Please do not include any test entities.\n\n'
|
|
else:
|
|
valid_entities.append(root)
|
|
|
|
return valid_entities, hints
|
|
|
|
|
|
def explore_tree_structure(
|
|
start_entities: List[str],
|
|
direction: str = 'downstream',
|
|
traversal_depth: int = 2,
|
|
entity_type_filter: Optional[List[str]] = None,
|
|
dependency_type_filter: Optional[List[str]] = None,
|
|
):
|
|
"""Analyzes and displays the dependency structure around specified entities in a code graph.
|
|
This function searches and presents relationships and dependencies for the specified entities (such as classes, functions, files, or directories) in a code graph.
|
|
It explores how the input entities relate to others, using defined types of dependencies, including 'contains', 'imports', 'invokes' and 'inherits'.
|
|
The search can be controlled to traverse upstream (exploring dependencies that entities rely on) or downstream (exploring how entities impact others), with optional limits on traversal depth and filters for entity and dependency types.
|
|
Example Usage:
|
|
1. Exploring Outward Dependencies:
|
|
```
|
|
get_local_structure(
|
|
start_entities=['src/module_a.py:ClassA'],
|
|
direction='downstream',
|
|
traversal_depth=2,
|
|
entity_type_filter=['class', 'function'],
|
|
dependency_type_filter=['invokes', 'imports']
|
|
)
|
|
```
|
|
This retrieves the dependencies of `ClassA` up to 2 levels deep, focusing only on classes and functions with 'invokes' and 'imports' relationships.
|
|
2. Exploring Inward Dependencies:
|
|
```
|
|
get_local_structure(
|
|
start_entities=['src/module_b.py:FunctionY'],
|
|
direction='upstream',
|
|
traversal_depth=-1
|
|
)
|
|
```
|
|
This finds all entities that depend on `FunctionY` without restricting the traversal depth.
|
|
Notes:
|
|
* Traversal Control: The `traversal_depth` parameter specifies how deep the function should explore the graph starting from the input entities.
|
|
* Filtering: Use `entity_type_filter` and `dependency_type_filter` to narrow down the scope of the search, focusing on specific entity types and relationships.
|
|
* Graph Context: The function operates on a pre-built code graph containing entities (e.g., files, classes and functions) and dependencies representing their interactions and relationships.
|
|
Parameters:
|
|
----------
|
|
start_entities : list[str]
|
|
List of entities (e.g., class, function, file, or directory paths) to begin the search from.
|
|
- Entities representing classes or functions must be formatted as "file_path:QualifiedName"
|
|
(e.g., `interface/C.py:C.method_a.inner_func`).
|
|
- For files or directories, provide only the file or directory path (e.g., `src/module_a.py` or `src/`).
|
|
direction : str, optional
|
|
Direction of traversal in the code graph; allowed options are:
|
|
- 'upstream': Traversal to explore dependencies that the specified entities rely on (how they depend on others).
|
|
- 'downstream': Traversal to explore the effects or interactions of the specified entities on others
|
|
(how others depend on them).
|
|
- 'both': Traversal in both directions.
|
|
Default is 'downstream'.
|
|
traversal_depth : int, optional
|
|
Maximum depth of traversal. A value of -1 indicates unlimited depth (subject to a maximum limit).
|
|
Must be either `-1` or a non-negative integer (≥ 0).
|
|
Default is 2.
|
|
entity_type_filter : list[str], optional
|
|
List of entity types (e.g., 'class', 'function', 'file', 'directory') to include in the traversal.
|
|
If None, all entity types are included.
|
|
Default is None.
|
|
dependency_type_filter : list[str], optional
|
|
List of dependency types (e.g., 'contains', 'imports', 'invokes', 'inherits') to include in the traversal.
|
|
If None, all dependency types are included.
|
|
Default is None.
|
|
Returns:
|
|
-------
|
|
result : object
|
|
An object representing the traversal results, which includes discovered entities and their dependencies.
|
|
"""
|
|
parse_repo_index()
|
|
start_entities, hints = _validate_graph_explorer_inputs(
|
|
start_entities,
|
|
direction,
|
|
traversal_depth,
|
|
entity_type_filter,
|
|
dependency_type_filter,
|
|
)
|
|
G = get_graph()
|
|
|
|
# return_json = True
|
|
return_json = False
|
|
if return_json:
|
|
rtns = {
|
|
node: traverse_json_structure(
|
|
G,
|
|
node,
|
|
direction,
|
|
traversal_depth,
|
|
entity_type_filter,
|
|
dependency_type_filter,
|
|
)
|
|
for node in start_entities
|
|
}
|
|
rtn_str = json.dumps(rtns)
|
|
else:
|
|
rtns = [
|
|
traverse_tree_structure(
|
|
G,
|
|
node,
|
|
direction,
|
|
traversal_depth,
|
|
entity_type_filter,
|
|
dependency_type_filter,
|
|
)
|
|
for node in start_entities
|
|
]
|
|
rtn_str = '\n\n'.join(rtns)
|
|
|
|
if hints.strip():
|
|
rtn_str += '\n\n' + hints
|
|
return rtn_str.strip()
|
|
|
|
|
|
__all__ = [
|
|
'get_entity_contents',
|
|
'search_code_snippets',
|
|
'explore_tree_structure',
|
|
]
|