Skip to content

feat: Agentic GraphRAG toolkit (graphrag_sdk.tools) - #276

Open
gkorland wants to merge 10 commits into
mainfrom
feat/agent-toolkit
Open

feat: Agentic GraphRAG toolkit (graphrag_sdk.tools)#276
gkorland wants to merge 10 commits into
mainfrom
feat/agent-toolkit

Conversation

@gkorland

@gkorland gkorland commented Jul 12, 2026

Copy link
Copy Markdown
Contributor

Summary

Implements the 2026-Q3 roadmap item Agentic GraphRAG (Task 01): a framework-neutral agent surface so any agent framework (pydantic-ai, LangGraph, CrewAI, MCP) can expose a knowledge graph as tools. Targets v1.4.0; pydantic-ai-falkordb (Task 04) hard-depends on this release and on the tool_specs() contract shipped here.

  • GraphRAGToolkit — async search / answer / schema / entity / guarded cypher_read / remember / flush; knobs finalize_policy= (manual/on_write/never), read_only=, include=; generic call(name, arguments) dispatch; for_tenant(...) with owned-connection lifecycle.
  • tool_specs() — machine-readable tool definitions (name, LLM-ready description, Draft 2020-12 JSON-Schema input with additionalProperties: false, output hint). Single source of truth for adapters; JSON round-trip + meta-validation tested.
  • Typed results — pydantic v2 models with deterministic, budget-bounded to_llm_text(max_chars=…) (truncation only at item boundaries with an explicit …(N more) marker; control chars stripped via the ingestion sanitizer). Answers carry document_id/chunk_id citations.
  • Fail-closed Cypher guard — single-pass lexer (string/comment interleaving is exploitable with ordered regex stripping), dual raw+NFKC scan, full-name read-safe procedure allowlist, LIMIT injection, per-query timeout_ms; raises ReadOnlyViolation naming the offending token. ~40-case test table.
  • MultiPathRetrieval provenance (the one change under retrieval/)_execute now stashes {entities, chunks(+document_path), facts, relationships} into RetrieverResult.metadata["provenance"], captured before [Source:] tagging. Additive and behavior-preserving: a byte-stability test pins the exact section rendering before/after; this lets search()/answer() ride the real pipeline instead of forking it. Approved deviation from the task's out-of-scope fence (owner ruling, 2026-07-12).
  • Docs: docs/agentic.md (in mkdocs nav, --strict verified), README "Use with Agents" section + example row, examples/11_agent_toolkit.py, CHANGELOG under Unreleased.
  • Drive-by fixes: stale __version__ = "1.2.0"1.3.0 (+ pyproject drift test; the value is stamped into graphs), pre-existing broken doc link in incremental-updates.md that failed mkdocs build --strict, README MENTIONSMENTIONED_IN.

Test plan

  • Unit suite: 1114 passed, 33 skipped (pytest tests/ -q, py3.12)
  • Integration vs live FalkorDB, CI-exact command incl. the new file: 24 passed, 1 skipped (RUN_INTEGRATION=1 pytest -m integration tests/test_integration.py tests/test_tools_integration.py)
  • ruff check src/ + ruff format --check src/ clean (CI-exact); new test files clean
  • mypy --python-version=3.12 clean on tools/, core/exceptions.py, multi_path.py
  • mkdocs build --strict passes
  • tool_specs() JSON round-trip smoke

