"""Module for finding relevant code snippets and analyzing relationships."""

def __init__(self, db_manager: DatabaseManager):

self.db_manager = db_manager

self.driver = self.db_manager.get_driver()

self._is_falkordb = getattr(db_manager, 'get_backend_type', lambda: 'neo4j')() != 'neo4j'

def format_query(self, find_by: Literal["Class", "Function"], fuzzy_search:bool, repo_path: Optional[str] = None) -> str:

"""Format the search query based on the search type and fuzzy search settings."""

repo_filter = "AND node.path STARTS WITH $repo_path" if repo_path else ""

if self._is_falkordb:

# FalkorDB does not support CALL db.idx.fulltext.queryNodes.

# Fall back to a pure Cypher CONTAINS/toLower match on node name.

name_filter = "toLower(node.name) CONTAINS toLower($search_term)"

return f"""

return f"""

def find_by_function_name(self, search_term: str, fuzzy_search: bool, repo_path: Optional[str] = None) -> List[Dict]:

"""Find functions by name matching."""

with self.driver.session() as session:

if not fuzzy_search:

# Use simple match for exact search to avoid fulltext index dependency

result = session.run(f"""

""", name=search_term, repo_path=repo_path)

return result.data()

# Fuzzy search using fulltext index (Neo4j) or CONTAINS fallback (FalkorDB)

# On FalkorDB, format_query uses CONTAINS so we pass the raw term; on Neo4j

# we need the Lucene field-selector prefix.

formatted_search_term = search_term if self._is_falkordb else f"name:{search_term}"

result = session.run(self.format_query("Function", fuzzy_search, repo_path), search_term=formatted_search_term, repo_path=repo_path)

return result.data()

def find_by_class_name(self, search_term: str, fuzzy_search: bool, repo_path: Optional[str] = None) -> List[Dict]:

"""Find classes by name matching."""

with self.driver.session() as session:

if not fuzzy_search:

# Use simple match for exact search to avoid fulltext index dependency

result = session.run(f"""

""", name=search_term, repo_path=repo_path)

return result.data()

# Fuzzy search using fulltext index (Neo4j) or CONTAINS fallback (FalkorDB)

# On FalkorDB, format_query uses CONTAINS so we pass the raw term; on Neo4j

# we need the Lucene field-selector prefix.

formatted_search_term = search_term if self._is_falkordb else f"name:{search_term}"

result = session.run(self.format_query("Class", fuzzy_search, repo_path), search_term=formatted_search_term, repo_path=repo_path)

return result.data()

def find_by_variable_name(self, search_term: str, repo_path: Optional[str] = None) -> List[Dict]:

"""Find variables by name matching"""

with self.driver.session() as session:

result = session.run(f"""

""", search_term=search_term, repo_path=repo_path)

return result.data()

def find_by_content(self, search_term: str, repo_path: Optional[str] = None) -> List[Dict]:

"""Find code by content matching in source or docstrings using the full-text index."""

if self._is_falkordb:

return self._find_by_content_falkordb(search_term, repo_path)

with self.driver.session() as session:

result = session.run(f"""

""", search_term=search_term, repo_path=repo_path)

return result.data()

def _find_by_content_falkordb(self, search_term: str, repo_path: Optional[str] = None) -> List[Dict]:

"""FalkorDB-compatible content search using pure Cypher CONTAINS matching.

all_results = []

with self.driver.session() as session:

repo_filter = "AND node.path STARTS WITH $repo_path" if repo_path else ""

for label, type_name in [('Function', 'function'), ('Class', 'class')]:

try:

result = session.run(f"""

""", search_term=search_term, repo_path=repo_path)

all_results.extend(result.data())

except Exception:

logger.debug(f"FalkorDB content query failed for label {label}", exc_info=True)

return all_results[:20]

def find_by_module_name(self, search_term: str) -> List[Dict]:

"""Find modules by name matching"""

with self.driver.session() as session:

result = session.run("""

""", search_term=search_term)

return result.data()

def find_imports(self, search_term: str) -> List[Dict]:

"""Find imported symbols (aliases or original names)."""

with self.driver.session() as session:

result = session.run("""

""", search_term=search_term)

return result.data()

def find_related_code(self, user_query: str, fuzzy_search: bool, edit_distance: int, repo_path: Optional[str] = None) -> Dict[str, Any]:

"""Find code related to a query using multiple search strategies"""

# FalkorDB does not support Lucene-style fuzzy edit-distance syntax (e.g. term~2).

# On FalkorDB, always use the plain query so that the CONTAINS-based fallbacks work.

if fuzzy_search and self._is_falkordb:

logger.debug("FalkorDB backend: ignoring fuzzy edit-distance normalisation; using plain CONTAINS search.")

fuzzy_search = False

if fuzzy_search:

user_query_normalized = " ".join(map(lambda x: f"{x}~{edit_distance}", user_query.split(" ")))

else:

user_query_normalized = user_query

results: Dict[str, Any] = {

"query": user_query_normalized,

"functions_by_name": self.find_by_function_name(user_query_normalized, fuzzy_search, repo_path),

"classes_by_name": self.find_by_class_name(user_query_normalized, fuzzy_search, repo_path),

"variables_by_name": self.find_by_variable_name(user_query, repo_path), # no fuzzy for variables as they are not using full-text index

"content_matches": self.find_by_content(user_query_normalized, repo_path)

}

all_results: List[Dict[str, Any]] = []

for func in results["functions_by_name"]:

func["search_type"] = "function_name"

func["relevance_score"] = 0.9 if not func["is_dependency"] else 0.7

all_results.append(func)

for cls in results["classes_by_name"]:

cls["search_type"] = "class_name"

cls["relevance_score"] = 0.8 if not cls["is_dependency"] else 0.6

all_results.append(cls)

for var in results["variables_by_name"]:

var["search_type"] = "variable_name"

var["relevance_score"] = 0.7 if not var["is_dependency"] else 0.5

all_results.append(var)

for content in results["content_matches"]:

content["search_type"] = "content"

content["relevance_score"] = 0.6 if not content["is_dependency"] else 0.4

all_results.append(content)

all_results.sort(key=lambda x: x["relevance_score"], reverse=True)

results["ranked_results"] = all_results[:15]

results["total_matches"] = len(all_results)

return results

def find_functions_by_argument(self, argument_name: str, path: Optional[str] = None, repo_path: Optional[str] = None) -> List[Dict]:

"""Find functions that take a specific argument name."""

with self.driver.session() as session:

repo_filter = "AND f.path STARTS WITH $repo_path" if repo_path else ""

if path:

query = f"""

result = session.run(query, argument_name=argument_name, path=path, repo_path=repo_path)

else:

query = f"""

result = session.run(query, argument_name=argument_name, repo_path=repo_path)

return result.data()

def find_functions_by_decorator(self, decorator_name: str, path: Optional[str] = None, repo_path: Optional[str] = None) -> List[Dict]:

"""Find functions that have a specific decorator applied to them."""

with self.driver.session() as session:

repo_filter = "AND f.path STARTS WITH $repo_path" if repo_path else ""

if path:

query = f"""

result = session.run(query, decorator_name=decorator_name, path=path, repo_path=repo_path)

else:

query = f"""

result = session.run(query, decorator_name=decorator_name, repo_path=repo_path)

return result.data()

def who_calls_function(self, function_name: str, path: Optional[str] = None, repo_path: Optional[str] = None) -> List[Dict]:

"""Find what functions call a specific function using CALLS relationships with improved matching"""

with self.driver.session() as session:

repo_filter = "AND caller.path STARTS WITH $repo_path" if repo_path else ""

if path:

result = session.run(f"""

""", function_name=function_name, path=path, repo_path=repo_path)

results = result.data()

if not results:

result = session.run(f"""

""", function_name=function_name, repo_path=repo_path)

results = result.data()

else:

result = session.run(f"""

""", function_name=function_name, repo_path=repo_path)

results = result.data()

return results

def what_does_function_call(self, function_name: str, path: Optional[str] = None, repo_path: Optional[str] = None) -> List[Dict]:

"""Find what functions a specific function calls using CALLS relationships"""

with self.driver.session() as session:

if path:

# Convert path to absolute path

absolute_file_path = str(Path(path).resolve())

result = session.run(f"""

""", function_name=function_name, absolute_file_path=absolute_file_path, repo_path=repo_path)

else:

result = session.run(f"""

""", function_name=function_name, repo_path=repo_path)

return result.data()

def who_imports_module(self, module_name: str, repo_path: Optional[str] = None) -> List[Dict]:

"""Find what files import a specific module using IMPORTS relationships"""

with self.driver.session() as session:

repo_filter = "AND file.path STARTS WITH $repo_path" if repo_path else ""

result = session.run(f"""

