Skip to content

Commit 8e69a7d

Browse files
m1rl0km1rl0k
authored
Merge pull request #189 from m1rl0k/multi-granular
Improve symbol matching in graph queries and enforce JSON output
2 parents aa5db61 + b92daa1 commit 8e69a7d

2 files changed

Lines changed: 56 additions & 32 deletions

File tree

‎scripts/mcp_impl/context_search.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -594,6 +594,8 @@ def _maybe_dict(val: Any) -> Dict[str, Any]:
594594
pass
595595

596596
# First: run code search via internal repo_search for consistent behavior
597+
# Force output_format="json" to ensure we get raw results for internal processing
598+
# (TOON format returns results as a string which breaks result parsing)
597599
code_res = await repo_search_fn(
598600
query=queries if len(queries) > 1 else (queries[0] if queries else ""),
599601
limit=code_limit,
@@ -619,6 +621,7 @@ def _maybe_dict(val: Any) -> Dict[str, Any]:
619621
compact=False,
620622
repo=repo, # Cross-codebase isolation
621623
session=session,
624+
output_format="json", # Always use JSON for internal processing
622625
)
623626

624627
# Optional debug

‎scripts/mcp_impl/neo4j_graph.py‎

Lines changed: 53 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -183,30 +183,37 @@ async def _query_callers_async(
183183
"""Async simple caller query (depth 1).
184184
185185
Supports class-level queries: for "MyClass", also matches callers of "MyClass.method".
186+
Also supports suffix matching for symbols stored with full module paths
187+
(e.g., "RecursiveReranker" matches "scripts.rerank_recursive.RecursiveReranker").
186188
"""
187189
symbol_prefix = f"{symbol}."
190+
symbol_suffix = f".{symbol}" # For matching full module paths like "module.ClassName"
188191
if repo and repo != "*":
189192
query = """
190193
MATCH (caller:Symbol {collection: $collection})-[r:CALLS]->(callee:Symbol {collection: $collection})
191-
WHERE (callee.name = $symbol OR callee.name STARTS WITH $symbol_prefix)
194+
WHERE (callee.name = $symbol
195+
OR callee.name STARTS WITH $symbol_prefix
196+
OR callee.name ENDS WITH $symbol_suffix)
192197
AND r.collection = $collection AND (r.repo = $repo OR callee.repo = $repo)
193198
RETURN caller.name as symbol, r.caller_path as path,
194199
r.start_line as start_line, r.end_line as end_line,
195200
r.language as language, callee.repo as repo
196201
LIMIT $limit
197202
"""
198-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "repo": repo, "collection": collection, "limit": limit}
203+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "repo": repo, "collection": collection, "limit": limit}
199204
else:
200205
query = """
201206
MATCH (caller:Symbol {collection: $collection})-[r:CALLS]->(callee:Symbol {collection: $collection})
202-
WHERE (callee.name = $symbol OR callee.name STARTS WITH $symbol_prefix)
207+
WHERE (callee.name = $symbol
208+
OR callee.name STARTS WITH $symbol_prefix
209+
OR callee.name ENDS WITH $symbol_suffix)
203210
AND r.collection = $collection
204211
RETURN caller.name as symbol, r.caller_path as path,
205212
r.start_line as start_line, r.end_line as end_line,
206213
r.language as language, callee.repo as repo
207214
LIMIT $limit
208215
"""
209-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "collection": collection, "limit": limit}
216+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "collection": collection, "limit": limit}
210217

211218
return await backend.run_query_async(query, params)
212219

@@ -221,30 +228,36 @@ async def _query_callees_async(
221228
"""Async simple callee query (depth 1).
222229
223230
Supports class-level queries: for "MyClass", also matches callees of "MyClass.method".
231+
Also supports suffix matching for symbols stored with full module paths.
224232
"""
225233
symbol_prefix = f"{symbol}."
234+
symbol_suffix = f".{symbol}" # For matching full module paths
226235
if repo and repo != "*":
227236
query = """
228237
MATCH (caller:Symbol {collection: $collection})-[r:CALLS]->(callee:Symbol {collection: $collection})
229-
WHERE (caller.name = $symbol OR caller.name STARTS WITH $symbol_prefix)
238+
WHERE (caller.name = $symbol
239+
OR caller.name STARTS WITH $symbol_prefix
240+
OR caller.name ENDS WITH $symbol_suffix)
230241
AND r.collection = $collection AND (r.repo = $repo OR caller.repo = $repo)
231242
RETURN callee.name as symbol, r.caller_path as path,
232243
r.start_line as start_line, r.end_line as end_line,
233244
r.language as language, caller.repo as repo
234245
LIMIT $limit
235246
"""
236-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "repo": repo, "collection": collection, "limit": limit}
247+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "repo": repo, "collection": collection, "limit": limit}
237248
else:
238249
query = """
239250
MATCH (caller:Symbol {collection: $collection})-[r:CALLS]->(callee:Symbol {collection: $collection})
240-
WHERE (caller.name = $symbol OR caller.name STARTS WITH $symbol_prefix)
251+
WHERE (caller.name = $symbol
252+
OR caller.name STARTS WITH $symbol_prefix
253+
OR caller.name ENDS WITH $symbol_suffix)
241254
AND r.collection = $collection
242255
RETURN callee.name as symbol, r.caller_path as path,
243256
r.start_line as start_line, r.end_line as end_line,
244257
r.language as language, caller.repo as repo
245258
LIMIT $limit
246259
"""
247-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "collection": collection, "limit": limit}
260+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "collection": collection, "limit": limit}
248261

249262
return await backend.run_query_async(query, params)
250263

@@ -261,15 +274,17 @@ async def _query_transitive_callers_async(
261274
"""Async multi-hop caller traversal.
262275
263276
Supports class-level queries: for "MyClass", also matches callers of "MyClass.method".
277+
Also supports suffix matching for symbols stored with full module paths.
264278
"""
265279
safe_depth = max(1, min(10, int(depth)))
266280
symbol_prefix = f"{symbol}."
281+
symbol_suffix = f".{symbol}" # For matching full module paths
267282

268283
if include_paths:
269284
if repo and repo != "*":
270285
query = f"""
271286
MATCH path = (caller:Symbol {{collection: $collection}})-[:CALLS*1..{safe_depth}]->(target:Symbol {{collection: $collection}})
272-
WHERE (target.name = $symbol OR target.name STARTS WITH $symbol_prefix)
287+
WHERE (target.name = $symbol OR target.name STARTS WITH $symbol_prefix OR target.name ENDS WITH $symbol_suffix)
273288
AND all(r IN relationships(path) WHERE r.collection = $collection AND r.repo = $repo)
274289
WITH caller, path, length(path) as hop
275290
RETURN caller.name as symbol, hop,
@@ -278,11 +293,11 @@ async def _query_transitive_callers_async(
278293
ORDER BY hop
279294
LIMIT $limit
280295
"""
281-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "repo": repo, "collection": collection, "limit": limit}
296+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "repo": repo, "collection": collection, "limit": limit}
282297
else:
283298
query = f"""
284299
MATCH path = (caller:Symbol {{collection: $collection}})-[:CALLS*1..{safe_depth}]->(target:Symbol {{collection: $collection}})
285-
WHERE (target.name = $symbol OR target.name STARTS WITH $symbol_prefix)
300+
WHERE (target.name = $symbol OR target.name STARTS WITH $symbol_prefix OR target.name ENDS WITH $symbol_suffix)
286301
AND all(r IN relationships(path) WHERE r.collection = $collection)
287302
WITH caller, path, length(path) as hop
288303
RETURN caller.name as symbol, hop,
@@ -291,28 +306,28 @@ async def _query_transitive_callers_async(
291306
ORDER BY hop
292307
LIMIT $limit
293308
"""
294-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "collection": collection, "limit": limit}
309+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "collection": collection, "limit": limit}
295310
else:
296311
if repo and repo != "*":
297312
query = f"""
298313
MATCH path = (caller:Symbol {{collection: $collection}})-[:CALLS*1..{safe_depth}]->(target:Symbol {{collection: $collection}})
299-
WHERE (target.name = $symbol OR target.name STARTS WITH $symbol_prefix)
314+
WHERE (target.name = $symbol OR target.name STARTS WITH $symbol_prefix OR target.name ENDS WITH $symbol_suffix)
300315
AND all(r IN relationships(path) WHERE r.collection = $collection AND r.repo = $repo)
301316
WITH DISTINCT caller
302317
RETURN caller.name as symbol, caller.repo as repo
303318
LIMIT $limit
304319
"""
305-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "repo": repo, "collection": collection, "limit": limit}
320+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "repo": repo, "collection": collection, "limit": limit}
306321
else:
307322
query = f"""
308323
MATCH path = (caller:Symbol {{collection: $collection}})-[:CALLS*1..{safe_depth}]->(target:Symbol {{collection: $collection}})
309-
WHERE (target.name = $symbol OR target.name STARTS WITH $symbol_prefix)
324+
WHERE (target.name = $symbol OR target.name STARTS WITH $symbol_prefix OR target.name ENDS WITH $symbol_suffix)
310325
AND all(r IN relationships(path) WHERE r.collection = $collection)
311326
WITH DISTINCT caller
312327
RETURN caller.name as symbol, caller.repo as repo
313328
LIMIT $limit
314329
"""
315-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "collection": collection, "limit": limit}
330+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "collection": collection, "limit": limit}
316331