Notes for reviewers

  • Benchmark (CONTRIBUTING 85% rule) was not run: the only retrieval-path change is metadata-additive and byte-stability-tested; happy to run the corpus if you want it on record.
  • The repo mypy config (python_version = "3.10") fails on installed numpy 2.x stubs for any module importing numpy (pre-existing, repo-wide); checks here used --python-version=3.12.
  • examples/11_agent_toolkit.py compiles and mirrors the live-tested integration flow, but was not executed end-to-end in CI/dev (no OPENAI_API_KEY available).
  • Known gaps left as-is: docs/incremental-updates.md and docs/ontology-discovery.md are still absent from mkdocs nav; ~136 pre-existing ruff findings in old tests/ files (CI doesn't lint tests).
  • Deferred by design: structured-upsert API (upsert_delta/KnowledgeDelta) — Task 04 will file an upstream issue; AnswerResult.cypher_used is always None until the text-to-Cypher path surfaces its query.

🤖 Generated with Claude Code

Summary by CodeRabbit

  • New Features

    • Added an agent-ready GraphRAG toolkit with search, answers, schema, entity lookup, ingestion, flushing, and tenant-scoped access.
    • Added machine-readable tool specifications and typed, citation-ready results with deterministic text truncation.
    • Added guarded read-only Cypher operations with automatic limits and clear safety errors.
    • Search results now expose entity, relationship, fact, chunk, and document provenance.
  • Documentation

    • Added Agentic GraphRAG guidance, quickstarts, and an example.
  • Bug Fixes

    • Corrected the reported package version and added safeguards against version drift.

gkorland and others added 10 commits July 12, 2026 13:25
__version__ was left at 1.2.0 when 1.3.0 shipped; the value is stamped
into every graph as __GraphRAGConfig__.sdk_version, so 1.3.0 installs
were mislabeling graphs. Adds a drift test against pyproject.toml.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Typed pydantic-v2 results for the agent toolkit (SearchResult,
AnswerResult, SchemaResult, CypherResult, EntityResult, RememberResult
+ shared parts). to_llm_text() renders deterministic plain text bounded
by max_chars, truncating only at item boundaries with an explicit
…(N more) marker; every string passes the ingestion control-char
sanitizer. Adds ReadOnlyViolation to core exceptions. Golden files pin
the rendering contract.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Single-pass lexer strips comments and masks string/backtick literals
(regex-order stripping is exploitable via quote/comment interleaving),
then scans both the raw and NFKC-normalized copies for write keywords,
multi-statement separators, disallowed start keywords, and CALLs to
procedures outside a full-name read-safe allowlist. Rejections raise
ReadOnlyViolation naming the offending token. apply_limit() injects
LIMIT only when the (noise-stripped) query lacks one. The original
query text — never the normalized copy — is what executes.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…adata

RetrieverResult.metadata now carries a 'provenance' dict with entity
ids/names/descriptions, kept chunk ids/texts/document paths (in rerank
order, captured before [Source:] tagging), fact strings, and
relationship strings. Additive only: assemble_raw_result inputs are
untouched and a byte-stability test pins the exact section rendering
for a fixed input both before and after this change. Consumed by the
graphrag_sdk.tools agent toolkit for citations. Approved deviation from
the Task-01 out-of-scope fence (user ruling 2026-07-12).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Seven pydantic input models (extra='forbid' -> additionalProperties:
false) and a static registry binding tool names to toolkit methods with
LLM-ready descriptions and output hints. build_tool_specs() filters by
read_only / finalize_policy / include. Adapters and MCP servers generate
their tool definitions from this — the single source of truth. Adds
jsonschema (dev) for Draft 2020-12 meta-validation in tests.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The seven async operations (search/answer/schema/entity/cypher_read/
remember/flush) over a wrapped GraphRAG. search()/answer() drive a
per-call tuned MultiPathRetrieval through the real retrieve()/
completion() paths and map provenance metadata to typed results with
document/chunk-id citations. cypher_read guards via ensure_read_only +
LIMIT injection and enforces timeout_ms on the connection. call(name,
arguments) validates through the registry input models for adapters/
MCP. for_tenant() derives an owned tenant-scoped graph (validated
tenant_id); aclose()/async-with manage its lifecycle. Root package
re-exports GraphRAGToolkit/ToolSpec/ReadOnlyViolation. graph_ops
centralizes toolkit-owned parameterized Cypher; entity/chunk bulky
props (embedding, source_chunk_ids) never reach agent output.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Covers falkordb Node/Edge conversion with real driver objects (embedding
and source_chunk_ids stripped, truncated heuristic), frontier expansion
dedupe/caps, guard-before-connection ordering, LIMIT injection with
params/timeout forwarding, ontology+live-count merging, and ranked
entity matching with nearby suggestions.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Asserts per-call MultiPathRetrieval tuning (chunk_top_k/rel_top_k/
max_entities from top_k), provenance-order entity mapping with
enrichment fallback, top_k caps, include_chunks/expand_hops behavior,
citation construction (chunk id, document id/path, 200-char snippet),
case-insensitive entities_touched dedupe, and the graceful
no-provenance degradation path.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Round-trip over 3 scripted docs + remember/flush: typed search with
document ids, entity card with WORKS_AT neighbors and provenance docs,
schema counts, parameterized cypher_read with write rejection and live
LIMIT injection, answer() citations, and for_tenant graph isolation
with owned-connection cleanup. Real-LLM smoke stays behind
OPENAI_API_KEY. The CI integration job now runs this file alongside
test_integration.py (its test path is hardcoded).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
docs/agentic.md covers the tool surface, tool_specs contract for
adapter authors, finalize policies (finalize is O(graph size)), the
cypher_read guard semantics + allowlist and the treat-agent-Cypher-as-
untrusted guidance, and for_tenant. Wired into mkdocs nav (strict
build verified). README gains a 'Use with Agents' section + example
table row and marks the 2026-Q3 roadmap item; fixes the MENTIONS ->
MENTIONED_IN edge name. Also repairs a pre-existing broken relative
link in incremental-updates.md that failed mkdocs --strict (the docs
workflow triggers on any docs/** change).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@coderabbitai

coderabbitai Bot commented Jul 12, 2026

Copy link
Copy Markdown

Review Change Stack

📝 Walkthrough

Walkthrough

Adds GraphRAGToolkit, typed tool results and schemas, guarded read-only Cypher, retrieval provenance, tenant lifecycle support, documentation, examples, and extensive unit and real-backend integration tests.

Changes

Agentic GraphRAG toolkit

Layer / File(s) Summary
Tool contracts and guarded result surface
graphrag_sdk/src/graphrag_sdk/tools/models.py, specs.py, cypher_guard.py, core/exceptions.py
Adds Pydantic result/input models, deterministic bounded LLM text rendering, adapter tool specifications, and fail-closed Cypher validation with ReadOnlyViolation.
Graph operations and retrieval provenance
graphrag_sdk/src/graphrag_sdk/tools/graph_ops.py, retrieval/strategies/multi_path.py
Adds graph enrichment, traversal, document mapping, schema counts, JSON-safe query conversion, and structured retrieval provenance.
Toolkit lifecycle and operations
graphrag_sdk/src/graphrag_sdk/tools/toolkit.py, tools/__init__.py, graphrag_sdk/src/graphrag_sdk/__init__.py
Adds tenant-scoped toolkit construction, dispatch, lifecycle handling, search, answer, schema, entity, Cypher, ingest, flush, and public exports.
Behavioral and rendering validation
graphrag_sdk/tests/test_tools_*.py, graphrag_sdk/tests/test_multi_path_retrieval.py, graphrag_sdk/tests/golden/tools/*
Covers guard behavior, graph helpers, toolkit policies, provenance mapping, tool schemas, deterministic rendering, and golden outputs.
Integration coverage and release wiring
graphrag_sdk/tests/test_tools_integration.py, graphrag_sdk/examples/11_agent_toolkit.py, docs/*, README.md, CHANGELOG.md, .github/workflows/ci.yml
Adds real-FalkorDB integration tests, an agent toolkit example, documentation and navigation, release notes, version synchronization, dependency configuration, and CI execution.

Estimated code review effort: 5 (Critical) | ~120 minutes

Sequence Diagram(s)

sequenceDiagram
  participant Agent
  participant GraphRAGToolkit
  participant GraphRAG
  participant FalkorDB
  Agent->>GraphRAGToolkit: call graph_search or graph_answer
  GraphRAGToolkit->>GraphRAG: retrieve or complete with provenance
  GraphRAG->>FalkorDB: execute graph queries
  FalkorDB-->>GraphRAGToolkit: graph data and provenance
  GraphRAGToolkit-->>Agent: typed result and LLM text
Loading

Possibly related PRs

Suggested reviewers: galshubeli, Naseem77

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 27.74% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed Title clearly names the main change: adding the Agentic GraphRAG toolkit in graphrag_sdk.tools.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch feat/agent-toolkit

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.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

Comment on lines +26 to +27
"UNWIND $kws AS kw CALL { WITH kw MATCH (e:__Entity__) WHERE e.name = kw "
"RETURN e LIMIT 1 } RETURN e", # CALL subquery
assert store.query_raw.await_count == 1

store.query_raw.reset_mock()
triples = await expand_triples(store, ["e1"], hops=2, cap=10)
),
)

TOOL_NAMES: tuple[str, ...] = tuple(td.name for td in _TOOL_REGISTRY)

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (4)
graphrag_sdk/src/graphrag_sdk/tools/graph_ops.py (1)

148-177: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Cross-module use of GraphStore._sanitize_string (private member).

_to_jsonable depends on a leading-underscore method of a class in a different module. Consider exposing it as a public helper (or a shared sanitizer utility) so this dependency is part of GraphStore's supported contract rather than an implicit coupling to its internals.

#!/bin/bash
rg -n '_sanitize_string' graphrag_sdk/src/graphrag_sdk/storage/graph_store.py
🤖 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/tools/graph_ops.py` around lines 148 - 177,
Replace the cross-module calls to the private GraphStore._sanitize_string used
by _to_jsonable with a public or shared sanitizer helper, and update GraphStore
and _to_jsonable to use that supported symbol consistently. Preserve the
existing string-sanitization behavior.
graphrag_sdk/src/graphrag_sdk/tools/toolkit.py (1)

321-353: 🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win

Independent awaits could run concurrently.

schema()'s three calls (get_ontology, schema_counts, get_statistics) and entity()'s two calls (expand_triples, entity_documents) are all independent but awaited sequentially. asyncio.gather would reduce round-trip latency for these tools.

⚡ Proposed concurrency for `entity()`
-    neighbors = await graph_ops.expand_triples(store, [eid], hops=inp.hops)
-    documents = await graph_ops.entity_documents(store, eid)
+    neighbors, documents = await asyncio.gather(
+        graph_ops.expand_triples(store, [eid], hops=inp.hops),
+        graph_ops.entity_documents(store, eid),
+    )

(requires import asyncio at module top; same pattern applies to schema()'s three calls.)

Also applies to: 376-393

🤖 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/tools/toolkit.py` around lines 321 - 353,
Update schema() to run get_ontology(), graph_ops.schema_counts(), and
get_statistics() concurrently with asyncio.gather, then unpack the results
before building the response. Apply the same concurrency pattern in entity() to
expand_triples() and entity_documents(), preserving their existing result
handling and output behavior.
docs/agentic.md (1)

167-167: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Use a relative path for the example link for consistency with the README.

The README links the same example with a relative path (graphrag_sdk/examples/11_agent_toolkit.py), but this doc uses an absolute GitHub URL. Relative paths are more robust to repo moves and render correctly on GitHub.

🔗 Suggested fix
-See [`examples/11_agent_toolkit.py`](https://github.com/FalkorDB/GraphRAG-SDK/blob/main/graphrag_sdk/examples/11_agent_toolkit.py)
+See [`examples/11_agent_toolkit.py`](../graphrag_sdk/examples/11_agent_toolkit.py)
🤖 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 `@docs/agentic.md` at line 167, Update the example link in docs/agentic.md to
use the repository-relative path graphrag_sdk/examples/11_agent_toolkit.py,
matching the README link, and remove the absolute GitHub URL.
graphrag_sdk/tests/test_tools_toolkit.py (1)

91-99: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Narrow the exception assertion for the invalid-arg case.

pytest.raises(Exception) on line 98 would also pass if call() raised something other than a validation error (e.g. an AttributeError from a regression elsewhere), masking a real bug. Prefer asserting on pydantic.ValidationError specifically.

♻️ Suggested tightening
+from pydantic import ValidationError
+
 async def test_call_dispatch_and_validation():
     rag = make_stub_rag()
     tk = GraphRAGToolkit(rag)
     result = await tk.call("graph_remember", {"text": "hi", "document_id": "d"})
     assert isinstance(result, RememberResult)
     with pytest.raises(ValueError, match="Unknown tool"):
         await tk.call("nope", {})
-    with pytest.raises(Exception):  # pydantic ValidationError on extra arg
+    with pytest.raises(ValidationError):
         await tk.call("graph_remember", {"text": "hi", "bogus": 1})
🤖 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_tools_toolkit.py` around lines 91 - 99, Update
test_call_dispatch_and_validation to import and assert
pytest.raises(pydantic.ValidationError) for the graph_remember call with the
extra bogus argument, replacing the broad Exception assertion while leaving the
other dispatch and unknown-tool checks unchanged.
🤖 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/tools/toolkit.py`:
- Around line 195-206: Add read-only GraphRAG accessors for the connection,
graph store, vector store, and global ontology, then update GraphRAGToolkit
methods including _make_strategy to use those accessors instead of _conn,
_graph_store, _vector_store, and _global_ontology. Preserve the existing
returned handles and behavior.
- Around line 321-353: The schema method accepts ctx but ignores it; remove the
unused ctx parameter from schema and update its callers accordingly, unless the
downstream APIs support context propagation, in which case pass ctx through each
relevant call. Ensure the resulting API does not imply that tenant or logging
context affects schema retrieval.

---

Nitpick comments:
In `@docs/agentic.md`:
- Line 167: Update the example link in docs/agentic.md to use the
repository-relative path graphrag_sdk/examples/11_agent_toolkit.py, matching the
README link, and remove the absolute GitHub URL.

In `@graphrag_sdk/src/graphrag_sdk/tools/graph_ops.py`:
- Around line 148-177: Replace the cross-module calls to the private
GraphStore._sanitize_string used by _to_jsonable with a public or shared
sanitizer helper, and update GraphStore and _to_jsonable to use that supported
symbol consistently. Preserve the existing string-sanitization behavior.

In `@graphrag_sdk/src/graphrag_sdk/tools/toolkit.py`:
- Around line 321-353: Update schema() to run get_ontology(),
graph_ops.schema_counts(), and get_statistics() concurrently with
asyncio.gather, then unpack the results before building the response. Apply the
same concurrency pattern in entity() to expand_triples() and entity_documents(),
preserving their existing result handling and output behavior.

In `@graphrag_sdk/tests/test_tools_toolkit.py`:
- Around line 91-99: Update test_call_dispatch_and_validation to import and
assert pytest.raises(pydantic.ValidationError) for the graph_remember call with
the extra bogus argument, replacing the broad Exception assertion while leaving
the other dispatch and unknown-tool checks unchanged.
🪄 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: ddaec9c2-c63d-431c-bcbf-a69fc3a6f62d

📥 Commits

Reviewing files that changed from the base of the PR and between 0ab92ba and 2f03967.

📒 Files selected for processing (32)
  • .github/workflows/ci.yml
  • CHANGELOG.md
  • README.md
  • docs/agentic.md
  • docs/incremental-updates.md
  • graphrag_sdk/examples/11_agent_toolkit.py
  • graphrag_sdk/pyproject.toml
  • graphrag_sdk/src/graphrag_sdk/__init__.py
  • graphrag_sdk/src/graphrag_sdk/core/exceptions.py
  • graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py
  • graphrag_sdk/src/graphrag_sdk/tools/__init__.py
  • graphrag_sdk/src/graphrag_sdk/tools/cypher_guard.py
  • graphrag_sdk/src/graphrag_sdk/tools/graph_ops.py
  • graphrag_sdk/src/graphrag_sdk/tools/models.py
  • graphrag_sdk/src/graphrag_sdk/tools/specs.py
  • graphrag_sdk/src/graphrag_sdk/tools/toolkit.py
  • graphrag_sdk/tests/golden/tools/answer_result.txt
  • graphrag_sdk/tests/golden/tools/cypher_result.txt
  • graphrag_sdk/tests/golden/tools/entity_result.txt
  • graphrag_sdk/tests/golden/tools/remember_result.txt
  • graphrag_sdk/tests/golden/tools/schema_result.txt
  • graphrag_sdk/tests/golden/tools/search_result.txt
  • graphrag_sdk/tests/golden/tools/search_result_truncated.txt
  • graphrag_sdk/tests/test_multi_path_retrieval.py
  • graphrag_sdk/tests/test_tools_cypher_guard.py
  • graphrag_sdk/tests/test_tools_graph_ops.py
  • graphrag_sdk/tests/test_tools_integration.py
  • graphrag_sdk/tests/test_tools_models.py
  • graphrag_sdk/tests/test_tools_specs.py
  • graphrag_sdk/tests/test_tools_toolkit.py
  • graphrag_sdk/tests/test_version_sync.py
  • mkdocs.yml

Comment on lines +195 to +206
def _make_strategy(self, top_k: int) -> MultiPathRetrieval:
"""A per-call MultiPathRetrieval tuned to top_k (ctor is pure assignment)."""
return MultiPathRetrieval(
graph_store=self._rag._graph_store,
vector_store=self._rag._vector_store,
embedder=self._rag.embedder,
llm=self._rag.llm,
chunk_top_k=top_k,
rel_top_k=top_k,
max_entities=max(2 * top_k, 10),
ontology=self._rag._global_ontology,
)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
rg -n -A3 '`@property`' graphrag_sdk/src/graphrag_sdk/api/main.py
rg -n 'self\._graph_store|self\._vector_store|self\._global_ontology|self\._conn\b|def embedder|def llm' graphrag_sdk/src/graphrag_sdk/api/main.py

Repository: FalkorDB/GraphRAG-SDK

Length of output: 8806


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Map the relevant symbols in GraphRAG and the toolkit.
ast-grep outline graphrag_sdk/src/graphrag_sdk/api/main.py --view expanded
echo "---- TOOLKIT ----"
ast-grep outline graphrag_sdk/src/graphrag_sdk/tools/toolkit.py --view expanded

echo "---- PUBLIC/PRIVATE ACCESSORS IN GraphRAG ----"
rg -n -A3 -B2 '^(    `@property`|    def (ontology|schema|embedder|llm|graph_store|vector_store|connection)\b)' graphrag_sdk/src/graphrag_sdk/api/main.py

echo "---- PRIVATE ATTR USAGE IN TOOLKIT ----"
rg -n 'self\._rag\._graph_store|self\._rag\._vector_store|self\._rag\._global_ontology|self\._rag\._conn\b|self\._rag\.(embedder|llm|ontology|schema)|self\._[a-zA-Z_]+' graphrag_sdk/src/graphrag_sdk/tools/toolkit.py

Repository: FalkorDB/GraphRAG-SDK

Length of output: 8252


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Show the exact GraphRAG property/attribute region with line numbers.
sed -n '220,340p' graphrag_sdk/src/graphrag_sdk/api/main.py | cat -n

echo "----"
# Show the toolkit region around _make_strategy and the repeated call sites.
sed -n '180,240p' graphrag_sdk/src/graphrag_sdk/tools/toolkit.py | cat -n
echo "----"
sed -n '280,390p' graphrag_sdk/src/graphrag_sdk/tools/toolkit.py | cat -n

Repository: FalkorDB/GraphRAG-SDK

Length of output: 14421


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the public properties around ontology/schema and the internal fields they expose.
sed -n '300,445p' graphrag_sdk/src/graphrag_sdk/api/main.py | cat -n

# Show the toolkit's repeated internal field access in context.
sed -n '190,235p' graphrag_sdk/src/graphrag_sdk/tools/toolkit.py | cat -n
sed -n '285,385p' graphrag_sdk/src/graphrag_sdk/tools/toolkit.py | cat -n

Repository: FalkorDB/GraphRAG-SDK

Length of output: 15107


Expose the remaining GraphRAG internals
GraphRAGToolkit still reads GraphRAG’s private _conn, _graph_store, _vector_store, and _global_ontology in several methods. Add read-only accessors for those handles so the toolkit doesn’t depend on internal attribute names.

🤖 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/tools/toolkit.py` around lines 195 - 206, Add
read-only GraphRAG accessors for the connection, graph store, vector store, and
global ontology, then update GraphRAGToolkit methods including _make_strategy to
use those accessors instead of _conn, _graph_store, _vector_store, and
_global_ontology. Preserve the existing returned handles and behavior.

Comment on lines +321 to +353
async def schema(self, *, ctx: Context | None = None) -> SchemaResult:
"""Entity labels + relation types with declared metadata and live counts."""
ontology = await self._rag.get_ontology()
label_counts, rel_counts = await graph_ops.schema_counts(self._rag._graph_store)
stats = await self._rag.get_statistics()
declared_e = {e.label: e for e in ontology.entities}
declared_r = {r.label: r for r in ontology.relations}
entity_infos = [
EntityTypeInfo(
label=label,
description=declared_e[label].description if label in declared_e else None,
count=count,
properties=[a.name for a in declared_e[label].properties]
if label in declared_e
else [],
)
for label, count in sorted({**{e: 0 for e in declared_e}, **label_counts}.items())
]
relation_infos = [
RelationTypeInfo(
label=label,
description=declared_r[label].description if label in declared_r else None,
patterns=list(declared_r[label].patterns) if label in declared_r else [],
count=count,
)
for label, count in sorted({**{r: 0 for r in declared_r}, **rel_counts}.items())
]
return SchemaResult(
entities=entity_infos,
relations=relation_infos,
node_count=int(stats.get("node_count", 0)),
edge_count=int(stats.get("edge_count", 0)),
)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
rg -n -A6 'async def finalize|async def get_ontology|async def get_statistics' graphrag_sdk/src/graphrag_sdk/api/main.py
rg -n -A6 'async def query' graphrag_sdk/src/graphrag_sdk/core/connection.py

Repository: FalkorDB/GraphRAG-SDK

Length of output: 1480


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the relevant methods and surrounding context.
sed -n '321,440p' graphrag_sdk/src/graphrag_sdk/tools/toolkit.py

# Find the Context type and any helper that creates a new Context.
rg -n -A3 -B3 'class Context|def _new_ctx|ctx = ctx or self\._new_ctx\(\)|ctx:' graphrag_sdk/src/graphrag_sdk -g '!**/__pycache__/**'

Repository: FalkorDB/GraphRAG-SDK

Length of output: 50377


ctx is accepted here but never used in these tools.
Either thread it through the downstream calls or remove the parameter so callers don’t assume tenant/logging context has any effect.

🤖 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/tools/toolkit.py` around lines 321 - 353, The
schema method accepts ctx but ignores it; remove the unused ctx parameter from
schema and update its callers accordingly, unless the downstream APIs support
context propagation, in which case pass ctx through each relevant call. Ensure
the resulting API does not imply that tenant or logging context affects schema
retrieval.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant