diff --git a/mcp/src/graph/neo4j.ts b/mcp/src/graph/neo4j.ts index fd535f2f4..be6c49903 100644 --- a/mcp/src/graph/neo4j.ts +++ b/mcp/src/graph/neo4j.ts @@ -697,11 +697,21 @@ class Db { } } - async get_node_with_related(ref_id: string) { + async get_node_with_related( + ref_id: string, + depth: number = 100, + node_type: NodeType[] = [], + edge_type: EdgeType[] = [], + limit: number = 25, + ) { const session = this.resilientSession(); try { - const result = await session.run(Q.GET_NODE_WITH_RELATED_QUERY, { + const result = await session.run(Q.SUBGRAPH_ALL_QUERY, { ref_id, + depth, + limit, + node_filter: node_type.length > 0 ? node_type.join("|") : null, + edge_filter: edge_type.length > 0 ? edge_type.join("|") : null, }); if (result.records.length === 0) { @@ -709,7 +719,7 @@ class Db { } const record = result.records[0]; - const allNodesArray = record.get("allNodes"); + const allNodesArray = record.get("nodes"); const edgesArray = record.get("edges"); // Process nodes diff --git a/mcp/src/graph/queries.ts b/mcp/src/graph/queries.ts index e84f20f31..e6c470704 100644 --- a/mcp/src/graph/queries.ts +++ b/mcp/src/graph/queries.ts @@ -236,33 +236,29 @@ MATCH (p:Prompt {ref_id: $prompt_ref_id})-[r]->(h:Hint) RETURN h `; -export const GET_NODE_WITH_RELATED_QUERY = ` -// First, collect all related nodes -MATCH (h:${Data_Bank} {ref_id: $ref_id}) -OPTIONAL MATCH (h)-[e]-(m) -WITH h, collect(DISTINCT m) AS relatedNodes - -// Collect all nodes including the main node -WITH h, relatedNodes, [h] + relatedNodes AS allNodes - -// Extract ref_ids for finding interconnections -WITH h, allNodes, [node IN allNodes WHERE node IS NOT NULL | node.ref_id] AS allRefIds - -// Find all edges between any of the collected nodes -OPTIONAL MATCH (a)-[r]->(b) -WHERE a.ref_id IN allRefIds - AND b.ref_id IN allRefIds - AND a.ref_id <> b.ref_id - -// Return nodes and interconnecting edges -RETURN h, allNodes, - collect(DISTINCT { - source: a.ref_id, - target: b.ref_id, +// apoc.path.subgraphAll always includes the start node, expanding outward up to +// $depth hops and capping traversal at $limit paths. $node_filter/$edge_filter are +// null when the caller passes no types, which APOC treats as "no restriction". +export const SUBGRAPH_ALL_QUERY = ` +MATCH (start:${Data_Bank} {ref_id: $ref_id}) +CALL apoc.path.subgraphAll(start, { + maxLevel: $depth, + relationshipFilter: $edge_filter, + labelFilter: $node_filter, + limit: $limit +}) +YIELD nodes, relationships +WITH nodes, relationships +UNWIND CASE WHEN size(relationships) = 0 THEN [null] ELSE relationships END AS r +WITH nodes, collect( + CASE WHEN r IS NULL THEN null ELSE { + source: startNode(r).ref_id, + target: endNode(r).ref_id, edge_type: type(r), - properties: properties(r), - ref_id: r.ref_id - }) AS edges + properties: properties(r) + } END +) AS rawEdges +RETURN nodes, [e IN rawEdges WHERE e IS NOT NULL] AS edges `; export const GET_WORKFLOW_PUBLISHED_VERSION_SUBGRAPH_QUERY = ` diff --git a/mcp/src/graph/routes.ts b/mcp/src/graph/routes.ts index 6f08f6061..663def1ea 100644 --- a/mcp/src/graph/routes.ts +++ b/mcp/src/graph/routes.ts @@ -377,7 +377,7 @@ export async function get_learnings_og(req: Request, res: Response) { } export async function fetch_node_with_related(req: Request, res: Response) { - // curl "http://localhost:3355/node?ref_id=bcc79e17-fae9-41d6-8932-40ea60e34b54" + // curl "http://localhost:3355/subgraph?ref_id=bcc79e17-fae9-41d6-8932-40ea60e34b54&depth=3&node_type=Function,Endpoint&edge_type=CALLS&limit=50" try { const ref_id = req.query.ref_id as string; if (!ref_id) { @@ -385,7 +385,20 @@ export async function fetch_node_with_related(req: Request, res: Response) { return; } - const result = await db.get_node_with_related(ref_id); + const depth = parseInt(req.query.depth as string) || 100; + const node_type = parseNodeTypes(req.query); + const edge_type: EdgeType[] = req.query.edge_type + ? ((req.query.edge_type as string).split(",").map((s) => s.trim()) as EdgeType[]) + : []; + const limit = parseLimit(req.query) ?? 25; + + const result = await db.get_node_with_related( + ref_id, + depth, + node_type, + edge_type, + limit, + ); res.json(result); } catch (error) {