""", module_name=module_name, repo_path=repo_path)

return result.data()

def who_modifies_variable(self, variable_name: str, repo_path: Optional[str] = None) -> List[Dict]:

"""Find what functions contain or modify a specific variable"""

with self.driver.session() as session:

repo_filter = "AND container.path STARTS WITH $repo_path" if repo_path else ""

result = session.run(f"""

""", variable_name=variable_name, repo_path=repo_path)

return result.data()

def find_class_hierarchy(self, class_name: str, path: Optional[str] = None, repo_path: Optional[str] = None) -> Dict[str, Any]:

"""Find class inheritance relationships using INHERITS relationships"""

with self.driver.session() as session:

repo_filter = "AND parent.path STARTS WITH $repo_path" if repo_path else ""

if path:

match_clause = "MATCH (child:Class {name: $class_name, path: $path})"

else:

match_clause = "MATCH (child:Class {name: $class_name})"

parents_query = f"""

parents_result = session.run(parents_query, class_name=class_name, path=path, repo_path=repo_path)

repo_filter_child = "AND grandchild.path STARTS WITH $repo_path" if repo_path else ""

children_query = f"""

children_result = session.run(children_query, class_name=class_name, path=path, repo_path=repo_path)

repo_filter_method = "WHERE method.path STARTS WITH $repo_path" if repo_path else ""

methods_query = f"""

methods_result = session.run(methods_query, class_name=class_name, path=path, repo_path=repo_path)

return {

"class_name": class_name,

"parent_classes": parents_result.data(),

"child_classes": children_result.data(),

"methods": methods_result.data()

}

def find_function_overrides(self, function_name: str, repo_path: Optional[str] = None) -> List[Dict]:

"""Find all implementations of a function across different classes"""

with self.driver.session() as session:

repo_filter = "AND class.path STARTS WITH $repo_path" if repo_path else ""

result = session.run(f"""

""", function_name=function_name, repo_path=repo_path)

return result.data()

def find_dead_code(self, exclude_decorated_with: Optional[List[str]] = None, repo_path: Optional[str] = None) -> Dict[str, Any]:

"""Find potentially unused functions (not called by other functions in the project), optionally excluding those with specific decorators."""

if exclude_decorated_with is None:

exclude_decorated_with = []

with self.driver.session() as session:

repo_filter = "AND func.path STARTS WITH $repo_path" if repo_path else ""

decorator_filter = "AND ALL(decorator_name IN $exclude_decorated_with WHERE NOT decorator_name IN func.decorators)" if exclude_decorated_with else ""

func_ignore = cypher_path_not_under_ignore_dirs("func.path")

caller_ignore = cypher_path_not_under_ignore_dirs("caller.path")

query = f"""

params = {}

if repo_path:

params["repo_path"] = repo_path

if exclude_decorated_with:

params["exclude_decorated_with"] = exclude_decorated_with

result = session.run(query, **params)

return {

"potentially_unused_functions": result.data(),

"note": "These functions might be unused, but could be entry points, callbacks, or called dynamically"

}

def find_all_callers(self, function_name: str, path: Optional[str] = None, repo_path: Optional[str] = None) -> List[Dict]:

"""Find all direct and indirect callers of a specific function."""

with self.driver.session() as session:

repo_filter = "AND f.path STARTS WITH $repo_path" if repo_path else ""

if path:

# KùzuDB-compatible: Use anonymous end node and filter with WHERE

query = f"""

result = session.run(query, function_name=function_name, path=path, repo_path=repo_path)

else:

# KùzuDB-compatible: Use anonymous end node and filter with WHERE

query = f"""

result = session.run(query, function_name=function_name, repo_path=repo_path)

return result.data()

def find_all_callees(self, function_name: str, path: Optional[str] = None, repo_path: Optional[str] = None) -> List[Dict]:

"""Find all direct and indirect callees of a specific function."""

with self.driver.session() as session:

repo_filter = "WHERE f.path STARTS WITH $repo_path" if repo_path else ""

if path:

# KùzuDB-compatible: Use anonymous end node and extract from path

query = f"""

result = session.run(query, function_name=function_name, path=path, repo_path=repo_path)

else:

# KùzuDB-compatible: Use anonymous end node and extract from path

