feat: agentic retrieval - graph-walk, skills, MCP server, and MultiPath path router#274
Open
Naseem77 wants to merge 29 commits into
Open
feat: agentic retrieval - graph-walk, skills, MCP server, and MultiPath path router#274Naseem77 wants to merge 29 commits into
Naseem77 wants to merge 29 commits into
Conversation
Add AgentStep and AgentTrace (Thought→Action→Observation record for the agentic loop), ScoredPath (graph-walk output), and SkillResult (structured output of high-level reasoning skills). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Expose two graph-algorithm helpers used by the graph-walk retrieval strategy: pagerank() for node importance weights and weighted_neighbors() to fetch a node's neighbors with edge weights. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add DynamicGraphWalk (beam search + bidirectional path-finding over the knowledge graph using node weights) and the GraphWalkRetrieval strategy, plus score_path scoring helper. Includes unit tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add AgenticRetrieval, a ReAct-style loop where an LLM iteratively chooses tools (search, cypher, traverse, plus skills) to answer a question, with a guard-railed ToolRegistry. Includes the loop, tool wrappers, prompts, and unit tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add a Skill base plus SKILL_REGISTRY/build_skill factory and five built-in skills: entity_comparison, impact_analysis, contradiction_detection, gap_analysis, and timeline_reconstruction. Each wraps graph queries into a structured SkillResult. Includes unit tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add an async run_skill(name, **params) method on the GraphRAG facade that resolves a registered skill via build_skill and runs it against the graph, returning a SkillResult. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add an MCP server (graphrag_sdk.mcp) with a GraphRAGToolset wrapping retrieval and skills as MCP tools, a runnable __main__ entrypoint, the graphrag-mcp console script, and an optional [mcp] dependency extra. Includes tool unit tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add use_cypher/use_fulltext flags (default True) to discover_entities so its two sub-paths can be toggled independently. Enables the path router to prune entity discovery without changing default behavior. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Add an opt-in path router that lets an agent select which of MultiPath's five retrieval paths (relates, entity_cypher, entity_fulltext, expansion, chunks) to run per question instead of always running all. Provides LLMPathRouter (one cheap LLM call) and HeuristicPathRouter, with parse_paths and all_paths helpers. Wire router=None into MultiPathRetrieval and gate stages 3-6 accordingly; any empty plan, router error, or timeout falls back to all paths so recall is never lost. enable_cypher (text-to-Cypher) stays independent. Includes unit tests covering parsing, routing, fallbacks, and live gating. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Re-export the new agentic retrieval, graph-walk, skills, and path-router symbols (and their model types) from the package __init__ files so they are importable from the top-level graphrag_sdk namespace. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (12)
🚧 Files skipped from review as they are similar to previous changes (12)
📝 WalkthroughWalkthroughThis PR adds shared execution models, graph-walk traversal, agentic retrieval, path routing, built-in skills, MCP serving, and public export wiring across the SDK. ChangesPhase 3: Agentic Retrieval, Graph Walk, Skills, MCP Server, and Path Routing
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 11
🧹 Nitpick comments (5)
graphrag_sdk/tests/test_mcp_tools.py (1)
124-132: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winAdd handler coverage for
get_ontologyto complete the MCP toolset contract.
get_ontologyis included in the public tool list but currently has no handler-level test, so one MCP endpoint path remains unverified.Suggested test addition
class TestHandlers: @@ async def test_run_skill(self, toolset: GraphRAGToolset): out = await toolset.by_name("run_skill").handler( {"skill": "gap_analysis", "params": {}} ) assert json.loads(out)["skill"] == "gap_analysis" + + async def test_get_ontology(self, toolset: GraphRAGToolset): + out = await toolset.by_name("get_ontology").handler({}) + payload = json.loads(out) + assert "entities" in payload🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/tests/test_mcp_tools.py` around lines 124 - 132, The test class is missing coverage for the get_ontology MCP tool handler. Add a new async test method named test_get_ontology that follows the same pattern as test_get_statistics and test_run_skill: call toolset.by_name("get_ontology").handler() with appropriate parameters (likely an empty dict based on the get_statistics pattern), parse the JSON output using json.loads(), and assert on a relevant field from the ontology response to verify the handler works correctly.graphrag_sdk/src/graphrag_sdk/api/main.py (1)
3017-3043: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winAdd
run_skill_sync()for API parity with the sync facade.Line [3017] adds a new async public method, but the sync convenience surface is missing its counterpart in the section that explicitly keeps wrappers in sync.
🔧 Suggested addition
+ def run_skill_sync( + self, + skill: str, + *, + ctx: Context | None = None, + **params: Any, + ) -> SkillResult: + """Synchronous run_skill convenience method. + + Keep in sync with :meth:`run_skill`. + """ + return asyncio.run(self.run_skill(skill, ctx=ctx, **params))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/src/graphrag_sdk/api/main.py` around lines 3017 - 3043, The async method run_skill() has been added to the API but its synchronous counterpart is missing from the sync facade wrapper methods. Add a run_skill_sync() method that provides the same functionality as run_skill() but as a synchronous wrapper, following the established pattern used by other sync wrapper methods in the sync facade class. This method should accept the same parameters (skill, ctx, and **params) and return the same SkillResult type.graphrag_sdk/tests/test_agentic_retrieval.py (1)
64-69: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winAdd a regression case for
CALLin read-only Cypher tests.The write-rejection matrix should include a
CALL ...case so the read-only contract is pinned by tests.🔧 Suggested test update
`@pytest.mark.parametrize`( "cypher", - ["CREATE (n)", "MATCH (n) DELETE n", "MERGE (n)", "MATCH (n) SET n.x = 1"], + [ + "CREATE (n)", + "MATCH (n) DELETE n", + "MERGE (n)", + "MATCH (n) SET n.x = 1", + "CALL custom.writeProcedure()", + ], ) def test_rejects_writes(self, cypher: str): assert not is_read_only_cypher(cypher)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/tests/test_agentic_retrieval.py` around lines 64 - 69, Add a regression test case for the CALL statement to the parametrized test_rejects_writes method. Include a CALL cypher string in the parametrize decorator's cypher parameter list (alongside the existing CREATE, DELETE, MERGE, and SET cases) to ensure that is_read_only_cypher correctly identifies CALL statements as non-read-only operations.graphrag_sdk/src/graphrag_sdk/skills/base.py (1)
46-50: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winUse the
query_raw().result_setcontract directly in_rows.
getattr(..., "result_set", [])hides contract violations and can silently turn integration faults into “no data” results.Proposed patch
async def _rows(self, cypher: str, params: dict[str, Any] | None = None) -> list[list[Any]]: """Run a Cypher query and return its raw ``result_set`` rows.""" try: result = await self._graph.query_raw(cypher, params or {}) - return list(getattr(result, "result_set", []) or []) + return list(result.result_set or []) except Exception as exc: logger.warning("Skill %s query failed: %s", self.name, exc) return []Based on learnings:
query_raw()is expected to always return an object withresult_set, and failures should surface via exceptions.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/src/graphrag_sdk/skills/base.py` around lines 46 - 50, The _rows method currently uses getattr with a fallback to hide missing result_set attributes, which masks contract violations and silently turns integration faults into empty results. Replace the getattr call with direct access to result.result_set, removing the defensive default value. This will ensure that any failure to access result_set properly raises an AttributeError, allowing integration faults to surface as exceptions rather than being silently converted to no data.Source: Learnings
graphrag_sdk/tests/test_skills.py (1)
100-179: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winAdd regression tests for impact dedupe and timeline label safety.
Current tests miss two critical edges: longer-path overwrite in impact ranking and unsafe
labelhandling in timeline reconstruction.Suggested test additions
class TestImpactAnalysis: @@ async def test_requires_entity(self, ctx: Context): skill = ImpactAnalysisSkill(FakeGraphStore()) with pytest.raises(ValueError): await skill.run(ctx) + + async def test_keeps_shortest_distance_when_entity_seen_multiple_paths(self, ctx: Context): + store = FakeGraphStore( + neighbors={ + "root": [("x", 1.0, "R"), ("mid", 10.0, "R")], + "mid": [("x", 10.0, "R")], + }, + pagerank={"x": 1.0, "mid": 1.0}, + ) + skill = ImpactAnalysisSkill(store) + result = await skill.run(ctx, entity="root", max_depth=3) + impacted = {r["entity"]: r["distance"] for r in result.data["impacted"]} + assert impacted["x"] == 1 @@ class TestTimelineReconstruction: @@ async def test_orders_events_by_date(self, ctx: Context): @@ assert order == ["e1", "e2"] assert result.data["num_events"] == 2 + + async def test_rejects_unsafe_label(self, ctx: Context): + skill = TimelineReconstructionSkill(FakeGraphStore()) + with pytest.raises(ValueError): + await skill.run(ctx, label="__Entity__) MATCH (n) RETURN n //")🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/tests/test_skills.py` around lines 100 - 179, Add two regression tests to cover critical edge cases: First, add a test to the TestImpactAnalysis class that verifies ImpactAnalysisSkill correctly handles longer-path overwrites by setting up a graph where an entity can be reached via multiple paths with different distances and asserting that the impact ranking reflects the correct distance calculation. Second, add a test to the TestTimelineReconstruction class that verifies TimelineReconstructionSkill safely handles events without a label field or with unsafe label values by configuring the FakeGraphStore to return events missing the label property and asserting the skill gracefully handles these cases without errors.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@graphrag_sdk/pyproject.toml`:
- Around line 83-84: The graphrag-mcp command entry point defined in
pyproject.toml references the main function in
graphrag_sdk/src/graphrag_sdk/mcp/__main__.py, but this module has optional
dependencies that may not be installed. To fix this, wrap the imports in the
main function (or at the module level) in a try-except block that catches
ImportError exceptions. When an ImportError is caught, display a user-friendly
error message that clearly instructs the user to install the optional MCP
dependencies by running "pip install graphrag-sdk[mcp]" instead of letting the
raw traceback propagate.
In `@graphrag_sdk/src/graphrag_sdk/mcp/server.py`:
- Around line 68-74: The _call_tool function at line 73 awaits
tool.handler(arguments or {}) without exception handling, which allows handler
exceptions to escape the MCP call path. Wrap the tool.handler invocation in a
try-except block to catch any exceptions raised by the handler and return them
as TextContent error responses with an appropriate error message, ensuring
exceptions are handled at the tool level rather than propagating up the call
stack.
In `@graphrag_sdk/src/graphrag_sdk/retrieval/agentic/tools.py`:
- Around line 71-75: The is_read_only_cypher() function currently allows CALL
procedure queries which could invoke write-capable procedures despite the tool's
read-only contract. Add 'CALL' to the _WRITE_KEYWORDS collection that is checked
in the is_read_only_cypher() function to ensure that any Cypher query attempting
to invoke procedures is rejected as a potential write operation, closing the
vulnerability in the read-only gate.
In `@graphrag_sdk/src/graphrag_sdk/retrieval/graph_walk.py`:
- Around line 191-193: The issue is that line 191 in the graph walk expansion
loop is regenerating fwd_frontier from all previously discovered nodes in the
fwd map, causing neighbor_fn to be called repeatedly for the same nodes across
iterations. Instead of repopulating fwd_frontier from the entire fwd map, you
should only use the newly discovered nodes from the previous iteration as the
frontier for the next depth expansion. Modify the logic to track and reuse only
the latest layer of frontier nodes rather than rebuilding from all accumulated
nodes in fwd each iteration.
In `@graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py`:
- Around line 338-350: The variable `chunk_sources` is assigned in both the
retrieve_chunks function call and in the else clause but is never used
afterwards, causing a linter F841 error. Replace `chunk_sources` with `_`
(underscore) in both places where it appears in the tuple unpacking to indicate
it's a deliberately unused variable. This change should be made in both the if
block where retrieve_chunks is awaited and the else block where it's assigned an
empty dict.
- Around line 39-42: The import statement for path_router from
graphrag_sdk.retrieval.strategies is not sorted alphabetically within the import
group. Review all imports from graphrag_sdk.retrieval.strategies submodules in
the file and reorder the path_router import to its correct alphabetical position
relative to other imports in that group. Ensure the import block follows proper
alphabetical sorting conventions to resolve the I001 linter error.
In `@graphrag_sdk/src/graphrag_sdk/retrieval/strategies/path_router.py`:
- Line 67: The regex pattern in the re.search() call in the path_router.py file
exceeds the 100 character line length limit required by ruff linting. Extract
the long regex pattern string into a separate variable defined above the
re.search() call, then pass that variable to re.search() instead of the inline
pattern. This will break the line into multiple shorter lines that comply with
the E501 linting rule while preserving the exact same functionality.
In `@graphrag_sdk/src/graphrag_sdk/skills/contradiction_detection.py`:
- Around line 31-38: The limit parameter extracted and converted to an integer
is not validated before being embedded in the Cypher query string via the
f-string. Add validation after the limit conversion to reject values that are
less than or equal to zero, raising an appropriate exception or setting a
reasonable minimum value. This ensures that only valid positive limits are used
in the Cypher query to prevent the query from returning no results and
incorrectly reporting no contradictions found.
In `@graphrag_sdk/src/graphrag_sdk/skills/impact_analysis.py`:
- Around line 54-58: The condition checking whether to update an entry in the
`impacted` dictionary only compares `path.score` against the existing entry's
score, which can cause a longer-hop route to overwrite a shorter-hop route if it
has a higher score. To fix this, modify the overwrite condition in the logic
around the `impacted[node]` assignment to also consider the distance metric.
Ensure that when consolidating impacted entities, a path is only allowed to
overwrite an existing entry if it has both a higher score AND an equal or
shorter distance (hop count), so that the shortest routes are always preserved.
In `@graphrag_sdk/src/graphrag_sdk/skills/timeline_reconstruction.py`:
- Around line 33-37: The label parameter is being directly interpolated into the
Cypher query string in the f-string that constructs the MATCH clause with MATCH
(e:{label}), creating a Cypher injection vulnerability. Before constructing the
cypher variable, validate and sanitize the label to ensure it contains only safe
identifier characters (alphanumeric and underscores). Restrict the label to a
whitelist of allowed characters or use a validation function that ensures it
matches a safe pattern before embedding it in the query string.
In `@graphrag_sdk/src/graphrag_sdk/storage/graph_store.py`:
- Around line 380-383: The pagerank() method directly interpolates entity_label
and relationship_type into the Cypher query string without sanitization, which
creates a query injection vulnerability. Sanitize both the entity_label and
relationship_type parameters before interpolating them into the query
construction on lines 380-383. You should validate these identifiers against a
whitelist of allowed values or implement proper escaping to prevent malformed or
crafted input from altering the query structure.
---
Nitpick comments:
In `@graphrag_sdk/src/graphrag_sdk/api/main.py`:
- Around line 3017-3043: The async method run_skill() has been added to the API
but its synchronous counterpart is missing from the sync facade wrapper methods.
Add a run_skill_sync() method that provides the same functionality as
run_skill() but as a synchronous wrapper, following the established pattern used
by other sync wrapper methods in the sync facade class. This method should
accept the same parameters (skill, ctx, and **params) and return the same
SkillResult type.
In `@graphrag_sdk/src/graphrag_sdk/skills/base.py`:
- Around line 46-50: The _rows method currently uses getattr with a fallback to
hide missing result_set attributes, which masks contract violations and silently
turns integration faults into empty results. Replace the getattr call with
direct access to result.result_set, removing the defensive default value. This
will ensure that any failure to access result_set properly raises an
AttributeError, allowing integration faults to surface as exceptions rather than
being silently converted to no data.
In `@graphrag_sdk/tests/test_agentic_retrieval.py`:
- Around line 64-69: Add a regression test case for the CALL statement to the
parametrized test_rejects_writes method. Include a CALL cypher string in the
parametrize decorator's cypher parameter list (alongside the existing CREATE,
DELETE, MERGE, and SET cases) to ensure that is_read_only_cypher correctly
identifies CALL statements as non-read-only operations.
In `@graphrag_sdk/tests/test_mcp_tools.py`:
- Around line 124-132: The test class is missing coverage for the get_ontology
MCP tool handler. Add a new async test method named test_get_ontology that
follows the same pattern as test_get_statistics and test_run_skill: call
toolset.by_name("get_ontology").handler() with appropriate parameters (likely an
empty dict based on the get_statistics pattern), parse the JSON output using
json.loads(), and assert on a relevant field from the ontology response to
verify the handler works correctly.
In `@graphrag_sdk/tests/test_skills.py`:
- Around line 100-179: Add two regression tests to cover critical edge cases:
First, add a test to the TestImpactAnalysis class that verifies
ImpactAnalysisSkill correctly handles longer-path overwrites by setting up a
graph where an entity can be reached via multiple paths with different distances
and asserting that the impact ranking reflects the correct distance calculation.
Second, add a test to the TestTimelineReconstruction class that verifies
TimelineReconstructionSkill safely handles events without a label field or with
unsafe label values by configuring the FakeGraphStore to return events missing
the label property and asserting the skill gracefully handles these cases
without errors.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 5abe2bef-a872-4c0e-91be-ff2802139cdf
📒 Files selected for processing (31)
graphrag_sdk/pyproject.tomlgraphrag_sdk/src/graphrag_sdk/__init__.pygraphrag_sdk/src/graphrag_sdk/api/main.pygraphrag_sdk/src/graphrag_sdk/core/models.pygraphrag_sdk/src/graphrag_sdk/mcp/__init__.pygraphrag_sdk/src/graphrag_sdk/mcp/__main__.pygraphrag_sdk/src/graphrag_sdk/mcp/server.pygraphrag_sdk/src/graphrag_sdk/mcp/tools.pygraphrag_sdk/src/graphrag_sdk/retrieval/__init__.pygraphrag_sdk/src/graphrag_sdk/retrieval/agentic/__init__.pygraphrag_sdk/src/graphrag_sdk/retrieval/agentic/loop.pygraphrag_sdk/src/graphrag_sdk/retrieval/agentic/prompts.pygraphrag_sdk/src/graphrag_sdk/retrieval/agentic/tools.pygraphrag_sdk/src/graphrag_sdk/retrieval/graph_walk.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/__init__.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/entity_discovery.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/path_router.pygraphrag_sdk/src/graphrag_sdk/skills/__init__.pygraphrag_sdk/src/graphrag_sdk/skills/base.pygraphrag_sdk/src/graphrag_sdk/skills/contradiction_detection.pygraphrag_sdk/src/graphrag_sdk/skills/entity_comparison.pygraphrag_sdk/src/graphrag_sdk/skills/gap_analysis.pygraphrag_sdk/src/graphrag_sdk/skills/impact_analysis.pygraphrag_sdk/src/graphrag_sdk/skills/timeline_reconstruction.pygraphrag_sdk/src/graphrag_sdk/storage/graph_store.pygraphrag_sdk/tests/test_agentic_retrieval.pygraphrag_sdk/tests/test_graph_walk.pygraphrag_sdk/tests/test_mcp_tools.pygraphrag_sdk/tests/test_path_router.pygraphrag_sdk/tests/test_skills.py
Naseem77
changed the title
The agentic search tool capped MultiPath output at 1500 chars by default, discarding ~95% of retrieved context (e.g. ~29K chars for Novel-Lite) before the agent could read it. MultiPath already bounds its own output via chunk_top_k plus entity/relationship limits, so the hard char cap only crippled answer quality. Default max_chars to None (no truncation); callers can still pass max_chars to force a cap. Lifts let-the-agent-decide quality on Novel-Lite from 4.00 to 6.10 (overall, LLM-as-judge), with no other change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Reorder imports (isort), rename unused chunk_sources to _chunk_sources, and wrap a long regex in path_router. No behavior change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Reformat 14 files to satisfy 'ruff format --check src/' in CI. Formatting only, no behavior change. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py (1)
183-183: 🧹 Nitpick | 🔵 Trivial | 💤 Low valueConsider a Protocol type for the router interface.
The
Anytype hint doesn't communicate the expected interface. AProtocolwould document that routers need an asyncplan(query, ctx) -> Iterable[str]method and would help catch integration bugs when mypy is re-enabled.from typing import Protocol class PathRouter(Protocol): async def plan(self, query: str, ctx: Context) -> list[str]: ...Then use
router: PathRouter | None = None.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py` at line 183, Create a Protocol type called PathRouter that documents the expected interface for router objects by defining an async method plan that takes a query string and context and returns a list of strings. Import Protocol from typing module at the top of the file. Then replace the router parameter type hint from Any to PathRouter in the function/method signature where router is defined as an optional parameter, ensuring type safety and better documentation of the expected router interface.Source: Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@graphrag_sdk/src/graphrag_sdk/retrieval/agentic/tools.py`:
- Around line 22-31: The `_WRITE_KEYWORDS` check uses plain substring matching
with the `in` operator against the lowered query string, which causes false
positives when keywords appear as parts of other words. Replace the substring
matching logic with word boundary-based regex matching to ensure keywords only
match as complete words, not as substrings within other words. This will prevent
legitimate queries containing words like "Asset", "Offset", "Preset", "created",
"deleted", "dropped" from being incorrectly flagged as write operations.
---
Nitpick comments:
In `@graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py`:
- Line 183: Create a Protocol type called PathRouter that documents the expected
interface for router objects by defining an async method plan that takes a query
string and context and returns a list of strings. Import Protocol from typing
module at the top of the file. Then replace the router parameter type hint from
Any to PathRouter in the function/method signature where router is defined as an
optional parameter, ensuring type safety and better documentation of the
expected router interface.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: a8e11344-7125-43f1-b00e-a57826a80503
📒 Files selected for processing (15)
graphrag_sdk/src/graphrag_sdk/mcp/__main__.pygraphrag_sdk/src/graphrag_sdk/mcp/server.pygraphrag_sdk/src/graphrag_sdk/mcp/tools.pygraphrag_sdk/src/graphrag_sdk/retrieval/agentic/loop.pygraphrag_sdk/src/graphrag_sdk/retrieval/agentic/tools.pygraphrag_sdk/src/graphrag_sdk/retrieval/graph_walk.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/entity_discovery.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/path_router.pygraphrag_sdk/src/graphrag_sdk/skills/__init__.pygraphrag_sdk/src/graphrag_sdk/skills/contradiction_detection.pygraphrag_sdk/src/graphrag_sdk/skills/entity_comparison.pygraphrag_sdk/src/graphrag_sdk/skills/gap_analysis.pygraphrag_sdk/src/graphrag_sdk/skills/timeline_reconstruction.pygraphrag_sdk/src/graphrag_sdk/storage/graph_store.py
💤 Files with no reviewable changes (1)
- graphrag_sdk/src/graphrag_sdk/storage/graph_store.py
🚧 Files skipped from review as they are similar to previous changes (11)
- graphrag_sdk/src/graphrag_sdk/retrieval/strategies/entity_discovery.py
- graphrag_sdk/src/graphrag_sdk/skills/gap_analysis.py
- graphrag_sdk/src/graphrag_sdk/skills/contradiction_detection.py
- graphrag_sdk/src/graphrag_sdk/skills/entity_comparison.py
- graphrag_sdk/src/graphrag_sdk/skills/timeline_reconstruction.py
- graphrag_sdk/src/graphrag_sdk/skills/init.py
- graphrag_sdk/src/graphrag_sdk/retrieval/strategies/path_router.py
- graphrag_sdk/src/graphrag_sdk/retrieval/graph_walk.py
- graphrag_sdk/src/graphrag_sdk/retrieval/agentic/loop.py
- graphrag_sdk/src/graphrag_sdk/mcp/tools.py
- graphrag_sdk/src/graphrag_sdk/mcp/server.py
- is_read_only_cypher now blocks every CALL procedure (algo./db./dbms./apoc.) via word-boundary regex, closing a write-capable procedure path behind the read-only tool contract; word boundaries also avoid false positives on identifiers like 'recall'/'asset'. - DynamicGraphWalk bidirectional search no longer repopulates the frontier from the full visited map each depth (which re-expanded every discovered node); _expand_frontier already maintains the next layer in place. Removes the now unused _latest_layer no-op. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…te inputs - timeline_reconstruction: sanitize the user-provided label before interpolating it into MATCH (Cypher injection); use backtick-quoted identifier. - graph_store.pagerank: validate entity_label/relationship_type as identifiers before embedding them in the algo.pageRank procedure string args. - impact_analysis: preserve the shortest-hop distance when consolidating an impacted entity instead of overwriting it whenever a higher-score (but longer) path is seen; track best score independently. - contradiction_detection: reject limit < 1 early. - skills.base: replace no-op '...' abstract body with an explicit NotImplementedError. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
- MCP _call_tool now catches handler exceptions and returns a tool-level error response instead of letting them escape the call path. - Remove unused Context import in test_path_router. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
MultiPath._execute reads chunk_top_k/max_entities/max_relationships/ rel_top_k and output caps from kwargs, falling back to defaults so normal flow is unchanged. result_assembly caps are now parameters. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
search/cypher/traverse tools accept per-call numeric overrides (top_k, max_rows, beam_width, max_depth). ReAct prompt now routes count/list-all queries to cypher and suggests widening limits when search lacks breadth. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py (1)
260-274: 🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy liftAdd the routed empty-result fallback before returning.
The router fallback currently covers empty/invalid plans and planning errors, but not a routed plan that executes successfully and produces no data sections. A narrow custom/LLM plan like
{"entity_cypher"}can miss and return only the hint instead of retrying the default all-path flow, which loses recall despite the PR contract.Consider extracting the planned execution body and retrying with
all_paths()whenself._router is not None, the selected plan is not already all paths, and the assembled result has no non-hint records. Based on PR objectives, routed MultiPath should fall back to all paths on empty results.Also applies to: 407-419
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py` around lines 260 - 274, The routed MultiPath flow in `MultiPath.retrieve` does not retry the default path set when a custom router plan succeeds but yields no real data, so add a post-execution empty-result fallback. Extract the planned execution/assembly logic around the router plan handling, and after running a non-all-path plan, detect when the combined result contains only hint records or no non-hint records; in that case, rerun with `all_paths()` unless the plan was already all paths. Keep the existing behavior for planning errors and `LatencyBudgetExceededError`, and apply the same fix in the second MultiPath branch referenced by the duplicate path handling.
🧹 Nitpick comments (1)
graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py (1)
225-245: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick winClamp agent-tunable limits to safe upper bounds.
_pos_intaccepts any positive integer, so agent/tool kwargs can drive very largetop_k, entity, relationship, rerank, or final-context caps. Add per-key ceilings to prevent accidental expensive vector/graph calls and oversized LLM context.Example direction
- def _pos_int(key: str, default: int) -> int: + max_overrides = { + "chunk_top_k": 100, + "max_entities": 200, + "max_relationships": 200, + "rel_top_k": 100, + "max_cypher_out": 100, + "max_entities_out": 100, + "max_relationships_out": 100, + "max_facts_out": 100, + "max_passages_out": 100, + } + + def _pos_int(key: str, default: int) -> int: val = kwargs.get(key, default) try: val = int(val) except (TypeError, ValueError): return default - return val if val > 0 else default + if val <= 0: + return default + return min(val, max_overrides.get(key, val))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py` around lines 225 - 245, Clamp the agent-tunable limits in `MultiPathStrategy` so `kwargs` cannot request arbitrarily large values. Update the `_pos_int` handling in `multi_path.py` to apply per-key maximum ceilings for `chunk_top_k`, `max_entities`, `max_relationships`, `rel_top_k`, and the final-context caps like `max_cypher_out`, `max_entities_out`, `max_relationships_out`, `max_facts_out`, and `max_passages_out`. Keep the existing positive-int parsing, but bound each resolved value to a safe upper limit before it is used by the retrieval flow.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@graphrag_sdk/src/graphrag_sdk/retrieval/strategies/result_assembly.py`:
- Around line 162-167: In assemble_raw_result, the
max_cypher/max_entities/max_relationships/max_facts/max_passages caps are used
directly for slicing, which allows negative values and still emits empty header
sections when a cap is 0. Add upfront validation in the assemble_raw_result path
to reject negative cap values, then only append each section when its capped
list is non-empty so empty sections are skipped entirely.
---
Outside diff comments:
In `@graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py`:
- Around line 260-274: The routed MultiPath flow in `MultiPath.retrieve` does
not retry the default path set when a custom router plan succeeds but yields no
real data, so add a post-execution empty-result fallback. Extract the planned
execution/assembly logic around the router plan handling, and after running a
non-all-path plan, detect when the combined result contains only hint records or
no non-hint records; in that case, rerun with `all_paths()` unless the plan was
already all paths. Keep the existing behavior for planning errors and
`LatencyBudgetExceededError`, and apply the same fix in the second MultiPath
branch referenced by the duplicate path handling.
---
Nitpick comments:
In `@graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py`:
- Around line 225-245: Clamp the agent-tunable limits in `MultiPathStrategy` so
`kwargs` cannot request arbitrarily large values. Update the `_pos_int` handling
in `multi_path.py` to apply per-key maximum ceilings for `chunk_top_k`,
`max_entities`, `max_relationships`, `rel_top_k`, and the final-context caps
like `max_cypher_out`, `max_entities_out`, `max_relationships_out`,
`max_facts_out`, and `max_passages_out`. Keep the existing positive-int parsing,
but bound each resolved value to a safe upper limit before it is used by the
retrieval flow.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 27724ca0-6f6b-4d18-aafb-beb34edb74dc
📒 Files selected for processing (4)
graphrag_sdk/src/graphrag_sdk/retrieval/agentic/prompts.pygraphrag_sdk/src/graphrag_sdk/retrieval/agentic/tools.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/result_assembly.py
🚧 Files skipped from review as they are similar to previous changes (2)
- graphrag_sdk/src/graphrag_sdk/retrieval/agentic/prompts.py
- graphrag_sdk/src/graphrag_sdk/retrieval/agentic/tools.py
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
… format - mcp/__main__: catch ImportError and tell users to install graphrag-sdk[mcp] instead of crashing with a raw traceback. - result_assembly: clamp max_* caps to >=0 and guard each section on the capped list so 0 omits a section and negatives no longer truncate from tail. - apply ruff format to agentic/tools.py and result_assembly.py. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Naseem77
changed the title
galshubeli
reviewed
Jun 30, 2026
| rag = self.rag | ||
|
|
||
| async def ingest(args: dict[str, Any]) -> str: | ||
| result = await rag.ingest(args["text"], document_id=args.get("document_id")) |
Collaborator
There was a problem hiding this comment.
MCP ingest can never ingest text. GraphRAG.ingest(source, *, text=...) treats the first positional arg as a file path — raw text must go through text=. As written, an MCP client calling ingest with {"text": "..."} makes the loader try to open the prose as a path and the tool returns Error running tool 'ingest'.
Suggested change
| result = await rag.ingest(args["text"], document_id=args.get("document_id")) | |
| result = await rag.ingest(text=args["text"], document_id=args.get("document_id")) |
| """ | ||
| query = ( | ||
| "MATCH (n:__Entity__ {id: $node_id})-[r]-(m:__Entity__) " | ||
| "RETURN m.id AS id, type(r) AS label, " |
Collaborator
There was a problem hiding this comment.
type(r) collapses every edge to "RELATES". The schema stores a single RELATES edge type with the actual relation in the r.rel_type property (see relationship_expansion.py:23,42). Using type(r) here means every ScoredPath.edges label and traverse result reads "RELATES" instead of BORN_IN/WORKS_AT/etc., silently losing the semantic relation.
Suggested change
| "RETURN m.id AS id, type(r) AS label, " | |
| "RETURN m.id AS id, coalesce(r.rel_type, type(r)) AS label, " |
|
|
||
| rows = await self._rows( | ||
| "MATCH (e:__Entity__ {id: $id})-[r]-(m:__Entity__) " | ||
| "RETURN type(r) AS rel, m.id AS other, " |
Collaborator
There was a problem hiding this comment.
Same type(r) issue as in graph_store.weighted_neighbors: every fact line becomes "RELATES <other>" instead of the semantic relation, so a contradiction that hinges on the relation type (e.g. BORN_IN Paris vs BORN_IN London) is invisible to the LLM. Prefer coalesce(r.rel_type, type(r)) AS rel.
| logger = logging.getLogger(__name__) | ||
|
|
||
| _ACTION_RE = re.compile(r"Action\s*:\s*(.+)", re.IGNORECASE) | ||
| _ACTION_INPUT_RE = re.compile(r"Action\s*Input\s*:\s*(\{.*\})", re.IGNORECASE | re.DOTALL) |
Collaborator
There was a problem hiding this comment.
Greedy Action Input capture drops valid tool args. (\{.*\}) with DOTALL spans from the first { to the last } in the whole response. When the model emits anything containing } after the JSON (a trailing Observation: line, prose), json.loads fails → action_input = {} → the search tool then returns "Error: 'query' is required." Verified locally with a trailing Observation: {...}. Consider a non-greedy / brace-balanced parse.
| # word boundaries so substrings of legitimate identifiers (e.g. "recall", | ||
| # "asset") are not falsely flagged. ``call`` blocks every stored-procedure path | ||
| # (algo.*, db.*, dbms.*, apoc.*) since procedures can mutate the graph. | ||
| _WRITE_KEYWORD_RE = re.compile( |
Collaborator
There was a problem hiding this comment.
Read-only gate false-positively rejects valid queries. _WRITE_KEYWORD_RE matches write keywords anywhere in the raw text, including inside string literals. Verified: MATCH (n) WHERE n.name = 'call center' RETURN n and ... CONTAINS "delete" ... are both rejected as non-read-only (by the agent and the MCP cypher_query tool). Note the codebase already has validate_cypher() in cypher_generation.py — worth reusing rather than maintaining a second allow/deny list.
| "via": path.nodes, | ||
| } | ||
| continue | ||
| if hop < existing["distance"]: |
Collaborator
There was a problem hiding this comment.
score and distance/via can come from different paths. distance/via update on hop < existing["distance"] while score updates independently on path.score > existing["score"]. A node first seen at depth 2 (score 0.9) then depth 1 (score 0.3) ends up distance=1, score=0.9 — a mismatched pair that the (distance, -score) ranking key then sorts on. Update score/via together with distance.
| val = int(val) | ||
| except (TypeError, ValueError): | ||
| return default | ||
| return val if val > 0 else default |
Collaborator
There was a problem hiding this comment.
_pos_int coerces 0 to the default, so the documented "0 = omit section" is unreachable. result_assembly.assemble_raw_result clamps negatives and documents 0 = omit it entirely, but every value reaches it through this return val if val > 0 else default. An agent passing max_passages_out=0 to drop a section silently gets the full 15; same for the breadth knobs (chunk_top_k=0, etc.). Allow >= 0 for the output caps.
| if re.search(rel_pattern, q): | ||
| plan |= {"expansion", "entity_cypher"} | ||
| # Quoted or capitalized proper nouns in the original query → name match. | ||
| if '"' in query or "'" in query or re.search(r"\b[A-Z][a-z]+\b", query or ""): |
Collaborator
There was a problem hiding this comment.
Heuristic rarely prunes the entity paths. re.search(r"\b[A-Z][a-z]+\b", query) matches the sentence-cased leading word of nearly every question ("What…", "How…"), so entity_cypher + entity_fulltext get added on almost all queries — undercutting the router's stated purpose of skipping irrelevant paths. Consider excluding a leading stop-word / requiring a capitalized token that isn't the first word.
| self._graph, entity_list, self._max_relationships, ctx=ctx | ||
| ) | ||
| entity_list = list(found_entities.items())[:max_entities] | ||
| if "expansion" in plan: |
Collaborator
There was a problem hiding this comment.
A router plan of expansion alone yields zero relationships. expansion is gated independently of the paths that populate entity_list. parse_paths accepts {"expansion"} with no entity path, leaving entity_list empty so expand_relationships([], …) returns [] — the exact data the path exists to provide is dropped on a "how are X and Y connected" question. HeuristicPathRouter couples them, but nothing in the strategy enforces it; consider auto-including an entity path whenever expansion is selected.
GraphRAG.ingest treats a positional first arg as a file path, so the MCP ingest tool tried to open the prose as a path (and could read an arbitrary local file if the text happened to be a valid path). Pass text= explicitly and align the test fake's signature with the real API so the mismatch can't mask this again. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…pe(r) All edges are stored as a single RELATES type with the semantic relation in the rel_type property, so type(r) collapsed every neighbor label and contradiction fact to 'RELATES'. Use coalesce(r.rel_type, type(r)) in weighted_neighbors and contradiction fact collection so BORN_IN/WORKS_AT/etc. reach the traverse results and the contradiction LLM. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The greedy (\{.*\}) DOTALL capture spanned from the first { to the last }
in the whole completion, so any trailing text containing } (a hallucinated
Observation: line, prose) broke json.loads and the tool silently ran with
empty args. Use json.JSONDecoder.raw_decode from the opening brace so parsing
stops at the object's matching close brace.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Write keywords inside quoted data values ('call center', CONTAINS "delete")
or backtick identifiers falsely tripped the gate. Strip quoted segments
before the keyword scan; an unterminated quote leaves its tail scanned, so
the failure mode stays rejection, never bypass.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
score updated independently of distance/via, so a node could report the distance of one path with the score of another, and the (distance, -score) ranking sorted on that mismatched pair. Replace the record atomically: shortest hop wins, ties broken by higher score. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The documented '0 = omit this section' in result_assembly was unreachable: _pos_int coerced 0 back to the default before it ever reached assembly. Use a non-negative parser for the five *_out display caps; breadth knobs (chunk_top_k, max_entities, ...) remain strictly positive. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
A plan of {expansion} alone left entity_list empty, so expand_relationships
returned nothing — the exact data the path exists to provide. Enforce the
invariant where the plan is applied (covers every router): if expansion is
selected without a seed-producing path, add entity_cypher. Also note the
coupling in the router's path guide so the LLM picks compatible sets.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
…estion
\b[A-Z][a-z]+\b matched the sentence-cased leading word ('What', 'Who',
'How'...), so entity_cypher+entity_fulltext were added to nearly all queries
and the router never pruned. Only count capitalized tokens after the first
word as proper-noun signals; quoted names still qualify.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
A beam entry whose tail had no unvisited neighbors was silently dropped: it produced no candidates and only goal-reaching paths were retained. A high-value leaf one hop from the start vanished while deeper branches survived. Record unextendable paths as terminal and include them in goal-less results (skipping the all-dead-end round, where the surviving beam already covers them). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Phase-3 work adding agentic retrieval on top of the existing MultiPath pipeline, plus supporting graph algorithms, a skills framework, an MCP server, and a new agentic path router.
Split into 10 focused commits (dependency-ordered; every commit is importable, exports are wired last).
What's included
AgentStep,AgentTrace,ScoredPath,SkillResult.GraphStore.pagerank()andweighted_neighbors().DynamicGraphWalk(beam + bidirectional) andGraphWalkRetrieval.AgenticRetrieval, a ReAct loop where the LLM picks tools (search,cypher,traverse, + skills) via a guard-railedToolRegistry.Skillbase + registry and 5 built-ins (entity_comparison, impact_analysis, contradiction_detection, gap_analysis, timeline_reconstruction).GraphRAG.run_skill(name, **params).graphrag-mcpconsole script, optional[mcp]extra.relates,entity_cypher,entity_fulltext,expansion,chunks) to run per question.LLMPathRouter+HeuristicPathRouter, with all-paths fallback on empty/error/timeout so recall is never lost. Default behavior unchanged.graphrag_sdk.Testing
New unit tests:
test_graph_walk.py,test_agentic_retrieval.py,test_skills.py,test_mcp_tools.py,test_path_router.py.Full suite: 1083 passed, 29 skipped (1 pre-existing, unrelated LiteLLM timing test failure; no provider code touched here).
Summary by CodeRabbit
graphrag-mcpCLI plus an MCP server with stdio and SSE transports.run_skill()support, and expanded SDK exports.max_*limits).