Skip to content

Commit 495216b

Browse files
Aisha630Aisha630
committed
feat: Add fuzzy search capabilites to find code tool
1 parent a95bd25 commit 495216b

2 files changed

Lines changed: 86 additions & 40 deletions

File tree

‎src/codegraphcontext/server.py‎

Lines changed: 19 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@
2424
from .tools.package_resolver import get_local_package_path
2525
from .utils.debug_log import debug_log, info_logger, error_logger, warning_logger, debug_logger
2626

27+
DEFAULT_EDIT_DISTANCE = 2
28+
DEFAULT_FUZZY_SEARCH = False
29+
2730
class MCPServer:
2831
"""
2932
The main MCP Server class.
@@ -105,15 +108,16 @@ def _init_tools(self):
105108
"description": "List all background jobs and their current status.",
106109
"inputSchema": {"type": "object", "properties": {}}
107110
},
108-
"find_code": {
111+
"find_code": {
109112
"name": "find_code",
110113
"description": "Find relevant code snippets related to a keyword (e.g., function name, class name, or content).",
111114
"inputSchema": {
112115
"type": "object",
113-
"properties": { "query": {"type": "string", "description": "Keyword or phrase to search for"} },
116+
"properties": { "query": {"type": "string", "description": "Keyword or phrase to search for"}, "fuzzy_search": {"type": "boolean", "description": "Whether to use fuzzy search", "default": False}, "edit_distance": {"type": "number", "description": "Edit distance for fuzzy search (between 0-2)", "default": 2}},
114117
"required": ["query"]
115118
}
116119
},
120+
117121
"analyze_code_relationships": {
118122
"name": "analyze_code_relationships",
119123
"description": "Analyze code relationships like 'who calls this function' or 'class hierarchy'. Supported query types include: find_callers, find_callees, find_all_callers, find_all_callees, find_importers, who_modifies, class_hierarchy, overrides, dead_code, call_chain, module_deps, variable_scope, find_complexity, find_functions_by_argument, find_functions_by_decorator.",
@@ -683,15 +687,24 @@ def analyze_code_relationships_tool(self, **args) -> Dict[str, Any]:
683687
except Exception as e:
684688
debug_log(f"Error analyzing relationships: {str(e)}")
685689
return {"error": f"Failed to analyze relationships: {str(e)}"}
690+
691+
@staticmethod
692+
def _normalize(text: str) -> str:
693+
return text.lower().replace("_", " ").strip()
686694

687695
def find_code_tool(self, **args) -> Dict[str, Any]:
688696
"""Tool to find relevant code snippets"""
689697
query = args.get("query")
690-
691-
try:
692-
debug_log(f"Finding code for query: {query}")
693-
results = self.code_finder.find_related_code(query)
698+
fuzzy_search = args.get("fuzzy_search", DEFAULT_FUZZY_SEARCH)
699+
edit_distance = args.get("edit_distance", DEFAULT_EDIT_DISTANCE)
700+
701+
if fuzzy_search:
702+
query = self._normalize(query)
694703

704+
try:
705+
debug_log(f"Finding code for query: {query} with fuzzy_search={fuzzy_search}, edit_distance={edit_distance}")
706+
results = self.code_finder.find_related_code(query, fuzzy_search, edit_distance)
707+
695708
return {"success": True, "query": query, "results": results}
696709

697710
except Exception as e:

‎src/codegraphcontext/tools/code_finder.py‎

Lines changed: 67 additions & 34 deletions
Original file line numberDiff line numberDiff line change
@@ -15,33 +15,61 @@ def __init__(self, db_manager: DatabaseManager):
1515
self.db_manager = db_manager
1616
self.driver = self.db_manager.get_driver()
1717

18-
def find_by_function_name(self, search_term: str) -> List[Dict]:
18+
def find_by_function_name(self, search_term: str, fuzzy_search: bool) -> List[Dict]:
1919
"""Find functions by name matching using the full-text index."""
20-
with self.driver.session() as session:
21-
result = session.run("""
22-
CALL db.index.fulltext.queryNodes("code_search_index", $search_term) YIELD node, score
23-
WITH node, score
24-
WHERE node:Function AND node.name CONTAINS $search_term
25-
RETURN node.name as name, node.file_path as file_path, node.line_number as line_number,
26-
node.source as source, node.docstring as docstring, node.is_dependency as is_dependency
27-
ORDER BY score DESC
28-
LIMIT 20
29-
""", search_term=search_term)
30-
return [dict(record) for record in result]
31-
32-
def find_by_class_name(self, search_term: str) -> List[Dict]:
20+
if fuzzy_search:
21+
with self.driver.session() as session:
22+
formatted_search_term = f"name:{search_term}"
23+
result = session.run("""
24+
CALL db.index.fulltext.queryNodes("code_search_index", $search_term) YIELD node, score
25+
WITH node, score
26+
WHERE node:Function
27+
RETURN node.name as name, node.file_path as file_path, node.line_number as line_number,
28+
node.source as source, node.docstring as docstring, node.is_dependency as is_dependency
29+
ORDER BY score DESC
30+
LIMIT 20
31+
""", search_term=formatted_search_term)
32+
return [dict(record) for record in result]
33+
else:
34+
with self.driver.session() as session:
35+
result = session.run("""
36+
CALL db.index.fulltext.queryNodes("code_search_index", $search_term) YIELD node, score
37+
WITH node, score
38+
WHERE node:Function AND node.name CONTAINS $search_term
39+
RETURN node.name as name, node.file_path as file_path, node.line_number as line_number,
40+
node.source as source, node.docstring as docstring, node.is_dependency as is_dependency
41+
ORDER BY score DESC
42+
LIMIT 20
43+
""", search_term=search_term)
44+
return [dict(record) for record in result]
45+
46+
def find_by_class_name(self, search_term: str, fuzzy_search: bool) -> List[Dict]:
3347
"""Find classes by name matching using the full-text index."""
34-
with self.driver.session() as session:
35-
result = session.run("""
36-
CALL db.index.fulltext.queryNodes("code_search_index", $search_term) YIELD node, score
37-
WITH node, score
38-
WHERE node:Class AND node.name CONTAINS $search_term
39-
RETURN node.name as name, node.file_path as file_path, node.line_number as line_number,
40-
node.source as source, node.docstring as docstring, node.is_dependency as is_dependency
41-
ORDER BY score DESC
42-
LIMIT 20
43-
""", search_term=search_term)
44-
return [dict(record) for record in result]
48+
if fuzzy_search:
49+
with self.driver.session() as session:
50+
formatted_search_term = f"name:{search_term}"
51+
result = session.run("""
52+
CALL db.index.fulltext.queryNodes("code_search_index", $search_term) YIELD node, score
53+
WITH node, score
54+
WHERE node:Class
55+
RETURN node.name as name, node.file_path as file_path, node.line_number as line_number,
56+
node.source as source, node.docstring as docstring, node.is_dependency as is_dependency
57+
ORDER BY score DESC
58+
LIMIT 20
59+
""", search_term=formatted_search_term)
60+
return [dict(record) for record in result]
61+
else:
62+
with self.driver.session() as session:
63+
result = session.run("""
64+
CALL db.index.fulltext.queryNodes("code_search_index", $search_term) YIELD node, score
65+
WITH node, score
66+
WHERE node:Class AND node.name CONTAINS $search_term
67+
RETURN node.name as name, node.file_path as file_path, node.line_number as line_number,
68+
node.source as source, node.docstring as docstring, node.is_dependency as is_dependency
69+
ORDER BY score DESC
70+
LIMIT 20
71+
""", search_term=search_term)
72+
return [dict(record) for record in result]
4573

4674
def find_by_variable_name(self, search_term: str) -> List[Dict]:
4775
"""Find variables by name matching"""
@@ -56,7 +84,7 @@ def find_by_variable_name(self, search_term: str) -> List[Dict]:
5684
""", search_term=search_term, regex_pattern=f"(?i).*{re.escape(search_term)}.*")
5785

