diff --git a/graphify/__main__.py b/graphify/__main__.py index d8f87827..256e0574 100644 --- a/graphify/__main__.py +++ b/graphify/__main__.py @@ -1016,6 +1016,7 @@ def main() -> None: print(" --no-viz skip graph.html generation (useful for >5000 node graphs / CI)") print(" query \"\" BFS traversal of graph.json for a question") print(" --dfs use depth-first instead of breadth-first") + print(" --context C explicit edge-context filter (repeatable)") print(" --budget N cap output at N tokens (default 2000)") print(" --graph path to graph.json (default graphify-out/graph.json)") print(" save-result save a Q&A result to graphify-out/memory/ for graph feedback loop") @@ -1209,15 +1210,16 @@ def main() -> None: sys.exit(1) elif cmd == "query": if len(sys.argv) < 3: - print("Usage: graphify query \"\" [--dfs] [--budget N] [--graph path]", file=sys.stderr) + print("Usage: graphify query \"\" [--dfs] [--context C] [--budget N] [--graph path]", file=sys.stderr) sys.exit(1) - from graphify.serve import _score_nodes, _bfs, _dfs, _subgraph_to_text + from graphify.serve import _query_graph_text from graphify.security import sanitize_label from networkx.readwrite import json_graph question = sys.argv[2] use_dfs = "--dfs" in sys.argv budget = 2000 graph_path = "graphify-out/graph.json" + context_filters: list[str] = [] args = sys.argv[3:] i = 0 while i < len(args): @@ -1235,6 +1237,12 @@ def main() -> None: print(f"error: --budget must be an integer", file=sys.stderr) sys.exit(1) i += 1 + elif args[i] == "--context" and i + 1 < len(args): + context_filters.append(args[i + 1]) + i += 2 + elif args[i].startswith("--context="): + context_filters.append(args[i].split("=", 1)[1]) + i += 1 elif args[i] == "--graph" and i + 1 < len(args): graph_path = args[i + 1]; i += 2 else: @@ -1257,14 +1265,16 @@ def main() -> None: except Exception as exc: print(f"error: could not load graph: {exc}", file=sys.stderr) sys.exit(1) - terms = [t.lower() for t in question.split() if len(t) > 2] - scored = _score_nodes(G, terms) - if not scored: - print("No matching nodes found.") - sys.exit(0) - start = [nid for _, nid in scored[:5]] - nodes, edges = (_dfs if use_dfs else _bfs)(G, start, depth=2) - print(_subgraph_to_text(G, nodes, edges, token_budget=budget)) + print( + _query_graph_text( + G, + question, + mode="dfs" if use_dfs else "bfs", + depth=2, + token_budget=budget, + context_filters=context_filters, + ) + ) elif cmd == "save-result": # graphify save-result --question Q --answer A --type T [--nodes N1 N2 ...] import argparse as _ap diff --git a/graphify/extract.py b/graphify/extract.py index 9ba9ed56..94ddd200 100644 --- a/graphify/extract.py +++ b/graphify/extract.py @@ -150,6 +150,7 @@ def _import_python(node, source: bytes, file_nid: str, stem: str, edges: list, s "source": file_nid, "target": tgt_nid, "relation": "imports", + "context": "import", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", @@ -174,6 +175,7 @@ def _import_python(node, source: bytes, file_nid: str, stem: str, edges: list, s "source": file_nid, "target": tgt_nid, "relation": "imports_from", + "context": "import", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", @@ -221,6 +223,7 @@ def _import_js(node, source: bytes, file_nid: str, stem: str, edges: list, str_p "source": file_nid, "target": tgt_nid, "relation": "imports_from", + "context": "import", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", @@ -285,6 +288,7 @@ def _import_java(node, source: bytes, file_nid: str, stem: str, edges: list, str "source": file_nid, "target": tgt_nid, "relation": "imports", + "context": "import", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", @@ -304,6 +308,7 @@ def _import_c(node, source: bytes, file_nid: str, stem: str, edges: list, str_pa "source": file_nid, "target": tgt_nid, "relation": "imports", + "context": "import", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", @@ -323,6 +328,7 @@ def _import_csharp(node, source: bytes, file_nid: str, stem: str, edges: list, s "source": file_nid, "target": tgt_nid, "relation": "imports", + "context": "import", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", @@ -342,6 +348,7 @@ def _import_kotlin(node, source: bytes, file_nid: str, stem: str, edges: list, s "source": file_nid, "target": tgt_nid, "relation": "imports", + "context": "import", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", @@ -357,6 +364,7 @@ def _import_kotlin(node, source: bytes, file_nid: str, stem: str, edges: list, s "source": file_nid, "target": tgt_nid, "relation": "imports", + "context": "import", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", @@ -376,6 +384,7 @@ def _import_scala(node, source: bytes, file_nid: str, stem: str, edges: list, st "source": file_nid, "target": tgt_nid, "relation": "imports", + "context": "import", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", @@ -395,6 +404,7 @@ def _import_php(node, source: bytes, file_nid: str, stem: str, edges: list, str_ "source": file_nid, "target": tgt_nid, "relation": "imports", + "context": "import", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", @@ -696,6 +706,7 @@ def _import_lua(node, source: bytes, file_nid: str, stem: str, edges: list, str_ "source": file_nid, "target": module_name, "relation": "imports", + "context": "import", "confidence": "EXTRACTED", "confidence_score": 1.0, "source_file": str_path, @@ -730,6 +741,7 @@ def _import_swift(node, source: bytes, file_nid: str, stem: str, edges: list, st "source": file_nid, "target": tgt_nid, "relation": "imports", + "context": "import", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{node.start_point[0] + 1}", @@ -756,6 +768,27 @@ def _import_vbnet(node, source: bytes, file_nid: str, stem: str, edges: list, st }) +def _read_csharp_type_name(node, source: bytes) -> str | None: + """Resolve a readable C# type name from a field/type node.""" + if node is None: + return None + if node.type in ("identifier", "predefined_type"): + return _read_text(node, source) + if node.type == "qualified_name": + return _read_text(node, source).split(".")[-1] + if node.type == "generic_name": + name_node = node.child_by_field_name("name") + if name_node is not None: + return _read_text(name_node, source) + for child in node.children: + if not child.is_named: + continue + name = _read_csharp_type_name(child, source) + if name: + return name + return None + + _SWIFT_CONFIG = LanguageConfig( ts_module="tree_sitter_swift", class_types=frozenset({"class_declaration", "protocol_declaration"}), @@ -833,8 +866,9 @@ def _extract_generic(path: Path, config: LanguageConfig) -> dict: }) def add_edge(src: str, tgt: str, relation: str, line: int, - confidence: str = "EXTRACTED", weight: float = 1.0) -> None: - edges.append({ + confidence: str = "EXTRACTED", weight: float = 1.0, + context: str | None = None) -> None: + edge = { "source": src, "target": tgt, "relation": relation, @@ -842,7 +876,19 @@ def _extract_generic(path: Path, config: LanguageConfig) -> dict: "source_file": str_path, "source_location": f"L{line}", "weight": weight, - }) + } + if context: + edge["context"] = context + edges.append(edge) + + def ensure_named_node(name: str, line: int) -> str: + nid = _make_id(stem, name) + if nid in seen_ids: + return nid + nid = _make_id(name) + if nid not in seen_ids: + add_node(nid, name, line) + return nid file_nid = _make_id(str(path)) add_node(file_nid, path.name, 1) @@ -1078,6 +1124,23 @@ def _extract_generic(path: Path, config: LanguageConfig) -> dict: break return + if (config.ts_module == "tree_sitter_c_sharp" + and t == "field_declaration" + and parent_class_nid): + type_node = node.child_by_field_name("type") + if type_node is None: + for child in node.children: + if child.type == "variable_declaration": + type_node = child.child_by_field_name("type") + if type_node is not None: + break + type_name = _read_csharp_type_name(type_node, source) + if type_name: + line = node.start_point[0] + 1 + add_edge(parent_class_nid, ensure_named_node(type_name, line), + "references", line, context="field") + return + # Function types if t in config.function_types: # Swift deinit/subscript have no name field — resolve before generic fallback @@ -1304,6 +1367,7 @@ def _extract_generic(path: Path, config: LanguageConfig) -> dict: "source": caller_nid, "target": tgt_nid, "relation": "calls", + "context": "call", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{line}", @@ -2032,8 +2096,9 @@ def extract_julia(path: Path) -> dict: }) def add_edge(src: str, tgt: str, relation: str, line: int, - confidence: str = "EXTRACTED", weight: float = 1.0) -> None: - edges.append({ + confidence: str = "EXTRACTED", weight: float = 1.0, + context: str | None = None) -> None: + edge = { "source": src, "target": tgt, "relation": relation, @@ -2041,7 +2106,10 @@ def extract_julia(path: Path) -> dict: "source_file": str_path, "source_location": f"L{line}", "weight": weight, - }) + } + if context: + edge["context"] = context + edges.append(edge) file_nid = _make_id(str(path)) add_node(file_nid, path.name, 1) @@ -2068,14 +2136,14 @@ def extract_julia(path: Path) -> dict: callee_name = _read_text(callee, source) target_nid = _make_id(stem, callee_name) add_edge(func_nid, target_nid, "calls", body_node.start_point[0] + 1, - confidence="EXTRACTED") + confidence="EXTRACTED", context="call") # Method call: obj.method(...) elif callee.type == "field_expression" and len(callee.children) >= 3: method_node = callee.children[-1] method_name = _read_text(method_node, source) target_nid = _make_id(stem, method_name) add_edge(func_nid, target_nid, "calls", body_node.start_point[0] + 1, - confidence="EXTRACTED") + confidence="EXTRACTED", context="call") for child in body_node.children: walk_calls(child, func_nid) @@ -2176,14 +2244,14 @@ def extract_julia(path: Path) -> dict: mod_name = _read_text(child, source) imp_nid = _make_id(mod_name) add_node(imp_nid, mod_name, line) - add_edge(scope_nid, imp_nid, "imports", line) + add_edge(scope_nid, imp_nid, "imports", line, context="import") elif child.type == "selected_import": identifiers = [c for c in child.children if c.type == "identifier"] if identifiers: pkg_name = _read_text(identifiers[0], source) pkg_nid = _make_id(pkg_name) add_node(pkg_nid, pkg_name, line) - add_edge(scope_nid, pkg_nid, "imports", line) + add_edge(scope_nid, pkg_nid, "imports", line, context="import") return for child in node.children: @@ -2248,8 +2316,9 @@ def extract_go(path: Path) -> dict: }) def add_edge(src: str, tgt: str, relation: str, line: int, - confidence: str = "EXTRACTED", weight: float = 1.0) -> None: - edges.append({ + confidence: str = "EXTRACTED", weight: float = 1.0, + context: str | None = None) -> None: + edge = { "source": src, "target": tgt, "relation": relation, @@ -2257,7 +2326,10 @@ def extract_go(path: Path) -> dict: "source_file": str_path, "source_location": f"L{line}", "weight": weight, - }) + } + if context: + edge["context"] = context + edges.append(edge) file_nid = _make_id(str(path)) add_node(file_nid, path.name, 1) @@ -2331,7 +2403,7 @@ def extract_go(path: Path) -> dict: # Prefix with go_pkg_ so stdlib names (e.g. "context") # don't collide with local files of the same basename. tgt_nid = _make_id("go", "pkg", raw) - add_edge(file_nid, tgt_nid, "imports_from", spec.start_point[0] + 1) + add_edge(file_nid, tgt_nid, "imports_from", spec.start_point[0] + 1, context="import") # Track local name (alias or last path segment) alias = spec.child_by_field_name("name") local_name = _read_text(alias, source) if alias else raw.split("/")[-1] @@ -2342,7 +2414,7 @@ def extract_go(path: Path) -> dict: if path_node: raw = _read_text(path_node, source).strip('"') tgt_nid = _make_id("go", "pkg", raw) - add_edge(file_nid, tgt_nid, "imports_from", child.start_point[0] + 1) + add_edge(file_nid, tgt_nid, "imports_from", child.start_point[0] + 1, context="import") alias = child.child_by_field_name("name") local_name = _read_text(alias, source) if alias else raw.split("/")[-1] if local_name and local_name != "_" and local_name != ".": @@ -2393,6 +2465,7 @@ def extract_go(path: Path) -> dict: "source": caller_nid, "target": tgt_nid, "relation": "calls", + "context": "call", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{line}", @@ -2460,8 +2533,9 @@ def extract_rust(path: Path) -> dict: }) def add_edge(src: str, tgt: str, relation: str, line: int, - confidence: str = "EXTRACTED", weight: float = 1.0) -> None: - edges.append({ + confidence: str = "EXTRACTED", weight: float = 1.0, + context: str | None = None) -> None: + edge = { "source": src, "target": tgt, "relation": relation, @@ -2469,7 +2543,10 @@ def extract_rust(path: Path) -> dict: "source_file": str_path, "source_location": f"L{line}", "weight": weight, - }) + } + if context: + edge["context"] = context + edges.append(edge) file_nid = _make_id(str(path)) add_node(file_nid, path.name, 1) @@ -2526,7 +2603,7 @@ def extract_rust(path: Path) -> dict: module_name = clean.split("::")[-1].strip() if module_name: tgt_nid = _make_id(module_name) - add_edge(file_nid, tgt_nid, "imports_from", node.start_point[0] + 1) + add_edge(file_nid, tgt_nid, "imports_from", node.start_point[0] + 1, context="import") return for child in node.children: @@ -2573,6 +2650,7 @@ def extract_rust(path: Path) -> dict: "source": caller_nid, "target": tgt_nid, "relation": "calls", + "context": "call", "confidence": "EXTRACTED", "source_file": str_path, "source_location": f"L{line}", @@ -2635,10 +2713,14 @@ def extract_zig(path: Path) -> dict: "source_file": str_path, "source_location": f"L{line}"}) def add_edge(src: str, tgt: str, relation: str, line: int, - confidence: str = "EXTRACTED", weight: float = 1.0) -> None: - edges.append({"source": src, "target": tgt, "relation": relation, - "confidence": confidence, "source_file": str_path, - "source_location": f"L{line}", "weight": weight}) + confidence: str = "EXTRACTED", weight: float = 1.0, + context: str | None = None) -> None: + edge = {"source": src, "target": tgt, "relation": relation, + "confidence": confidence, "source_file": str_path, + "source_location": f"L{line}", "weight": weight} + if context: + edge["context"] = context + edges.append(edge) file_nid = _make_id(str(path)) add_node(file_nid, path.name, 1) @@ -2801,10 +2883,14 @@ def extract_powershell(path: Path) -> dict: "source_file": str_path, "source_location": f"L{line}"}) def add_edge(src: str, tgt: str, relation: str, line: int, - confidence: str = "EXTRACTED", weight: float = 1.0) -> None: - edges.append({"source": src, "target": tgt, "relation": relation, - "confidence": confidence, "source_file": str_path, - "source_location": f"L{line}", "weight": weight}) + confidence: str = "EXTRACTED", weight: float = 1.0, + context: str | None = None) -> None: + edge = {"source": src, "target": tgt, "relation": relation, + "confidence": confidence, "source_file": str_path, + "source_location": f"L{line}", "weight": weight} + if context: + edge["context"] = context + edges.append(edge) file_nid = _make_id(str(path)) add_node(file_nid, path.name, 1) @@ -3196,10 +3282,14 @@ def extract_objc(path: Path) -> dict: "source_file": str_path, "source_location": f"L{line}"}) def add_edge(src: str, tgt: str, relation: str, line: int, - confidence: str = "EXTRACTED", weight: float = 1.0) -> None: - edges.append({"source": src, "target": tgt, "relation": relation, - "confidence": confidence, "source_file": str_path, - "source_location": f"L{line}", "weight": weight}) + confidence: str = "EXTRACTED", weight: float = 1.0, + context: str | None = None) -> None: + edge = {"source": src, "target": tgt, "relation": relation, + "confidence": confidence, "source_file": str_path, + "source_location": f"L{line}", "weight": weight} + if context: + edge["context"] = context + edges.append(edge) file_nid = _make_id(str(path)) add_node(file_nid, path.name, 1) @@ -3223,7 +3313,7 @@ def extract_objc(path: Path) -> dict: module = raw.split("/")[-1].replace(".h", "") if module: tgt_nid = _make_id(module) - add_edge(file_nid, tgt_nid, "imports", line) + add_edge(file_nid, tgt_nid, "imports", line, context="import") elif child.type == "string_literal": # recurse into string_literal to find string_content for sub in child.children: @@ -3232,7 +3322,7 @@ def extract_objc(path: Path) -> dict: module = raw.split("/")[-1].replace(".h", "") if module: tgt_nid = _make_id(module) - add_edge(file_nid, tgt_nid, "imports", line) + add_edge(file_nid, tgt_nid, "imports", line, context="import") return if t == "class_interface": @@ -3263,7 +3353,7 @@ def extract_objc(path: Path) -> dict: for s in sub.children: if s.type == "type_identifier": proto_nid = _make_id(_read(s)) - add_edge(cls_nid, proto_nid, "imports", line) + add_edge(cls_nid, proto_nid, "imports", line, context="import") elif child.type == "method_declaration": walk(child, cls_nid) return @@ -3355,7 +3445,7 @@ def extract_objc(path: Path) -> dict: if pair not in seen_calls and caller_nid != candidate: seen_calls.add(pair) add_edge(caller_nid, candidate, "calls", body_node.start_point[0] + 1, - confidence="EXTRACTED", weight=1.0) + confidence="EXTRACTED", weight=1.0, context="call") for child in n.children: walk_calls(child) walk_calls(body_node) @@ -3394,10 +3484,14 @@ def extract_elixir(path: Path) -> dict: "source_file": str_path, "source_location": f"L{line}"}) def add_edge(src: str, tgt: str, relation: str, line: int, - confidence: str = "EXTRACTED", weight: float = 1.0) -> None: - edges.append({"source": src, "target": tgt, "relation": relation, - "confidence": confidence, "source_file": str_path, - "source_location": f"L{line}", "weight": weight}) + confidence: str = "EXTRACTED", weight: float = 1.0, + context: str | None = None) -> None: + edge = {"source": src, "target": tgt, "relation": relation, + "confidence": confidence, "source_file": str_path, + "source_location": f"L{line}", "weight": weight} + if context: + edge["context"] = context + edges.append(edge) file_nid = _make_id(str(path)) add_node(file_nid, path.name, 1) @@ -3476,7 +3570,7 @@ def extract_elixir(path: Path) -> dict: module_name = _get_alias_text(arguments_node) if module_name: tgt_nid = _make_id(module_name) - add_edge(file_nid, tgt_nid, "imports", line) + add_edge(file_nid, tgt_nid, "imports", line, context="import") return for child in node.children: @@ -3531,7 +3625,8 @@ def extract_elixir(path: Path) -> dict: if pair not in seen_call_pairs: seen_call_pairs.add(pair) add_edge(caller_nid, tgt_nid, "calls", - node.start_point[0] + 1, confidence="EXTRACTED", weight=1.0) + node.start_point[0] + 1, confidence="EXTRACTED", weight=1.0, + context="call") else: raw_calls.append({ "caller_nid": caller_nid, @@ -3763,6 +3858,7 @@ def extract(paths: list[Path], cache_root: Path | None = None) -> dict: "source": caller, "target": tgt, "relation": "calls", + "context": "call", "confidence": "INFERRED", "confidence_score": 0.8, "source_file": rc.get("source_file", ""), diff --git a/graphify/serve.py b/graphify/serve.py index dd82bcb6..708f3f3c 100644 --- a/graphify/serve.py +++ b/graphify/serve.py @@ -63,6 +63,68 @@ def _score_nodes(G: nx.Graph, terms: list[str]) -> list[tuple[float, str]]: return sorted(scored, reverse=True) +_CONTEXT_HINTS: tuple[tuple[str, tuple[str, ...]], ...] = ( + ("call", ("call", "calls", "called", "invoke", "invokes", "invoked")), + ("import", ("import", "imports", "imported", "module", "modules")), + ("field", ("field", "fields", "member", "members", "property", "properties")), + ("parameter_type", ("parameter", "parameters", "param", "params", "argument", "arguments")), + ("return_type", ("return", "returns", "returned")), + ("generic_arg", ("generic", "generics", "template", "templates")), +) + + +def _normalize_context_filters(filters: list[str] | None) -> list[str]: + if not filters: + return [] + normalized: list[str] = [] + seen: set[str] = set() + for value in filters: + key = _strip_diacritics(str(value)).strip().lower() + if key and key not in seen: + seen.add(key) + normalized.append(key) + return normalized + + +def _infer_context_filters(question: str) -> list[str]: + lowered = { + _strip_diacritics(token).lower() + for token in question.replace("?", " ").replace(",", " ").split() + } + inferred: list[str] = [] + for context, hints in _CONTEXT_HINTS: + if any(hint in lowered for hint in hints): + inferred.append(context) + return inferred + + +def _resolve_context_filters(question: str, explicit_filters: list[str] | None = None) -> tuple[list[str], str | None]: + normalized = _normalize_context_filters(explicit_filters) + if normalized: + return normalized, "explicit" + inferred = _infer_context_filters(question) + if inferred: + return inferred, "heuristic" + return [], None + + +def _filter_graph_by_context(G: nx.Graph, context_filters: list[str] | None) -> nx.Graph: + filters = set(_normalize_context_filters(context_filters)) + if not filters: + return G + H = G.__class__() + H.add_nodes_from(G.nodes(data=True)) + if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)): + for u, v, key, data in G.edges(keys=True, data=True): + if data.get("context") in filters: + H.add_edge(u, v, key=key, **data) + else: + for u, v, data in G.edges(data=True): + if data.get("context") in filters: + H.add_edge(u, v, **data) + return H + + def _bfs(G: nx.Graph, start_nodes: list[str], depth: int) -> tuple[set[str], list[tuple]]: visited: set[str] = set(start_nodes) frontier = set(start_nodes) @@ -114,7 +176,13 @@ def _subgraph_to_text(G: nx.Graph, nodes: set[str], edges: list[tuple], token_bu if u in nodes and v in nodes: raw = G[u][v] d = next(iter(raw.values()), {}) if isinstance(G, (nx.MultiGraph, nx.MultiDiGraph)) else raw - line = f"EDGE {sanitize_label(G.nodes[u].get('label', u))} --{d.get('relation', '')} [{d.get('confidence', '')}]--> {sanitize_label(G.nodes[v].get('label', v))}" + context = d.get("context") + context_suffix = f" context={context}" if context else "" + line = ( + f"EDGE {sanitize_label(G.nodes[u].get('label', u))} " + f"--{d.get('relation', '')} [{d.get('confidence', '')}{context_suffix}]--> " + f"{sanitize_label(G.nodes[v].get('label', v))}" + ) lines.append(line) output = "\n".join(lines) if len(output) > char_budget: @@ -122,6 +190,34 @@ def _subgraph_to_text(G: nx.Graph, nodes: set[str], edges: list[tuple], token_bu return output +def _query_graph_text( + G: nx.Graph, + question: str, + *, + mode: str = "bfs", + depth: int = 3, + token_budget: int = 2000, + context_filters: list[str] | None = None, +) -> str: + terms = [t.lower() for t in question.split() if len(t) > 2] + scored = _score_nodes(G, terms) + start_nodes = [nid for _, nid in scored[:3]] + if not start_nodes: + return "No matching nodes found." + resolved_filters, filter_source = _resolve_context_filters(question, context_filters) + traversal_graph = _filter_graph_by_context(G, resolved_filters) + nodes, edges = _dfs(traversal_graph, start_nodes, depth) if mode == "dfs" else _bfs(traversal_graph, start_nodes, depth) + header_parts = [ + f"Traversal: {mode.upper()} depth={depth}", + f"Start: {[G.nodes[n].get('label', n) for n in start_nodes]}", + ] + if resolved_filters: + header_parts.append(f"Context: {', '.join(resolved_filters)} ({filter_source})") + header_parts.append(f"{len(nodes)} nodes found") + header = " | ".join(header_parts) + "\n\n" + return header + _subgraph_to_text(traversal_graph, nodes, edges, token_budget) + + def _find_node(G: nx.Graph, label: str) -> list[str]: """Return node IDs whose label or ID matches the search term (diacritic-insensitive).""" term = _strip_diacritics(label).lower() @@ -188,6 +284,11 @@ def serve(graph_path: str = "graphify-out/graph.json") -> None: "description": "bfs=broad context, dfs=trace a specific path"}, "depth": {"type": "integer", "default": 3, "description": "Traversal depth (1-6)"}, "token_budget": {"type": "integer", "default": 2000, "description": "Max output tokens"}, + "context_filter": { + "type": "array", + "items": {"type": "string"}, + "description": "Optional explicit edge-context filter, e.g. ['call', 'field']", + }, }, "required": ["question"], }, @@ -252,14 +353,15 @@ def serve(graph_path: str = "graphify-out/graph.json") -> None: mode = arguments.get("mode", "bfs") depth = min(int(arguments.get("depth", 3)), 6) budget = int(arguments.get("token_budget", 2000)) - terms = [t.lower() for t in question.split() if len(t) > 2] - scored = _score_nodes(G, terms) - start_nodes = [nid for _, nid in scored[:3]] - if not start_nodes: - return "No matching nodes found." - nodes, edges = _dfs(G, start_nodes, depth) if mode == "dfs" else _bfs(G, start_nodes, depth) - header = f"Traversal: {mode.upper()} depth={depth} | Start: {[G.nodes[n].get('label', n) for n in start_nodes]} | {len(nodes)} nodes found\n\n" - return header + _subgraph_to_text(G, nodes, edges, budget, seeds=start_nodes) + context_filter = arguments.get("context_filter") + return _query_graph_text( + G, + question, + mode=mode, + depth=depth, + token_budget=budget, + context_filters=context_filter, + ) def _tool_get_node(arguments: dict) -> str: label = arguments["label"].lower() diff --git a/tests/test_extract.py b/tests/test_extract.py index 3a26693d..06d371bc 100644 --- a/tests/test_extract.py +++ b/tests/test_extract.py @@ -124,6 +124,13 @@ def test_calls_edges_are_extracted(): assert edge["weight"] == 1.0 +def test_python_call_edges_have_call_context(): + result = extract_python(FIXTURES / "sample_calls.py") + call_edges = [e for e in result["edges"] if e["relation"] == "calls"] + assert call_edges + assert all(e.get("context") == "call" for e in call_edges) + + def test_calls_no_self_loops(): result = extract_python(FIXTURES / "sample_calls.py") for edge in result["edges"]: diff --git a/tests/test_languages.py b/tests/test_languages.py index 62ef1143..c4541de8 100644 --- a/tests/test_languages.py +++ b/tests/test_languages.py @@ -25,6 +25,22 @@ def _calls(r): } +def _references(r): + node_by_id = {n["id"]: n["label"] for n in r["nodes"]} + return [ + ( + node_by_id.get(e["source"], e["source"]), + node_by_id.get(e["target"], e["target"]), + e, + ) + for e in r["edges"] if e["relation"] == "references" + ] + + +def _edges_with_relation(r, *relations): + return [e for e in r["edges"] if e["relation"] in relations] + + # ── Java ────────────────────────────────────────────────────────────────────── def test_java_no_error(): @@ -49,6 +65,13 @@ def test_java_finds_imports(): r = extract_java(FIXTURES / "sample.java") assert "imports" in _relations(r) + +def test_java_import_edges_have_import_context(): + r = extract_java(FIXTURES / "sample.java") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + def test_java_no_dangling_edges(): r = extract_java(FIXTURES / "sample.java") node_ids = {n["id"] for n in r["nodes"]} @@ -83,6 +106,20 @@ def test_c_calls_are_extracted(): assert e["confidence"] == "EXTRACTED" +def test_c_import_edges_have_import_context(): + r = extract_c(FIXTURES / "sample.c") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + + +def test_c_call_edges_have_call_context(): + r = extract_c(FIXTURES / "sample.c") + call_edges = _edges_with_relation(r, "calls") + assert call_edges + assert all(e.get("context") == "call" for e in call_edges) + + # ── C++ ─────────────────────────────────────────────────────────────────────── def test_cpp_no_error(): @@ -104,6 +141,13 @@ def test_cpp_finds_includes(): assert "imports" in _relations(r) +def test_cpp_import_edges_have_import_context(): + r = extract_cpp(FIXTURES / "sample.cpp") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + + # ── Ruby ───────────────────────────────────────────────────────────────────── def test_ruby_no_error(): @@ -164,6 +208,33 @@ def test_csharp_inherits_iprocessor(): assert found, "DataProcessor should have inherits edge to IProcessor" +def test_csharp_field_type_references_have_field_context(): + r = extract_csharp(FIXTURES / "sample.cs") + refs = _references(r) + assert any( + "DataProcessor" in src and "HttpClient" in tgt and edge.get("context") == "field" + for src, tgt, edge in refs + ), "DataProcessor field declarations should reference HttpClient with field context" + + +def test_csharp_call_edges_have_call_context(): + r = extract_csharp(FIXTURES / "sample.cs") + node_by_id = {n["id"]: n["label"] for n in r["nodes"]} + assert any( + "Process" in node_by_id.get(e["source"], "") + and "Validate" in node_by_id.get(e["target"], "") + and e.get("context") == "call" + for e in r["edges"] if e["relation"] == "calls" + ), "C# call edges should retain call context" + + +def test_csharp_import_edges_have_import_context(): + r = extract_csharp(FIXTURES / "sample.cs") + import_edges = [e for e in r["edges"] if e["relation"] == "imports"] + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + + # ── Kotlin ─────────────────────────────────────────────────────────────────── def test_kotlin_no_error(): @@ -222,6 +293,20 @@ def test_scala_finds_methods(): assert any("post" in l for l in labels) +def test_scala_import_edges_have_import_context(): + r = extract_scala(FIXTURES / "sample.scala") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + + +def test_scala_call_edges_have_call_context(): + r = extract_scala(FIXTURES / "sample.scala") + call_edges = _edges_with_relation(r, "calls") + assert call_edges + assert all(e.get("context") == "call" for e in call_edges) + + # ── PHP ─────────────────────────────────────────────────────────────────────── def test_php_no_error(): @@ -246,6 +331,20 @@ def test_php_finds_imports(): r = extract_php(FIXTURES / "sample.php") assert "imports" in _relations(r) + +def test_php_import_edges_have_import_context(): + r = extract_php(FIXTURES / "sample.php") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + + +def test_php_call_edges_have_call_context(): + r = extract_php(FIXTURES / "sample.php") + call_edges = _edges_with_relation(r, "calls") + assert call_edges + assert all(e.get("context") == "call" for e in call_edges) + def test_php_finds_static_property_access(): r = extract_php(FIXTURES / "sample_php_static_prop.php") assert "uses_static_prop" in _relations(r) @@ -331,6 +430,13 @@ def test_swift_finds_imports(): r = extract_swift(FIXTURES / "sample.swift") assert "imports" in _relations(r) + +def test_swift_import_edges_have_import_context(): + r = extract_swift(FIXTURES / "sample.swift") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + def test_swift_no_dangling_edges(): r = extract_swift(FIXTURES / "sample.swift") node_ids = {n["id"] for n in r["nodes"]} @@ -489,6 +595,13 @@ def test_vbnet_no_dangling_edges(): for e in r["edges"]: assert e["source"] in node_ids +def test_swift_call_edges_have_call_context(): + r = extract_swift(FIXTURES / "sample.swift") + call_edges = _edges_with_relation(r, "calls") + assert call_edges + assert all(e.get("context") == "call" for e in call_edges) + + # ── Elixir ──────────────────────────────────────────────────────────────────── from graphify.extract import extract_elixir @@ -511,12 +624,26 @@ def test_elixir_finds_imports(): import_edges = [e for e in r["edges"] if e["relation"] == "imports"] assert len(import_edges) >= 2 + +def test_elixir_import_edges_have_import_context(): + r = extract_elixir(FIXTURES / "sample.ex") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + def test_elixir_finds_calls(): r = extract_elixir(FIXTURES / "sample.ex") calls = {(e["source"], e["target"]) for e in r["edges"] if e["relation"] == "calls"} labels = {n["id"]: n["label"] for n in r["nodes"]} assert any("create" in labels.get(src, "") and "validate" in labels.get(tgt, "") for src, tgt in calls) + +def test_elixir_call_edges_have_call_context(): + r = extract_elixir(FIXTURES / "sample.ex") + call_edges = _edges_with_relation(r, "calls") + assert call_edges + assert all(e.get("context") == "call" for e in call_edges) + def test_elixir_method_edges(): r = extract_elixir(FIXTURES / "sample.ex") methods = [e for e in r["edges"] if e["relation"] == "method"] @@ -551,6 +678,13 @@ def test_objc_finds_imports(): assert len(import_edges) >= 1 +def test_objc_import_edges_have_import_context(): + r = extract_objc(FIXTURES / "sample.m") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + + def test_objc_inherits_edge(): r = extract_objc(FIXTURES / "sample.m") inherits = [e for e in r["edges"] if e["relation"] == "inherits"] @@ -626,6 +760,13 @@ def test_julia_finds_imports(): assert len(import_edges) >= 1 +def test_julia_import_edges_have_import_context(): + r = extract_julia(FIXTURES / "sample.jl") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + + def test_julia_finds_inherits(): r = extract_julia(FIXTURES / "sample.jl") inherits = [e for e in r["edges"] if e["relation"] == "inherits"] @@ -638,6 +779,13 @@ def test_julia_finds_calls(): assert len(call_edges) >= 1 +def test_julia_call_edges_have_call_context(): + r = extract_julia(FIXTURES / "sample.jl") + call_edges = _edges_with_relation(r, "calls") + assert call_edges + assert all(e.get("context") == "call" for e in call_edges) + + def test_julia_no_dangling_edges(): r = extract_julia(FIXTURES / "sample.jl") node_ids = {n["id"] for n in r["nodes"]} diff --git a/tests/test_multilang.py b/tests/test_multilang.py index 3264122c..e4620034 100644 --- a/tests/test_multilang.py +++ b/tests/test_multilang.py @@ -24,6 +24,10 @@ def _confidences(result): return {e["confidence"] for e in result["edges"]} +def _edges_with_relation(result, *relations): + return [e for e in result["edges"] if e["relation"] in relations] + + # ── TypeScript ──────────────────────────────────────────────────────────────── def test_ts_finds_class(): @@ -53,6 +57,20 @@ def test_ts_calls_are_extracted(): if e["relation"] == "calls": assert e["confidence"] == "EXTRACTED" + +def test_ts_import_edges_have_import_context(): + r = extract_js(FIXTURES / "sample.ts") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + + +def test_ts_call_edges_have_call_context(): + r = extract_js(FIXTURES / "sample.ts") + call_edges = _edges_with_relation(r, "calls") + assert call_edges + assert all(e.get("context") == "call" for e in call_edges) + def test_ts_no_dangling_edges(): r = extract_js(FIXTURES / "sample.ts") node_ids = {n["id"] for n in r["nodes"]} @@ -87,6 +105,20 @@ def test_go_has_extracted_calls(): r = extract_go(FIXTURES / "sample.go") assert "EXTRACTED" in _confidences(r) + +def test_go_import_edges_have_import_context(): + r = extract_go(FIXTURES / "sample.go") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + + +def test_go_call_edges_have_call_context(): + r = extract_go(FIXTURES / "sample.go") + call_edges = _edges_with_relation(r, "calls") + assert call_edges + assert all(e.get("context") == "call" for e in call_edges) + def test_go_no_dangling_edges(): r = extract_go(FIXTURES / "sample.go") node_ids = {n["id"] for n in r["nodes"]} @@ -123,6 +155,20 @@ def test_rust_calls_are_extracted(): if e["relation"] == "calls": assert e["confidence"] == "EXTRACTED" + +def test_rust_import_edges_have_import_context(): + r = extract_rust(FIXTURES / "sample.rs") + import_edges = _edges_with_relation(r, "imports", "imports_from") + assert import_edges + assert all(e.get("context") == "import" for e in import_edges) + + +def test_rust_call_edges_have_call_context(): + r = extract_rust(FIXTURES / "sample.rs") + call_edges = _edges_with_relation(r, "calls") + assert call_edges + assert all(e.get("context") == "call" for e in call_edges) + def test_rust_no_dangling_edges(): r = extract_rust(FIXTURES / "sample.rs") node_ids = {n["id"] for n in r["nodes"]} diff --git a/tests/test_query_cli.py b/tests/test_query_cli.py new file mode 100644 index 00000000..39d016f8 --- /dev/null +++ b/tests/test_query_cli.py @@ -0,0 +1,51 @@ +"""Tests for graphify query CLI context filtering.""" +from __future__ import annotations + +import json + +import networkx as nx +from networkx.readwrite import json_graph + +import graphify.__main__ as mainmod + + +def _write_graph(tmp_path): + G = nx.Graph() + G.add_node("n1", label="extract", source_file="extract.py", source_location="L10", community=0) + G.add_node("n2", label="cluster", source_file="cluster.py", source_location="L5", community=0) + G.add_node("n3", label="build", source_file="build.py", source_location="L1", community=1) + G.add_edge("n1", "n2", relation="calls", confidence="EXTRACTED", context="call") + G.add_edge("n2", "n3", relation="imports", confidence="EXTRACTED", context="import") + graph_path = tmp_path / "graph.json" + graph_path.write_text(json.dumps(json_graph.node_link_data(G, edges="links"))) + return graph_path + + +def test_query_cli_explicit_context_filter(monkeypatch, tmp_path, capsys): + graph_path = _write_graph(tmp_path) + monkeypatch.setattr(mainmod, "_check_skill_version", lambda _: None) + monkeypatch.setattr( + mainmod.sys, + "argv", + ["graphify", "query", "extract", "--context", "call", "--graph", str(graph_path)], + ) + mainmod.main() + out = capsys.readouterr().out + assert "Context: call (explicit)" in out + assert "cluster" in out + assert "build" not in out + + +def test_query_cli_heuristic_context_filter(monkeypatch, tmp_path, capsys): + graph_path = _write_graph(tmp_path) + monkeypatch.setattr(mainmod, "_check_skill_version", lambda _: None) + monkeypatch.setattr( + mainmod.sys, + "argv", + ["graphify", "query", "who calls extract", "--graph", str(graph_path)], + ) + mainmod.main() + out = capsys.readouterr().out + assert "Context: call (heuristic)" in out + assert "cluster" in out + assert "build" not in out diff --git a/tests/test_serve.py b/tests/test_serve.py index 6457ac50..49d0a200 100644 --- a/tests/test_serve.py +++ b/tests/test_serve.py @@ -9,6 +9,10 @@ from graphify.serve import ( _score_nodes, _bfs, _dfs, + _filter_graph_by_context, + _infer_context_filters, + _query_graph_text, + _resolve_context_filters, _subgraph_to_text, _load_graph, ) @@ -21,8 +25,8 @@ def _make_graph() -> nx.Graph: G.add_node("n3", label="build", source_file="build.py", source_location="L1", community=1) G.add_node("n4", label="report", source_file="report.py", source_location="L1", community=1) G.add_node("n5", label="isolated", source_file="other.py", source_location="L1", community=2) - G.add_edge("n1", "n2", relation="calls", confidence="INFERRED") - G.add_edge("n2", "n3", relation="imports", confidence="EXTRACTED") + G.add_edge("n1", "n2", relation="calls", confidence="INFERRED", context="call") + G.add_edge("n2", "n3", relation="imports", confidence="EXTRACTED", context="import") G.add_edge("n3", "n4", relation="uses", confidence="EXTRACTED") return G @@ -73,6 +77,16 @@ def test_score_nodes_source_file_partial(): assert "n2" in nids +def test_infer_context_filters_for_calls_question(): + assert _infer_context_filters("who calls extract") == ["call"] + + +def test_resolve_context_filters_explicit_overrides_heuristic(): + filters, source = _resolve_context_filters("who calls extract", ["field"]) + assert filters == ["field"] + assert source == "explicit" + + # --- _bfs --- def test_bfs_depth_1(): @@ -99,6 +113,15 @@ def test_bfs_returns_edges(): assert any(u == "n1" or v == "n1" for u, v in edges) +def test_filter_graph_by_context_limits_traversal(): + G = _make_graph() + filtered = _filter_graph_by_context(G, ["call"]) + visited, edges = _bfs(filtered, ["n1"], depth=2) + assert "n2" in visited + assert "n3" not in visited + assert edges == [("n1", "n2")] + + # --- _dfs --- def test_dfs_depth_1(): @@ -135,6 +158,28 @@ def test_subgraph_to_text_edge_included(): assert "calls" in text +def test_subgraph_to_text_includes_edge_context(): + G = _make_graph() + text = _subgraph_to_text(G, {"n1", "n2"}, [("n1", "n2")]) + assert "context=call" in text + + +def test_query_graph_text_explicit_context_filter_changes_traversal(): + G = _make_graph() + text = _query_graph_text(G, "extract", mode="bfs", depth=2, token_budget=2000, context_filters=["call"]) + assert "Context: call (explicit)" in text + assert "cluster" in text + assert "build" not in text + + +def test_query_graph_text_heuristic_context_filter_changes_traversal(): + G = _make_graph() + text = _query_graph_text(G, "who calls extract", mode="bfs", depth=2, token_budget=2000) + assert "Context: call (heuristic)" in text + assert "cluster" in text + assert "build" not in text + + # --- _load_graph --- def test_load_graph_roundtrip(tmp_path):