query = f"""

result = session.run(query, function_name=function_name, repo_path=repo_path)

return result.data()

def find_function_call_chain(self, start_function: str, end_function: str, max_depth: int = 5, start_file: Optional[str] = None, end_file: Optional[str] = None, repo_path: Optional[str] = None) -> List[Dict]:

"""Find call chains between two functions"""

with self.driver.session() as session:

# Build match clauses based on whether files are specified

start_props = "{name: $start_function" + (", path: $start_file}" if start_file else "}")

end_props = "{name: $end_function" + (", path: $end_file}" if end_file else "}")

# KùzuDB-compatible: Use anonymous end node and filter

repo_filter = "WHERE 1=1 AND (start.path IS NULL OR start.path STARTS WITH $repo_path) AND (end_target.path IS NULL OR end_target.path STARTS WITH $repo_path)" if repo_path else ""

query = f"""

# Prepare parameters

params = {

"start_function": start_function,

"end_function": end_function,

"start_file": start_file,

"end_file": end_file,

"repo_path": repo_path

}

result = session.run(query, **params)

# Post-process Node/Rel objects into plain dicts so CLI output stays stable

rows = result.data()

transformed: List[Dict[str, Any]] = []

for row in rows:

func_nodes = row.get("function_nodes") or []

rel_nodes = row.get("call_nodes") or []

chain_len = row.get("chain_length", 0)

function_chain = []

for n in func_nodes:

# Depending on KùzuDB + driver wrapping, list elements can arrive

# either as Node/Rel objects or already-materialized dicts.

if isinstance(n, dict):

props = n

else:

props = None

try:

props = n.get_properties()

except Exception:

props = getattr(n, "properties", None)

if props is None:

props = {}

function_chain.append(

{

"name": props.get("name"),

"path": props.get("path"),

"line_number": props.get("line_number"),

"is_dependency": props.get("is_dependency"),

}

)

call_details = []

for r in rel_nodes:

if isinstance(r, dict):

props = r

else:

props = None

try:

props = r.get_properties()

except Exception:

props = getattr(r, "properties", None)

if props is None:

props = {}

call_details.append(

{

"call_line": props.get("line_number"),

"args": props.get("args"),

"full_call_name": props.get("full_call_name"),

}

)

transformed.append(

{

"function_chain": function_chain,

"call_details": call_details,

"chain_length": chain_len,

}

)

return transformed

def find_by_type(self, element_type: str, limit: int = 50) -> List[Dict]:

"""Find all elements of a specific type (Function, Class, File, Module)."""

# Map input type to node label

type_map = {

"function": "Function",

"class": "Class",

"file": "File",

"module": "Module"

}

label = type_map.get(element_type.lower())

if not label:

return []

with self.driver.session() as session:

if label == "File":

query = f"""

elif label == "Module":

query = f"""

else:

query = f"""

result = session.run(query, limit=limit)

return result.data()

def find_module_dependencies(self, module_name: str, repo_path: Optional[str] = None) -> Dict[str, Any]:

"""Find all dependencies and dependents of a module"""

with self.driver.session() as session:

repo_filter = "AND file.path STARTS WITH $repo_path" if repo_path else ""

# Find files that import this module (who imports this module)

importers_result = session.run(f"""

""", module_name=module_name, repo_path=repo_path)

# Find modules that are imported by files that also import the target module

# This helps understand what this module is typically used with

imports_result = session.run(f"""

""", module_name=module_name, repo_path=repo_path)

return {

"module_name": module_name,

"importers": importers_result.data(),

"imports": imports_result.data()

}

def find_variable_usage_scope(self, variable_name: str, path: Optional[str] = None, repo_path: Optional[str] = None) -> Dict[str, Any]:

"""Find the scope and usage patterns of a variable, optional file path filtering"""

with self.driver.session() as session:

repo_filter = "AND var.path STARTS WITH $repo_path" if repo_path else ""

if path:

variable_instances = session.run(f"""

""", variable_name=variable_name, path=path, repo_path=repo_path)

else:

variable_instances = session.run(f"""