317332
return await backend.run_query_async(query, params)
318333

@@ -329,15 +344,17 @@ async def _query_transitive_callees_async(
329344
"""Async multi-hop callee traversal.
330345
331346
Supports class-level queries: for "MyClass", also matches callees of "MyClass.method".
347+
Also supports suffix matching for symbols stored with full module paths.
332348
"""
333349
safe_depth = max(1, min(10, int(depth)))
334350
symbol_prefix = f"{symbol}."
351+
symbol_suffix = f".{symbol}" # For matching full module paths
335352

336353
if include_paths:
337354
if repo and repo != "*":
338355
query = f"""
339356
MATCH path = (source:Symbol {{collection: $collection}})-[:CALLS*1..{safe_depth}]->(callee:Symbol {{collection: $collection}})
340-
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix)
357+
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix OR source.name ENDS WITH $symbol_suffix)
341358
AND all(r IN relationships(path) WHERE r.collection = $collection AND r.repo = $repo)
342359
WITH callee, path, length(path) as hop
343360
RETURN callee.name as symbol, hop,
@@ -346,11 +363,11 @@ async def _query_transitive_callees_async(
346363
ORDER BY hop
347364
LIMIT $limit
348365
"""
349-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "repo": repo, "collection": collection, "limit": limit}
366+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "repo": repo, "collection": collection, "limit": limit}
350367
else:
351368
query = f"""
352369
MATCH path = (source:Symbol {{collection: $collection}})-[:CALLS*1..{safe_depth}]->(callee:Symbol {{collection: $collection}})
353-
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix)
370+
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix OR source.name ENDS WITH $symbol_suffix)
354371
AND all(r IN relationships(path) WHERE r.collection = $collection)
355372
WITH callee, path, length(path) as hop
356373
RETURN callee.name as symbol, hop,
@@ -359,28 +376,28 @@ async def _query_transitive_callees_async(
359376
ORDER BY hop
360377
LIMIT $limit
361378
"""
362-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "collection": collection, "limit": limit}
379+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "collection": collection, "limit": limit}
363380
else:
364381
if repo and repo != "*":
365382
query = f"""
366383
MATCH path = (source:Symbol {{collection: $collection}})-[:CALLS*1..{safe_depth}]->(callee:Symbol {{collection: $collection}})
367-
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix)
384+
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix OR source.name ENDS WITH $symbol_suffix)
368385
AND all(r IN relationships(path) WHERE r.collection = $collection AND r.repo = $repo)
369386
WITH DISTINCT callee
370387
RETURN callee.name as symbol, callee.repo as repo
371388
LIMIT $limit
372389
"""
373-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "repo": repo, "collection": collection, "limit": limit}
390+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "repo": repo, "collection": collection, "limit": limit}
374391
else:
375392
query = f"""
376393
MATCH path = (source:Symbol {{collection: $collection}})-[:CALLS*1..{safe_depth}]->(callee:Symbol {{collection: $collection}})
377-
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix)
394+
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix OR source.name ENDS WITH $symbol_suffix)
378395
AND all(r IN relationships(path) WHERE r.collection = $collection)
379396
WITH DISTINCT callee
380397
RETURN callee.name as symbol, callee.repo as repo
381398
LIMIT $limit
382399
"""
383-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "collection": collection, "limit": limit}
400+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "collection": collection, "limit": limit}
384401

385402
return await backend.run_query_async(query, params)
386403

@@ -397,31 +414,33 @@ async def _query_dependencies_async(
397414
"""Async query both calls and imports for full dependency analysis.
398415
399416
Supports class-level queries: for "MyClass", also matches dependencies of "MyClass.method".
417+
Also supports suffix matching for symbols stored with full module paths.
400418
"""
401419
safe_depth = max(1, min(10, int(depth)))
402420
symbol_prefix = f"{symbol}."
421+
symbol_suffix = f".{symbol}" # For matching full module paths
403422
_ = include_paths # Reserved for future use
404423

405424
if repo and repo != "*":
406425
query = f"""
407426
MATCH path = (source:Symbol {{collection: $collection}})-[:CALLS|IMPORTS*1..{safe_depth}]->(dep:Symbol {{collection: $collection}})
408-
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix)
427+
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix OR source.name ENDS WITH $symbol_suffix)
409428
AND all(r IN relationships(path) WHERE r.collection = $collection AND r.repo = $repo)
410429
WITH DISTINCT dep
411430
RETURN dep.name as symbol, dep.repo as repo
412431
LIMIT $limit
413432
"""
414-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "repo": repo, "collection": collection, "limit": limit}
433+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "repo": repo, "collection": collection, "limit": limit}
415434
else:
416435
query = f"""
417436
MATCH path = (source:Symbol {{collection: $collection}})-[:CALLS|IMPORTS*1..{safe_depth}]->(dep:Symbol {{collection: $collection}})
418-
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix)
437+
WHERE (source.name = $symbol OR source.name STARTS WITH $symbol_prefix OR source.name ENDS WITH $symbol_suffix)
419438
AND all(r IN relationships(path) WHERE r.collection = $collection)
420439
WITH DISTINCT dep
421440
RETURN dep.name as symbol, dep.repo as repo
422441
LIMIT $limit
423442
"""
424-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "collection": collection, "limit": limit}
443+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "collection": collection, "limit": limit}
425444

426445
return await backend.run_query_async(query, params)
427446

@@ -436,12 +455,14 @@ async def _query_cycles_async(
436455
"""Async detect circular dependencies involving the symbol.
437456
438457
Supports class-level queries: for "MyClass", also detects cycles involving "MyClass.method".
458+
Also supports suffix matching for symbols stored with full module paths.
439459
"""
440460
symbol_prefix = f"{symbol}."
461+
symbol_suffix = f".{symbol}" # For matching full module paths
441462
if repo and repo != "*":
442463
query = """
443464
MATCH path = (s:Symbol {collection: $collection})-[:CALLS*2..10]->(s)
444-
WHERE (s.name = $symbol OR s.name STARTS WITH $symbol_prefix)
465+
WHERE (s.name = $symbol OR s.name STARTS WITH $symbol_prefix OR s.name ENDS WITH $symbol_suffix)
445466
AND all(r IN relationships(path) WHERE r.collection = $collection AND r.repo = $repo)
446467
WITH s, path, length(path) as cycle_length
447468
RETURN [n in nodes(path) | n.name] as cycle_path,
@@ -450,11 +471,11 @@ async def _query_cycles_async(
450471
ORDER BY cycle_length
451472
LIMIT $limit
452473
"""
453-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "repo": repo, "collection": collection, "limit": limit}
474+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "repo": repo, "collection": collection, "limit": limit}
454475
else:
455476
query = """
456477
MATCH path = (s:Symbol {collection: $collection})-[:CALLS*2..10]->(s)
457-
WHERE (s.name = $symbol OR s.name STARTS WITH $symbol_prefix)
478+
WHERE (s.name = $symbol OR s.name STARTS WITH $symbol_prefix OR s.name ENDS WITH $symbol_suffix)
458479
AND all(r IN relationships(path) WHERE r.collection = $collection)
459480
WITH s, path, length(path) as cycle_length
460481
RETURN [n in nodes(path) | n.name] as cycle_path,
@@ -463,7 +484,7 @@ async def _query_cycles_async(
463484
ORDER BY cycle_length
464485
LIMIT $limit
465486
"""
466-
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "collection": collection, "limit": limit}
487+
params = {"symbol": symbol, "symbol_prefix": symbol_prefix, "symbol_suffix": symbol_suffix, "collection": collection, "limit": limit}
467488

468489
return await backend.run_query_async(query, params)
469490

0 commit comments

Comments
 (0)