5886
return [dict(record) for record in result]
59-
87+
6088
def find_by_content(self, search_term: str) -> List[Dict]:
6189
"""Find code by content matching in source or docstrings using the full-text index."""
6290
with self.driver.session() as session:
@@ -77,15 +105,20 @@ def find_by_content(self, search_term: str) -> List[Dict]:
77105
LIMIT 20
78106
""", search_term=search_term)
79107
return [dict(record) for record in result]
80-
81-
def find_related_code(self, user_query: str) -> Dict[str, Any]:
108+
109+
def find_related_code(self, user_query: str, fuzzy_search: bool, edit_distance: int) -> Dict[str, Any]:
82110
"""Find code related to a query using multiple search strategies"""
111+
if fuzzy_search:
112+
user_query_normalized = f"{" ".join(map(lambda x: f'{x}~{edit_distance}', user_query.split(' ')))}"
113+
else:
114+
user_query_normalized = user_query
115+
83116
results = {
84-
"query": user_query,
85-
"functions_by_name": self.find_by_function_name(user_query),
86-
"classes_by_name": self.find_by_class_name(user_query),
87-
"variables_by_name": self.find_by_variable_name(user_query),
88-
"content_matches": self.find_by_content(user_query)
117+
"query": user_query_normalized,
118+
"functions_by_name": self.find_by_function_name(user_query_normalized, fuzzy_search),
119+
"classes_by_name": self.find_by_class_name(user_query_normalized, fuzzy_search),
120+
"variables_by_name": self.find_by_variable_name(user_query), # no fuzzy for variables as they are not using full-text index
121+
"content_matches": self.find_by_content(user_query_normalized)
89122
}
90123

91124
all_results = []
@@ -763,4 +796,4 @@ def list_indexed_repositories(self) -> List[Dict]:
763796
RETURN r.name as name, r.path as path, r.is_dependency as is_dependency
764797
ORDER BY r.name
765798
""")
766-
return [dict(record) for record in result]
799+
return [dict(record) for record in result]

0 commit comments

Comments
 (0)