""", variable_name=variable_name, repo_path=repo_path)

return {

"variable_name": variable_name,

"instances": variable_instances.data()

}

def analyze_code_relationships(self, query_type: str, target: str, context: Optional[str] = None, repo_path: Optional[str] = None) -> Dict[str, Any]:

"""Main method to analyze different types of code relationships with fixed return types"""

query_type = query_type.lower().strip()

try:

if query_type == "find_callers":

results = self.who_calls_function(target, context, repo_path=repo_path)

return {

"query_type": "find_callers", "target": target, "context": context, "results": results,

"summary": f"Found {len(results)} functions that call '{target}'"

}

elif query_type == "find_callees":

results = self.what_does_function_call(target, context, repo_path=repo_path)

return {

"query_type": "find_callees", "target": target, "context": context, "results": results,

"summary": f"Function '{target}' calls {len(results)} other functions"

}

elif query_type == "find_importers":

results = self.who_imports_module(target, repo_path=repo_path)

return {

"query_type": "find_importers", "target": target, "results": results,

"summary": f"Found {len(results)} files that import '{target}'"

}

elif query_type == "find_functions_by_argument":

results = self.find_functions_by_argument(target, context, repo_path=repo_path)

return {

"query_type": "find_functions_by_argument", "target": target, "context": context, "results": results,

"summary": f"Found {len(results)} functions that take '{target}' as an argument"

}

elif query_type == "find_functions_by_decorator":

results = self.find_functions_by_decorator(target, context, repo_path=repo_path)

return {

"query_type": "find_functions_by_decorator", "target": target, "context": context, "results": results,

"summary": f"Found {len(results)} functions decorated with '{target}'"

}

elif query_type in ["who_modifies", "modifies", "mutations", "changes", "variable_usage"]:

results = self.who_modifies_variable(target, repo_path=repo_path)

return {

"query_type": "who_modifies", "target": target, "results": results,

"summary": f"Found {len(results)} containers that hold variable '{target}'"

}

elif query_type in ["class_hierarchy", "inheritance", "extends"]:

results = self.find_class_hierarchy(target, context, repo_path=repo_path)

return {

"query_type": "class_hierarchy", "target": target, "results": results,

"summary": f"Class '{target}' has {len(results['parent_classes'])} parents, {len(results['child_classes'])} children, and {len(results['methods'])} methods"

}

elif query_type in ["overrides", "implementations", "polymorphism"]:

results = self.find_function_overrides(target, repo_path=repo_path)

return {

"query_type": "overrides", "target": target, "results": results,

"summary": f"Found {len(results)} implementations of function '{target}'"

}

elif query_type in ["dead_code", "unused", "unreachable"]:

results = self.find_dead_code(repo_path=repo_path)

return {

"query_type": "dead_code", "results": results,

"summary": f"Found {len(results['potentially_unused_functions'])} potentially unused functions"

}

elif query_type == "find_complexity":

limit = int(context) if context and context.isdigit() else 10

results = self.find_most_complex_functions(limit, repo_path=repo_path)

return {

"query_type": "find_complexity", "limit": limit, "results": results,

"summary": f"Found the top {len(results)} most complex functions"

}

elif query_type == "find_all_callers":

results = self.find_all_callers(target, context, repo_path=repo_path)

return {

"query_type": "find_all_callers", "target": target, "context": context, "results": results,

"summary": f"Found {len(results)} direct and indirect callers of '{target}'"

}

elif query_type == "find_all_callees":

results = self.find_all_callees(target, context, repo_path=repo_path)

return {

"query_type": "find_all_callees", "target": target, "context": context, "results": results,

"summary": f"Found {len(results)} direct and indirect callees of '{target}'"

}

elif query_type in ["call_chain", "path", "chain"]:

if '->' in target:

start_func, end_func = target.split('->', 1)

# max_depth can be passed as context, default to 5 if not provided or invalid

max_depth = int(context) if context and context.isdigit() else 5

results = self.find_function_call_chain(start_func.strip(), end_func.strip(), max_depth, repo_path=repo_path)

return {

"query_type": "call_chain", "target": target, "results": results,

"summary": f"Found {len(results)} call chains from '{start_func.strip()}' to '{end_func.strip()}' (max depth: {max_depth})"

}

else:

return {

"error": "For call_chain queries, use format 'start_function->end_function'",

"example": "main->process_data"

}

elif query_type in ["module_deps", "module_dependencies", "module_usage"]:

results = self.find_module_dependencies(target, repo_path=repo_path)

return {

"query_type": "module_dependencies", "target": target, "results": results,

"summary": f"Module '{target}' is imported by {len(results['imported_by_files'])} files"

}

elif query_type in ["variable_scope", "var_scope", "variable_usage_scope"]:

results = self.find_variable_usage_scope(target, repo_path=repo_path)

return {

"query_type": "variable_scope", "target": target, "results": results,