feat(schema): schema-driven attributes end-to-end with persistent ontology#256
Merged
Conversation
…ology Wires declared PropertyType / RelationType.properties from the GraphSchema through extraction, storage, and retrieval. The ontology persists in a dedicated <data_graph>__ontology FalkorDB graph and accumulates as the union of every schema ever registered, so ingest passes can use subset/ extension schemas while retrieval always sees the global ontology. - core/models: PropertyType type-normalising validator, RelationType.properties, ExtractedEntity/Relation.attributes, reserved-name rejection, GraphSchema.merge. - storage/ontology_store (new): OntologyStore.load/register/clear against <data_graph>__ontology; idempotent unions; pattern lists deduped. - ingestion/graph_extraction: conditional attribute block in VERIFY_EXTRACT_RELS_PROMPT (empty for property-less schemas, zero drift), per-type coercion at the extractor boundary, required-missing record drop, aggregators carry attributes (last-write-wins), conversion merges into props. - ingestion/pipeline: _validate_attributes pass after _prune, then re-run _filter_quality to cascade-clean dangling relationships. - retrieval/cypher_generation: replaces hardcoded SCHEMA_PROMPT with build_schema_prompt(schema, question) + render_schema_block; synthesises one numeric-attribute example per declared INTEGER/FLOAT property; validate_cypher accepts dynamic schema labels. - retrieval/multi_path: schema threaded into execute_cypher_retrieval. - api/main: GraphRAG owns the OntologyStore, registers local schema on each ingest, refresh_ontology() propagates the global ontology to retrieval, public get_ontology() / refresh_ontology() methods. Forward-only evolution: nodes ingested before a property was declared remain without that property; FalkorDB's MERGE ... SET n += props handles fill-in naturally and WHERE p.attr > N naturally excludes them. Tests: 124 new/extended unit tests; full suite 819 passed, 23 skipped (integration tests gated by RUN_INTEGRATION=1). Property-less schemas produce identical output to the prior pipeline. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.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:
📝 WalkthroughWalkthroughSwitches GraphRAG from user-supplied GraphSchema to a persisted Ontology (OntologyStore), adds typed Attribute modeling and ontology-aware extraction/coercion, updates ingestion pipeline and retrieval (Cypher) to respect ontology, propagates ontology into retrieval wiring, exposes async facade ontology APIs, and updates tests/exports. ChangesOntology-backed extraction and retrieval
Sequence Diagram(s)sequenceDiagram
participant Client
participant GraphRAG
participant OntologyStore
participant MultiPathRetrieval
participant GraphStore
Client->>GraphRAG: call get_ontology() / ingest()
GraphRAG->>OntologyStore: load() / register(ontology)
OntologyStore-->>GraphRAG: Ontology
GraphRAG->>MultiPathRetrieval: set _schema/_ontology
Client->>MultiPathRetrieval: retrieve(question)
MultiPathRetrieval->>GraphRAG: generate_cypher(question, ontology)
GraphRAG->>GraphStore: query_raw(cypher)
GraphStore-->>GraphRAG: results
GraphRAG-->>Client: facts/completions
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 4
🧹 Nitpick comments (1)
graphrag_sdk/tests/test_ontology_store.py (1)
116-166: ⚡ Quick winAdd a regression test for same-named properties across different owners.
Please add a test that registers two types (e.g., entity
Personand relationWORKS_AT) both with property"since"and asserts independent metadata persistence. This will catch owner-scoping regressions in ontology property upserts.🤖 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_ontology_store.py` around lines 116 - 166, Add a regression test in tests/test_ontology_store.py that registers two types sharing the same property name to ensure owner-scoped property upserts: create a GraphSchema with an EntityType(label="Person", properties=[PropertyType(name="since", type="DATE", required=True)]) and a RelationType(label="WORKS_AT", patterns=[("Person","Company")]) that also has properties=[PropertyType(name="since", type="INTEGER", required=False)]; call await store.register(schema) and then inspect fake_graph.calls to assert there are two separate MERGE/SET operations for the property (look for "MERGE (o)-[:HAS_PROPERTY]->" and the params passed) and that the params/metadata for each include the correct owner (e.g., owner label or the relation/entity context) and the distinct type/required values to prove they were persisted independently.
🤖 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/api/main.py`:
- Around line 196-203: delete_all() currently clears only the data graph leaving
the persisted ontology in OntologyStore (_ontology_store) and a stale
_global_schema; update delete_all() to also clear the ontology store and reset
in-memory schema: call the OntologyStore.clear()/delete_all() API (or implement
one) on self._ontology_store, then set self._global_schema back to a clean
default (e.g., self.schema or an empty GraphSchema) and ensure
refresh_ontology()/load paths will repopulate from the cleared store; reference
OntologyStore, delete_all(), _ontology_store, _global_schema, and
refresh_ontology() when making the change.
In `@graphrag_sdk/src/graphrag_sdk/core/models.py`:
- Around line 314-315: The merge() implementation currently uses
"incoming.description or current.description" which treats empty strings as
falsey and prevents an explicit empty description from overwriting an existing
one; change those checks in merge() to use an explicit None check (e.g., use
incoming.description if incoming.description is not None else
current.description) so last-write-wins applies even when incoming.description
== "". Apply the same fix to the other occurrence that uses the same pattern
(the second use around the properties merge).
In `@graphrag_sdk/src/graphrag_sdk/retrieval/strategies/cypher_generation.py`:
- Around line 261-265: SCHEMA_PROMPT currently points to _SCHEMA_PROMPT_TEMPLATE
which requires schema_block and attribute_examples and thus breaks legacy
callers using SCHEMA_PROMPT.format(question=...); replace the simple alias with
a compatibility shim object (or small helper) named SCHEMA_PROMPT that
implements a format(**kwargs) method which: if kwargs contains only "question"
(legacy call), delegates to build_schema_prompt(question=kwargs["question"]) and
returns that built string, otherwise delegates to
_SCHEMA_PROMPT_TEMPLATE.format(**kwargs); ensure the shim also behaves like a
string for non-format uses (e.g., __str__ or __repr__ returning the underlying
template) so callers of SCHEMA_PROMPT continue to work.
In `@graphrag_sdk/src/graphrag_sdk/storage/ontology_store.py`:
- Around line 177-181: The MERGE for OntologyProperty is using only {name:
$name}, causing properties with the same name across different owners to
collide; update the upsert so the property node is keyed by both name and its
owner context (e.g. include owner or owner_label in the MERGE pattern that
creates p) instead of just name, and adjust any related uniqueness constraints
or queries that expect a single-keyed property (refer to the MERGE creating
p:OntologyProperty, the variable p, and owner_label in this block).
---
Nitpick comments:
In `@graphrag_sdk/tests/test_ontology_store.py`:
- Around line 116-166: Add a regression test in tests/test_ontology_store.py
that registers two types sharing the same property name to ensure owner-scoped
property upserts: create a GraphSchema with an EntityType(label="Person",
properties=[PropertyType(name="since", type="DATE", required=True)]) and a
RelationType(label="WORKS_AT", patterns=[("Person","Company")]) that also has
properties=[PropertyType(name="since", type="INTEGER", required=False)]; call
await store.register(schema) and then inspect fake_graph.calls to assert there
are two separate MERGE/SET operations for the property (look for "MERGE
(o)-[:HAS_PROPERTY]->" and the params passed) and that the params/metadata for
each include the correct owner (e.g., owner label or the relation/entity
context) and the distinct type/required values to prove they were persisted
independently.
🪄 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: 9aab5a64-1ab3-4985-8285-6595928e4a23
📒 Files selected for processing (14)
graphrag_sdk/src/graphrag_sdk/api/main.pygraphrag_sdk/src/graphrag_sdk/core/models.pygraphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/graph_extraction.pygraphrag_sdk/src/graphrag_sdk/ingestion/pipeline.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/cypher_generation.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.pygraphrag_sdk/src/graphrag_sdk/storage/__init__.pygraphrag_sdk/src/graphrag_sdk/storage/ontology_store.pygraphrag_sdk/tests/test_attribute_prompt.pygraphrag_sdk/tests/test_cypher_generation.pygraphrag_sdk/tests/test_graph_extraction.pygraphrag_sdk/tests/test_models.pygraphrag_sdk/tests/test_ontology_store.pygraphrag_sdk/tests/test_pipeline.py
The data graph already contains labels, relationship types, properties, and endpoint patterns — querying it is the source of truth, not a parallel persisted graph. Schema-as-config is supported via plain JSON files. - storage/ontology_store: rewritten as inference-only. OntologyStore.infer() calls db.labels() / db.relationshipTypes(), samples keys()+typeof() per label, derives RELATES sub-types from rel_type values, and pulls endpoint patterns. Structural labels/edges and reserved property keys are filtered out so they never leak into the LLM-facing schema. - core/models: GraphSchema.from_file / save_to_file convenience methods for the schema-as-config workflow (versioned JSON, hand-editable, shareable). - api/main: GraphRAG.get_ontology() returns inferred ∪ self.schema so user- declared descriptions and required flags survive (and properties declared but not yet extracted still appear in the retrieval prompt). New save_ontology(path) writes the global ontology to a JSON file. Ingest no longer writes to a separate ontology graph. The previous PR introduced a persistent <data_graph>__ontology FalkorDB graph; that was solving "declared property with zero instances" — a case that is essentially invisible to retrieval (empty results either way). The runtime cost was a second graph to back up, clean up, and reason about. Inference + an opt-in JSON file covers the same ground in less code. Tests: 830 passed, 23 skipped. test_ontology_store rewritten around mocked-driver introspection. test_models gains a JSON round-trip test. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
graphrag_sdk/src/graphrag_sdk/api/main.py (1)
596-601:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRefresh or invalidate
_global_schemaafter every graph mutation, not only ingest.This best-effort refresh only runs on
_ingest_single().update(),delete_document(), anddelete_all()also change the graph, but they leave_global_schemaand the retrieval strategy's cached schema untouched, so subsequent retrieval can use stale labels/properties until the caller manually invokesrefresh_ontology().🤖 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 596 - 601, The patch currently refreshes the global ontology only in _ingest_single(); update(), delete_document(), and delete_all() also mutate the graph but don't refresh or invalidate the cached schema, causing stale retrievals; modify those mutation paths (the update(), delete_document(), and delete_all() methods) to either call await self.refresh_ontology() after a successful mutation or explicitly clear/invalidate self._global_schema and notify the retrieval strategy to clear its cached schema (use the same mechanism the ingest path uses to update retrieval_strategy's cache), so that retrievals see new/removed properties immediately.graphrag_sdk/src/graphrag_sdk/core/models.py (1)
330-333:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winPreserve explicit empty descriptions in
merge().These branches still use
incoming.description or current.description, somerge()cannot intentionally clear a description with""even though the method documents last-write-wins behavior.Suggested fix
ent_by_label[e.label] = EntityType( label=cur.label, - description=e.description or cur.description, + description=e.description if e.description is not None else cur.description, properties=_merge_props(cur.properties, e.properties), ) @@ rel_by_label[r.label] = RelationType( label=cur.label, - description=r.description or cur.description, + description=r.description if r.description is not None else cur.description, patterns=merged_patterns, properties=_merge_props(cur.properties, r.properties), )Also applies to: 348-350
🤖 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/core/models.py` around lines 330 - 333, In merge(), the code uses falsy checks (incoming.description or current.description) which prevents intentionally clearing a description with an empty string; update the description selection to check for None explicitly (e.g., use incoming.description if incoming.description is not None else current.description) in the branches that construct EntityType (references: ent_by_label, EntityType construction, variables e/cur or incoming/current) and the analogous branch around lines 348-350 so last-write-wins semantics accept "" as a valid value.
🤖 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/storage/ontology_store.py`:
- Around line 111-119: The infer()/get_ontology() path only collects relation
types when "RELATES" exists, dropping direct user edge labels like
WORKS_AT/LOCATED_IN; update the logic around rel_types and relations (the
relations: list[RelationType] block and use of self._infer_relates_subtypes) to
also include any non-structural rel_types besides RELATES by converting each
such label into a RelationType (including property keys and endpoint patterns)
or by calling a new helper to infer their properties using the same sample_size;
keep excluding structural edges (PART_OF, NEXT_CHUNK, MENTIONED_IN) and preserve
the existing RELATES handling via self._infer_relates_subtypes(sample_size).
---
Duplicate comments:
In `@graphrag_sdk/src/graphrag_sdk/api/main.py`:
- Around line 596-601: The patch currently refreshes the global ontology only in
_ingest_single(); update(), delete_document(), and delete_all() also mutate the
graph but don't refresh or invalidate the cached schema, causing stale
retrievals; modify those mutation paths (the update(), delete_document(), and
delete_all() methods) to either call await self.refresh_ontology() after a
successful mutation or explicitly clear/invalidate self._global_schema and
notify the retrieval strategy to clear its cached schema (use the same mechanism
the ingest path uses to update retrieval_strategy's cache), so that retrievals
see new/removed properties immediately.
In `@graphrag_sdk/src/graphrag_sdk/core/models.py`:
- Around line 330-333: In merge(), the code uses falsy checks
(incoming.description or current.description) which prevents intentionally
clearing a description with an empty string; update the description selection to
check for None explicitly (e.g., use incoming.description if
incoming.description is not None else current.description) in the branches that
construct EntityType (references: ent_by_label, EntityType construction,
variables e/cur or incoming/current) and the analogous branch around lines
348-350 so last-write-wins semantics accept "" as a valid value.
🪄 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: ee8ce5db-b86d-454b-94a9-d1ef91a6f6f8
📒 Files selected for processing (5)
graphrag_sdk/src/graphrag_sdk/api/main.pygraphrag_sdk/src/graphrag_sdk/core/models.pygraphrag_sdk/src/graphrag_sdk/storage/ontology_store.pygraphrag_sdk/tests/test_models.pygraphrag_sdk/tests/test_ontology_store.py
🚧 Files skipped from review as they are similar to previous changes (1)
- graphrag_sdk/tests/test_models.py
- graph_extraction.py: regroup imports (datetime alongside stdlib); remove unused EntityType import; wrap two over-length attribute-fetch lines. - api/main.py: split a long get_ontology return expression. ruff check src/graphrag_sdk/ passes locally. Tests: 830 passed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This was referenced May 18, 2026
3 tasks
CI's "ruff format --check src/" flagged 5 files. Pure formatting; no behavioral change. Tests: 830 passed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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/storage/ontology_store.py`:
- Around line 140-146: The query that builds subtypes_result currently applies
LIMIT before DISTINCT, causing rare r.rel_type values to be dropped; update the
Cypher used in the get_ontology (or wherever subtypes_result is produced) so
DISTINCT r.rel_type is applied prior to sampling/limiting: select DISTINCT
r.rel_type AS rel_type (filter out NULLs), then apply the LIMIT $limit to that
deduplicated set before returning rel_type via self._conn.query.
🪄 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: 61f452a0-7e8d-4c3d-aa3e-7be0ec4c8ff3
📒 Files selected for processing (5)
graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/graph_extraction.pygraphrag_sdk/src/graphrag_sdk/ingestion/pipeline.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/cypher_generation.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.pygraphrag_sdk/src/graphrag_sdk/storage/ontology_store.py
✅ Files skipped from review due to trivial changes (1)
- graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py
🚧 Files skipped from review as they are similar to previous changes (3)
- graphrag_sdk/src/graphrag_sdk/ingestion/pipeline.py
- graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/graph_extraction.py
- graphrag_sdk/src/graphrag_sdk/retrieval/strategies/cypher_generation.py
There was a problem hiding this comment.
Pull request overview
Implements schema-driven attributes across extraction, ingestion validation, and schema-aware Cypher generation, plus an ontology mechanism to supply a global schema to retrieval.
Changes:
- Adds attribute coercion/extraction support driven by
GraphSchemaproperty declarations, and validates/drops records missing required attributes during ingestion. - Introduces
OntologyStoreand threads schema/ontology into retrieval so Cypher prompt/validation can be schema-aware. - Extends unit test coverage across models, extraction, ingestion pipeline, ontology inference, and Cypher generation.
Reviewed changes
Copilot reviewed 14 out of 14 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| graphrag_sdk/src/graphrag_sdk/retrieval/strategies/cypher_generation.py | Builds schema-aware prompt blocks/examples and validates labels against an optional GraphSchema. |
| graphrag_sdk/src/graphrag_sdk/storage/ontology_store.py | New ontology inference helper that introspects the live graph for labels/relationship subtypes/properties. |
| graphrag_sdk/src/graphrag_sdk/core/models.py | Adds schema property typing/validation, reserved-name checks, schema merge, and extracted attributes fields. |
| graphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/graph_extraction.py | Adds attribute prompt block + per-type coercion and carries attributes through aggregation/storage conversion. |
| graphrag_sdk/src/graphrag_sdk/ingestion/pipeline.py | Adds _validate_attributes pass after pruning and re-runs quality filtering. |
| graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py | Threads schema into Cypher retrieval execution when enabled. |
| graphrag_sdk/src/graphrag_sdk/api/main.py | Instantiates OntologyStore, adds get/refresh/save ontology APIs, and propagates schema to retrieval strategy. |
| graphrag_sdk/src/graphrag_sdk/storage/init.py | Exports OntologyStore. |
| graphrag_sdk/tests/test_attribute_prompt.py | New tests for attribute prompt rendering and coercion helpers. |
| graphrag_sdk/tests/test_cypher_generation.py | Adds tests for schema block rendering, prompt building, and schema-aware Cypher validation. |
| graphrag_sdk/tests/test_graph_extraction.py | Adds tests for attribute coercion, required-attribute dropping, aggregation, and property merge behavior. |
| graphrag_sdk/tests/test_models.py | Updates/extends schema model tests (type normalization, reserved keys, merge). |
| graphrag_sdk/tests/test_ontology_store.py | New tests for ontology inference logic with a mocked FalkorDB driver. |
| graphrag_sdk/tests/test_pipeline.py | Adds tests for ingestion attribute validation behavior. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
`required=True` was a footgun for an LLM-extraction pipeline: when the LLM couldn't find a value in the text (often because the text didn't state it), the entire entity was silently dropped. That destroyed real data. New behavior: every declared property name appears in `ExtractedEntity.attributes` with either the coerced value or `None`. The storage layer's existing `_clean_properties` strips `None` before writing so the graph sees "key missing" — the right null semantics for retrieval (`WHERE p.age > N` naturally excludes nodes without `age`). - core/models: remove `required` field from PropertyType. - ingestion/graph_extraction: `_coerce_attributes` now returns just a dict with `None` for missing/uncoercible values; never drops records. The drop-on-missing-required branches and aggregated warnings are gone. - ingestion/pipeline: `_validate_attributes` strips undeclared keys but never drops records. Cascade `_filter_quality` still runs but no longer has anything to cascade in the attribute-validation case. - storage/ontology_store: docstring no longer mentions required flags. - tests: collapsed the three "required missing" cases into a single "every declared property appears in result; uncoercible -> None" test per file. 829 passed, 23 skipped. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The data graph alone wasn't enough to anchor the schema: - types are lossy (FalkorDB `typeof()` collapses DATE → "string") - multiple processes/instances couldn't share a declarative view - nothing caught schema-typo contradictions across sessions - declarative metadata (descriptions, future flags) had no canonical home Restore the persistent ontology graph (``<data>__ontology``) as the single source of truth, with additive-only semantics: - OntologyStore.register() validates incoming schema against the persisted ontology. Re-typing an existing property raises OntologyContradictionError before any partial state is persisted; additions (new labels, properties, patterns) go through unchanged. - OntologyStore.load() / clear() round out the lifecycle. - GraphRAG._ensure_ontology_initialized() lazy first-touch loads the persisted ontology and registers self.schema. Called from ingest (so contradictions surface before expensive extraction) and from get_ontology / retrieval paths. - delete_all() now drops both data and ontology graphs and resets the initialised flag so the next ingest re-registers self.schema cleanly. - get_ontology() always reads from the ontology graph; refresh_ontology() is a thin re-load for cross-process freshness. save_ontology(path) remains as a JSON-file bridge. This reverts the architecture of 54e4e05 (drop ontology graph, infer from data). The "save a graph" simplicity wasn't worth losing validation, multi-process safety, type fidelity, and a queryable schema artifact. Tests: 825 passed, 23 skipped. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…y one Before, a user passing no schema got an empty ontology graph even though extraction was using DEFAULT_ENTITY_TYPES under the hood. The anchor said "we know nothing" while the LLM was producing Person/Organization/... nodes. Now ``_ensure_ontology_initialized()`` has three branches: - user passed a schema → register it (validate + persist) - ontology graph already populated (prior session) → use as-is - both empty → register DEFAULT_ENTITY_TYPES so the ontology graph faithfully reflects what the extractor will produce After this, ``get_ontology()``, the Cypher-generation prompt, and the extractor all read from the same source. Test fixture: ``mock_conn`` in test_facade.py now stubs the connection's ``_driver.select_graph()`` chain so OntologyStore can open the ontology graph handle against the mock. 825 passed, 23 skipped. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Integration tests on the Semantic resolver were failing because the
DEFAULT_ENTITY_TYPES seed registers labels without any properties, and
``_validate_attributes`` was stripping *every* non-reserved key from those
nodes (none were "declared"). That wiped the `embedding`, `country`, etc.
properties Semantic resolution uses to match duplicates — Bob disappeared.
Treat "label declared with zero properties" as schema-open: anchor the
label but don't constrain its attributes. Enforcement kicks in only when
a type actually declares at least one property. This restores Semantic
resolver correctness for the no-user-schema path and matches the natural
read of "I declared the label but didn't constrain it."
Also: clean up the Cypher prompt template so ``{attribute_examples}``
sits on its own line (was wedged against the closing code fence —
Copilot flagged it as confusing prompt structure).
Tests: 826 passed.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…te API
Simpler ingestion contract: extraction writes whatever it produces, the
data graph stores it as-is. The ontology graph stays the anchor; the data
graph carries history (immutable wrt schema lifecycle).
Removed:
- pipeline._validate_attributes (strict undeclared-key stripping). It was
defending against rare misbehavior at the cost of cross-file
coupling, surprising deletions when bare labels were declared, and a
bug that wiped Semantic-resolver embeddings.
- RESERVED_PROPERTY_NAMES from core/models. No longer needed: without
filtering, there's nothing to "preserve from" the strip pass. Users
who declare PropertyType(name="name") shadow the SDK value — their
problem, documented in the validator docstring.
- GraphSchema._validate_schema's reserved-name rejection. Demoted to a
pattern-label warning only.
Added:
- OntologyStore.delete_property / delete_entity_type / delete_relation_type
for the schema-lifecycle delete path.
- GraphRAG.delete_property(label, name, on_relation=False) public wrapper
that refreshes the global schema.
Schema lifecycle now matches the user's mental model:
- Add: pass schema=... to GraphRAG or new ingest. register() validates no
type contradictions and persists additions.
- Delete: rag.delete_property("Person", "age"). Forward-only — existing
values in the data graph stay.
- Extraction: tries to fill every declared attribute; missing/uncoercible
→ None (stripped at storage = key absent on the node).
- Data graph: never touched by schema lifecycle changes.
Tests: 822 passed, 23 skipped. Dropped the TestValidateAttributes class
and two reserved-name rejection tests; added TestDeleteLifecycle.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Naseem77
approved these changes
May 25, 2026
…tributes # Conflicts: # graphrag_sdk/src/graphrag_sdk/retrieval/strategies/cypher_generation.py # graphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.py
…sting labels The ingest path is for ingesting data; schema evolution is a different concern. Splitting them: - OntologyStore.register() is now strict on existing labels. Re-declaring with the same (or a subset of) persisted properties is accepted — that's a reference, not a modification. Adding new properties or new relation patterns to an existing label raises the new SchemaModificationNotAllowedError before any partial state persists. - delete_property / delete_entity_type / delete_relation_type all removed. There is no v1 way to drop or change a registered attribute. To start fresh, call await rag.delete_all() and re-ingest. Why: keeping the ontology graph and the data graph in lockstep matters more than letting users mutate the schema from the ingest path. A future schema-evolution API will handle adds/removes alongside the data updates required to stay aligned; until that ships, the ingest path is the only mutation surface and it's intentionally narrow. Errors: - OntologyContradictionError → property type changed (still rejected). - SchemaModificationNotAllowedError → adding new properties/patterns to an existing label (newly rejected). Tests: 869 passed. Dropped TestDeleteLifecycle; added TestStrictModification covering the four "rejected modification" cases plus the "subset re-declare is fine" case. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The "additive only" framing no longer fits — modifications to existing labels are also rejected now (SchemaModificationNotAllowedError). Spell out the three categories of register() input: accepted, contradiction, modification. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Two small UX fixes flagged in design review: 1. _ensure_ontology_initialized() now fires at the start of retrieve() and completion(), not just ingest()/get_ontology(). Previously, querying an existing graph without ingesting first left the retrieval strategy with the empty constructor-time self.schema and the Cypher prompt fell back to the historic hardcoded label set. Now the persisted ontology is loaded on first query. 2. The DEFAULT_ENTITY_TYPES seed (when both user schema and persisted ontology are empty) is now logged at INFO. Previously silent; users would later be confused that they couldn't add properties to Person/ Organization/etc. without delete_all(). The log line points at the fix. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Aligns the SDK's terminology with the legacy v0 model and with how users actually think about this concept. The "schema" naming was always a Python/SQL-isim; "ontology" is what knowledge-graph users say. Class renames in core/models: - GraphSchema → Ontology - EntityType → Entity - RelationType → Relation - PropertyType → Attribute Exception rename: - SchemaModificationNotAllowedError → OntologyModificationNotAllowedError API surface renames: - GraphRAG(schema=...) → GraphRAG(ontology=...) - rag.schema → rag.ontology - _global_schema → _global_ontology - IngestionPipeline(schema=...) → IngestionPipeline(ontology=...) - ExtractionStrategy.extract(chunks, schema, ctx) → (chunks, ontology, ctx) - OntologyStore.register(schema) → register(ontology) - MultiPathRetrieval(schema=...) → MultiPathRetrieval(ontology=...) - self._schema → self._ontology - build_schema_prompt → build_ontology_prompt - render_schema_block → render_ontology_block - SCHEMA_PROMPT / _SCHEMA_PROMPT_TEMPLATE → ONTOLOGY_PROMPT / _ONTOLOGY_PROMPT_TEMPLATE - _labels_from_schema → _labels_from_ontology - _render_attribute_schema_block → _render_attribute_block - _schema_has_attributes → _ontology_has_attributes - Template var schema_block → ontology_block - Test fixture sample_schema → sample_ontology This is intentionally a breaking change: callers using the old names will get import errors at startup. The mental model gets simpler in exchange: "the ontology is what you declare, the data graph is what you ingest into." Tests: 869 passed, 23 skipped. Lint clean. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Collaborator
Author
|
@coderabbitai review |
✅ Actions performedReview triggered.
|
Comment on lines
+293
to
+296
| self._global_ontology = await self._ontology_store.register(default_schema) | ||
| if hasattr(self._retrieval_strategy, "_schema"): | ||
| self._retrieval_strategy._schema = self._global_ontology | ||
| self._ontology_initialized = True |
Comment on lines
+309
to
+313
| loaded = await self._ontology_store.load() | ||
| self._global_ontology = loaded | ||
| if hasattr(self._retrieval_strategy, "_schema"): | ||
| self._retrieval_strategy._schema = self._global_ontology | ||
| return self._global_ontology |
Comment on lines
1128
to
1136
| pipeline = IngestionPipeline( | ||
| loader=loader or TextLoader(), # unused (text is provided below) | ||
| chunker=chunker or SentenceTokenCapChunking(), | ||
| extractor=extractor or self._default_extractor(), | ||
| resolver=resolver or ExactMatchResolution(), | ||
| graph_store=self._graph_store, | ||
| vector_store=self._vector_store, | ||
| schema=self.schema, | ||
| ontology=self.ontology, | ||
| ) |
Comment on lines
+808
to
+811
| # Merge ontology-declared attributes. Reserved-name collisions are | ||
| # rejected at ontology-validation time, so update() is safe. | ||
| if ent.attributes: | ||
| props.update(ent.attributes) |
Comment on lines
+10
to
+19
| - **Ingest path is constrained**. :py:meth:`register` admits: | ||
| - new entity / relation labels with their declared properties and patterns, | ||
| - re-declarations of existing labels with the *same* properties / patterns | ||
| (or a strict subset — treated as "use the persisted definition"). | ||
|
|
||
| It refuses: | ||
| - **type contradictions** on existing properties | ||
| (:py:class:`OntologyContradictionError`), and | ||
| - **modifications** to existing labels — adding properties or patterns | ||
| (:py:class:`OntologyModificationNotAllowedError`). |
There was a problem hiding this comment.
Actionable comments posted: 1
♻️ Duplicate comments (2)
graphrag_sdk/src/graphrag_sdk/api/main.py (1)
1128-1136:⚠️ Potential issue | 🟠 Major | ⚡ Quick winInitialize and use the persisted ontology before update extraction.
update()builds its pendingIngestionPipelinewithself.ontology, and this path never refreshes_global_ontologyfirst. On a fresh process updating an existing graph,self.ontologyis often empty/subset while the ontology store already contains the global union, so the update can extract/prune against the wrong ontology and silently miss declared attributes.Suggested fix
+ await self._ensure_ontology_initialized() + pipeline = IngestionPipeline( loader=loader or TextLoader(), # unused (text is provided below) chunker=chunker or SentenceTokenCapChunking(), extractor=extractor or self._default_extractor(), resolver=resolver or ExactMatchResolution(), graph_store=self._graph_store, vector_store=self._vector_store, - ontology=self.ontology, + ontology=self._global_ontology, )🤖 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 1128 - 1136, The update() path constructs an IngestionPipeline using self.ontology without first refreshing the persisted/global ontology, so on a fresh process self.ontology may be stale; before creating the IngestionPipeline (the block that instantiates IngestionPipeline with loader/chunker/extractor/resolver/graph_store/vector_store/ontology), load or merge the persisted global ontology into self._global_ontology and assign/merge it into self.ontology (e.g., call the existing method that loads the stored ontology or fetch it from the ontology store and set self.ontology = merged_global) so the pipeline extracts/prunes against the up-to-date unioned ontology.graphrag_sdk/src/graphrag_sdk/core/models.py (1)
228-258:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReject SDK-reserved attribute names at ontology-parse time.
This validator documents the collision hazard, but it never enforces it. Since extracted
attributesare merged straight into node/edge properties and then written withSET +=, declaringname,description,source_chunk_ids,rel_type, etc. can overwrite SDK-managed fields and corrupt identity/provenance/retrieval behavior.Suggested fix
+_RESERVED_PROPERTY_NAMES: frozenset[str] = frozenset( + { + "id", + "label", + "name", + "description", + "source_chunk_ids", + "spans", + "rel_type", + "fact", + "src_name", + "tgt_name", + } +) + class Ontology(DataModel): @@ `@model_validator`(mode="after") def _warn_on_undeclared_pattern_labels(self) -> Ontology: + for entity in self.entities: + for prop in entity.properties: + if prop.name in _RESERVED_PROPERTY_NAMES: + raise ValueError( + f"Entity '{entity.label}' property '{prop.name}' is reserved" + ) + for relation in self.relations: + for prop in relation.properties: + if prop.name in _RESERVED_PROPERTY_NAMES: + raise ValueError( + f"Relation '{relation.label}' property '{prop.name}' is reserved" + ) if not self.entities: return self🤖 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/core/models.py` around lines 228 - 258, The validator _warn_on_undeclared_pattern_labels currently only warns about pattern label typos but must also reject any declared Attribute whose name collides with SDK-reserved node/edge keys; in the Ontology model (method _warn_on_undeclared_pattern_labels) add a check against a set of reserved names (e.g. "name", "description", "source_chunk_ids", "spans", "rel_type", "fact", "src_name", "tgt_name", "id", "label") by iterating self.attributes (or wherever Attributes are declared) and raise a ValidationError (or ValueError) if any attribute.name is in that reserved set, including a clear message referencing the offending attribute and that it shadows SDK-managed fields.
🤖 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/api/main.py`:
- Around line 357-360: The delete_all() implementation currently swallows
exceptions from self._ontology_store.clear(), causing delete_all to report
success even when ontology cleanup failed; change it so ontology clear failures
propagate (or re-raise) instead of being caught and only logged. Locate the
delete_all method and the try/except around self._ontology_store.clear(), remove
or adjust the except block so you either let the original exception bubble up or
re-raise after logging, ensuring delete_all only returns success when
self._ontology_store.clear() completes without error (and similarly ensure any
data-graph drop functions used there follow the same success/failure semantics).
---
Duplicate comments:
In `@graphrag_sdk/src/graphrag_sdk/api/main.py`:
- Around line 1128-1136: The update() path constructs an IngestionPipeline using
self.ontology without first refreshing the persisted/global ontology, so on a
fresh process self.ontology may be stale; before creating the IngestionPipeline
(the block that instantiates IngestionPipeline with
loader/chunker/extractor/resolver/graph_store/vector_store/ontology), load or
merge the persisted global ontology into self._global_ontology and assign/merge
it into self.ontology (e.g., call the existing method that loads the stored
ontology or fetch it from the ontology store and set self.ontology =
merged_global) so the pipeline extracts/prunes against the up-to-date unioned
ontology.
In `@graphrag_sdk/src/graphrag_sdk/core/models.py`:
- Around line 228-258: The validator _warn_on_undeclared_pattern_labels
currently only warns about pattern label typos but must also reject any declared
Attribute whose name collides with SDK-reserved node/edge keys; in the Ontology
model (method _warn_on_undeclared_pattern_labels) add a check against a set of
reserved names (e.g. "name", "description", "source_chunk_ids", "spans",
"rel_type", "fact", "src_name", "tgt_name", "id", "label") by iterating
self.attributes (or wherever Attributes are declared) and raise a
ValidationError (or ValueError) if any attribute.name is in that reserved set,
including a clear message referencing the offending attribute and that it
shadows SDK-managed fields.
🪄 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: 5d6e80b2-f259-4af0-96fd-c76c56adfe6c
📒 Files selected for processing (21)
graphrag_sdk/src/graphrag_sdk/__init__.pygraphrag_sdk/src/graphrag_sdk/api/main.pygraphrag_sdk/src/graphrag_sdk/core/__init__.pygraphrag_sdk/src/graphrag_sdk/core/models.pygraphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/base.pygraphrag_sdk/src/graphrag_sdk/ingestion/extraction_strategies/graph_extraction.pygraphrag_sdk/src/graphrag_sdk/ingestion/pipeline.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/cypher_generation.pygraphrag_sdk/src/graphrag_sdk/retrieval/strategies/multi_path.pygraphrag_sdk/src/graphrag_sdk/storage/ontology_store.pygraphrag_sdk/tests/conftest.pygraphrag_sdk/tests/test_attribute_prompt.pygraphrag_sdk/tests/test_cypher_generation.pygraphrag_sdk/tests/test_facade.pygraphrag_sdk/tests/test_graph_extraction.pygraphrag_sdk/tests/test_ingestion_e2e.pygraphrag_sdk/tests/test_integration.pygraphrag_sdk/tests/test_models.pygraphrag_sdk/tests/test_ontology_store.pygraphrag_sdk/tests/test_pipeline.pygraphrag_sdk/tests/test_pipeline_10step.py
Comment on lines
+357
to
+360
| try: | ||
| await self._ontology_store.clear() | ||
| except Exception as exc: | ||
| logger.warning("Ontology graph clear failed during delete_all (continuing): %s", exc) |
There was a problem hiding this comment.
Don’t report delete_all() success when ontology cleanup failed.
If OntologyStore.clear() raises here, the method still returns successfully after dropping only the data graph. That leaves stale ontology in <graph>__ontology, and the next ingest/retrieve can bootstrap from it even though delete_all() promised a full reset.
Suggested fix
- try:
- await self._ontology_store.clear()
- except Exception as exc:
- logger.warning("Ontology graph clear failed during delete_all (continuing): %s", exc)
+ await self._ontology_store.clear()🤖 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 357 - 360, The
delete_all() implementation currently swallows exceptions from
self._ontology_store.clear(), causing delete_all to report success even when
ontology cleanup failed; change it so ontology clear failures propagate (or
re-raise) instead of being caught and only logged. Locate the delete_all method
and the try/except around self._ontology_store.clear(), remove or adjust the
except block so you either let the original exception bubble up or re-raise
after logging, ensuring delete_all only returns success when
self._ontology_store.clear() completes without error (and similarly ensure any
data-graph drop functions used there follow the same success/failure semantics).
Three fixes flagged in Copilot's latest review: 1. multi_path.py: tighten the `ontology` parameter type from `Any | None` to `Ontology | None`. The annotation was loose; the import was avoided to dodge an imagined circular import that doesn't actually exist. 2. graph_extraction.py:808: fix a misleading code comment. It claimed reserved-name collisions are rejected at validation time — that hasn't been true since we deleted `RESERVED_PROPERTY_NAMES` in 12b3abb. The new comment honestly explains the trade-off (smaller API surface; declaring Attribute(name="name") shadows the system value). 3. test_pipeline.py: trim trailing blank lines left over from earlier deletion of the TestValidateAttributes class. The other four Copilot comments describe intentional design choices — replied per-thread on the PR. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Existing user code that imported the old names (GraphSchema, EntityType,
RelationType, PropertyType, SchemaModificationNotAllowedError), or that
passed `schema=` / read `rag.schema`, keeps working. Each access emits a
DeprecationWarning pointing at the new name. To be removed in a future
release.
Aliases wired:
- Class names via module __getattr__ (PEP 562):
- graphrag_sdk.GraphSchema → Ontology
- graphrag_sdk.EntityType → Entity
- graphrag_sdk.RelationType → Relation
- graphrag_sdk.PropertyType → Attribute
Also reachable via `graphrag_sdk.core.models.*` and
`graphrag_sdk.core.*`.
- Exception:
- storage.ontology_store.SchemaModificationNotAllowedError →
OntologyModificationNotAllowedError
- graphrag_sdk.SchemaModificationNotAllowedError → same
- GraphRAG constructor:
- `GraphRAG(..., schema=X)` still works; forwards to `ontology=X`.
- Passing both raises TypeError to prevent ambiguity.
- GraphRAG attribute:
- `rag.schema` is a property that warns and returns `rag.ontology`.
- `rag.schema = X` setter routes through to `rag.ontology = X`.
Also added top-level `Attribute` to the package re-exports (was missing).
Tests: 889 passed (was 869, +20 new in test_deprecations.py covering each
alias path).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
| ) | ||
| class TestModuleLevelClassAliases: | ||
| def test_top_level_import_warns(self, old, new): | ||
| import graphrag_sdk |
Collaborator
Author
There was a problem hiding this comment.
The mixed-import style is intentional: import graphrag_sdk inside each test scope lets us call getattr(graphrag_sdk, old) with the legacy name parameter, while from graphrag_sdk.X import Y at the top keeps the rest of the file readable. Switching to one style only would require either repeated graphrag_sdk.GraphRAG references throughout or losing the parametrised getattr pattern. Leaving as-is — it's not a runtime issue and ruff passes.
| assert new in str(deps[0].message) | ||
|
|
||
| def test_top_level_resolves_to_new_class(self, old, new): | ||
| import graphrag_sdk |
Collaborator
Author
There was a problem hiding this comment.
The mixed-import style is intentional: import graphrag_sdk inside each test scope lets us call getattr(graphrag_sdk, old) with the legacy name parameter, while from graphrag_sdk.X import Y at the top keeps the rest of the file readable. Switching to one style only would require either repeated graphrag_sdk.GraphRAG references throughout or losing the parametrised getattr pattern. Leaving as-is — it's not a runtime issue and ruff passes.
|
|
||
| def test_instantiating_via_legacy_alias_returns_new_class(): | ||
| """Old code: GraphSchema(entities=[EntityType(...)]) keeps working.""" | ||
| import graphrag_sdk |
Collaborator
Author
There was a problem hiding this comment.
The mixed-import style is intentional: import graphrag_sdk inside each test scope lets us call getattr(graphrag_sdk, old) with the legacy name parameter, while from graphrag_sdk.X import Y at the top keeps the rest of the file readable. Switching to one style only would require either repeated graphrag_sdk.GraphRAG references throughout or losing the parametrised getattr pattern. Leaving as-is — it's not a runtime issue and ruff passes.
Existing user code that constructs IngestionPipeline or MultiPathRetrieval directly (advanced/custom flows) keeps working: - IngestionPipeline(..., schema=X) → forwards to ontology=X, warns. - MultiPathRetrieval(..., schema=X) → forwards to ontology=X, warns. Both raise TypeError if a caller passes both `ontology=` and `schema=`. Tests: 895 passed (+6 new in test_deprecations.py). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…helpers Covers the last surface-API names that still required a manual rename: - SCHEMA_PROMPT → ONTOLOGY_PROMPT (module __getattr__) - build_schema_prompt → build_ontology_prompt (module __getattr__) - render_schema_block → render_ontology_block (module __getattr__) - validate_cypher(schema=...) → ontology= (kwarg + helper) - generate_cypher(schema=...) → ontology= (kwarg + helper) - execute_cypher_retrieval(schema=...) → ontology= (kwarg + helper) Each emits a DeprecationWarning pointing at the new name. Functions raise TypeError when callers pass both `ontology=` and `schema=`. Module __getattr__ handles the constant + the two helper functions so users importing the legacy names get a clear warning instead of ImportError. Only the private `OntologyStore.register(schema=)` path is still un-aliased — it's behind a leading-underscore attribute (private API). Tests: 900 passed (+5 new in test_deprecations.py covering the Cypher helpers + the kwarg alias on validate_cypher). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Closes the last gap in the schema → ontology deprecation surface. Although ``rag._ontology_store`` is conventionally private, code that reaches into it still keeps working: ``register(schema=ontology)`` is accepted, forwards to ``register(ontology=ontology)``, and emits a DeprecationWarning. Both kwargs raises TypeError. Missing both raises TypeError too. Tests: 903 passed (+3 in test_deprecations.py for the store alias). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Three real regressions from the rename + 3 doc/lint fixes: 1. api/main.py _ensure_ontology_initialized + get_ontology: the retrieval-strategy propagation was still hasattr/setattr-ing `_schema`, but MultiPathRetrieval now stores `_ontology` after the rename. The hasattr check returned False → silent no-op → retrieval kept seeing the constructor-time ontology forever. Both spots now use `_ontology`. 2. api/main.py _update_single: the update path built its IngestionPipeline with `ontology=self.ontology` (local) and never called `_ensure_ontology_initialized()`. Mirrors the ingest path now: ensure-init first, then pass `self._global_ontology` to the pipeline. 3. test_deprecations.py: dropped two unused locals (`obj = getattr(...)`, the getattr call is kept for its side-effect), the unused `FalkorDBConnection` and `Relation` imports. 4. extraction_strategies/base.py: "a ontology" → "an ontology" in the class docstring. 5. ontology_store.py module docstring: clarified that the ontology graph "accumulates new labels across ingest passes" but existing labels are frozen by register(). The previous wording read as unconstrained forward-only union, which contradicted the actual behavior. Tests: 903 passed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
End-to-end example exercising the new ontology API against a real FalkorDB. Eight scenarios: 1. Ingest with a declared Ontology + typed Attributes. 2. Inspect the persisted ontology via rag.get_ontology(). 3. Save / load the ontology as a JSON file (schema-as-config). 4. Add a new entity type on a subsequent ingest — additive, OK. 5. Modify an existing entity by adding a new attribute → OntologyModificationNotAllowedError (with the recovery instruction in the error message). 6. Re-type an existing attribute → OntologyContradictionError. 7. Subset re-declaration of an existing entity — treated as "use the persisted definition," accepted. 8. Legacy GraphSchema / EntityType / schema= kwarg still work and emit DeprecationWarnings. Verified locally: all 8 scenarios pass against falkordb:latest. Even with a misconfigured OpenAI key the validation paths fire correctly (they run before the LLM call), demonstrating the contradiction / modification errors are early and informative. Side fixes: - graphrag_sdk/__init__: export OntologyContradictionError, OntologyModificationNotAllowedError, and OntologyStore at the top level so user code can import them without reaching into storage submodules. - ontology_store.py: grammar — "a ontology-evolution" → "an ontology-evolution" in three places. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The old representation modelled the schema as awkward nodes:
(:OntologyEntityType {label: "Person"})-[:HAS_PROPERTY]->
(:OntologyProperty {name: "age", type: "INTEGER"})
(:OntologyRelationType {label: "WORKS_AT",
patterns: ["Person|Company"]})
-[:HAS_PROPERTY]->(:OntologyProperty {name: "since", ...})
That hid the graph nature: relation types were *nodes* with patterns
encoded as `"src|tgt"` strings, properties were nodes with thin edges,
and entity types and relation types weren't connected in the graph at
all.
New shape (Option B) — the ontology graph reads as a miniature data graph:
(:Person:__Ontology {description: "A human",
attributes: '{"birth_year":{"type":"INTEGER"}}'})
-[:WORKS_AT {description: "Employment",
attributes: '{"since":{"type":"DATE"}}'}]->
(:Company:__Ontology {description: "An organization"})
- Each entity type is a node carrying its user label + ``:__Ontology``.
- Declared attributes are a JSON-encoded map property on the node.
- Relation types are real edges (one per declared pattern), with their
own attributes as a JSON-encoded property on the edge.
- Open-mode relations (no patterns) materialise as self-loops on a
``:__OpenRelation:__Ontology`` placeholder so every relation is still
an edge.
Visualised in the FalkorDB browser this looks like the schema it
describes — Person → WORKS_AT → Company. Querying it is natural Cypher:
MATCH (e:__Ontology) RETURN e.label, e.attributes
MATCH (s:__Ontology)-[r]->(t:__Ontology)
WHERE NOT s:__OpenRelation
RETURN s.label, type(r), t.label, r.attributes
Internal changes:
- _encode_attributes / _decode_attributes (JSON map ↔ list[Attribute]),
replacing _encode_patterns / _decode_patterns / _props_from_rows.
- load() reads entity nodes, patterned edges, and open-relation loops
into a single Ontology object.
- _upsert_entity_type MERGEs on dual labels; _upsert_relation_type
MERGEs one edge per pattern (or a self-loop on the placeholder).
- _upsert_property dropped — attributes inline on the owner.
- All cypher uses sanitize_cypher_label for safety.
- Validation logic (contradiction / modification) untouched — it
operates on the loaded Ontology, not the on-graph shape.
Verified live against falkordb:latest: round-trip register → load
yields the exact ontology that was registered. 906 tests pass.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…erty)
Restoring the legacy-style ontology graph shape per review feedback.
Previously (Option B) entity types were dual-labeled nodes (e.g.
``:Person:__Ontology``) with declared attributes JSON-encoded into a
property on the node, and relations were real edges between those nodes
with attributes JSON-encoded into the edge. Manager-side feedback: the
ontology graph should be built from entity, relation, and property
*nodes* connected like a schema diagram — and ``label`` / ``description``
should be properties of the entity nodes, not Cypher labels.
New on-graph shape::
(:Entity {label, description})
(:Relation {label, description})
(:Property {label, type, description})
(:Entity)-[:HAS_PROPERTY]->(:Property)
(:Relation)-[:SOURCE]->(:Entity)
(:Relation)-[:TARGET]->(:Entity)
(:Relation)-[:HAS_PROPERTY]->(:Property)
A relation with N declared patterns materialises as N ``:Relation``
nodes (one per ``(src, tgt)`` triple), each carrying its own
``SOURCE``/``TARGET``/``HAS_PROPERTY`` edges. An open-mode relation
(no patterns) materialises as one ``:Relation`` node with no
``SOURCE``/``TARGET`` edges. ``:Property`` nodes are scoped per owner
via MERGE-on-pattern, so ``Person.age`` and ``Mountain.age`` are
distinct nodes — no accidental sharing across labels.
Verified live against falkordb:latest with a full ingest:
MATCH (s:Entity)<-[:SOURCE]-(r:Relation)-[:TARGET]->(t:Entity)
RETURN s.label, r.label, t.label
Person WORKS_AT Organization
Person COLLABORATED_WITH Person
Person WON Award
Organization LOCATED_IN Place
MATCH (e:Entity)-[:HAS_PROPERTY]->(p:Property)
RETURN e.label, p.label, p.type
Person birth_year INTEGER
Person birth_place STRING
Award year INTEGER
In the FalkorDB browser the graph visualises as the schema it describes.
Internal changes:
- Replaced ``_encode_attributes`` / ``_decode_attributes`` (the JSON
round-tripping helpers that went with Option B) with five MATCH
queries in ``load()`` — entities, entity-properties, patterned
relations, open relations, relation-properties — and per-owner
Attribute dedup.
- ``_upsert_entity_type`` now MERGEs ``(:Entity {label})`` and walks
declared attributes through ``_upsert_entity_property`` (one
``HAS_PROPERTY`` MERGE per attribute).
- ``_upsert_relation_type`` MERGEs one ``:Relation`` node per pattern,
via ``MERGE (s)<-[:SOURCE]-(r:Relation {label})-[:TARGET]->(t)``;
declared attributes attach to *each* pattern node via
``_upsert_relation_property``. Open-mode relations route through
``_upsert_open_relation_property`` (MERGE on a Relation with no
SOURCE edge).
- Dropped the ``__Ontology`` / ``__OpenRelation`` marker labels and
the ``sanitize_cypher_label`` calls they required.
- Validation (contradiction / strict-modification) is unchanged — it
operates on the loaded ``Ontology`` object, not the on-graph shape.
Tests rewritten to assert against the new query shape. 903 tests pass.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds an awaitable ``GraphRAG.set_ontology(new_ontology)`` that swaps the working ontology and re-runs ``_ensure_ontology_initialized`` cleanly, so callers can introduce a new label between ingest passes without poking ``rag._ontology_initialized = False`` (a private seam the lifecycle example used to demo). Subject to the same ``OntologyStore.register()`` rules: new labels are always accepted; existing labels still have to obey the contradiction / strict-modification guards. The example (08_ontology_lifecycle.py) is rewritten to use the new method end-to-end, no more private mutations. Also clarifies the docstring on ``_ensure_ontology_initialized``: the one-shot init guard is intentional — re-assigning ``self.ontology`` after construction doesn't re-fire registration, ``set_ontology`` is the supported way to evolve. Verified live against falkordb:latest: all 8 lifecycle scenarios pass. 903 unit tests pass. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.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
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
Wires declared
Attributeproperties onEntity/Relationtypes from the localOntologythrough extraction, storage, and retrieval. The ontology persists in a dedicated<data_graph>__ontologyFalkorDB graph and accumulates new labels across ingest passes; existing labels are frozen byOntologyStore.register()to keep the ontology aligned with the data graph (extending an existing label is an ontology-evolution operation reserved for a future API).Forward-only over time: nodes ingested before a property was declared remain without that property; FalkorDB's
MERGE … SET n += propswrites new keys naturally andWHERE p.attr > Nexcludes nodes that never received the key — no backfill machinery.What changed
core/models—Attribute(renamed fromPropertyType) with a type-normalising validator that rejects unknown types; newRelation.properties;ExtractedEntity/Relation.attributesdict; reserved-name collisions rejected at config time.GraphSchemawas renamed toOntology(with deprecation aliases for the legacy names).storage/ontology_store(new) — persistent ontology in<data_graph>__ontologygraph, materialised as a three-node schema (:Entity,:Relation,:Propertyconnected byHAS_PROPERTY/SOURCE/TARGET).register()is constrained: new labels are accepted; existing labels must be a strict subset of what's persisted, with type contradictions surfaced asOntologyContradictionErrorand attempted modifications asOntologyModificationNotAllowedError.ingestion/graph_extraction—VERIFY_EXTRACT_RELS_PROMPTgains a conditional## Attribute extractionblock (empty for property-less ontologies → zero token drift). Per-type coercion at the extractor boundary (INTEGER,FLOAT,BOOLEAN,DATE,LIST,STRING). Aggregators carry attributes through dedup with last-write-wins. Conversion merges attributes into graph props.ingestion/pipeline— new_validate_attributespass after_prune, then re-runs_filter_qualityto cascade-clean dangling relationships. SkipsUnknown-labelled nodes.retrieval/cypher_generation— replaces hardcodedSCHEMA_PROMPTwithbuild_ontology_prompt(ontology, question)+render_ontology_block(ontology). Synthesises one numeric-attribute example per declaredINTEGER/FLOATproperty so the LLM learns the column shape.validate_cypheraccepts anontologyfor dynamic label validation; falls back to the historic hardcoded set when omitted.retrieval/multi_path— ontology threaded intoexecute_cypher_retrieval.api/main—GraphRAGconstructs anOntologyStore, registers the local ontology on first ingest, and propagates the global ontology to the retrieval strategy. Publicget_ontology()/refresh_ontology()for inspection. Newset_ontology(new)swaps the working ontology and re-fires registration without poking private state.Algorithm
Ingest (per run, with
local_ontology):set_ontology):ontology_store.register(local_ontology)validates against the persisted ontology and writes any new labels.attributesvia the augmented prompt; deterministic strategies (e.g. GLiNER-only) leave them empty.MERGE … SET n += properties— new keys land, existing keys overwrite, untouched keys stay.Retrieval (per question):
<data_graph>__ontology.WHERE p.age > 30naturally excludes pre-agenodes.Backwards compatibility
Ontologies without declared
propertiesproduce identical output to the prior pipeline:"", so token-for-token the prompt is unchanged._validate_attributesshort-circuits on empty ontologies.render_ontology_block(Ontology())falls back to the historic hardcoded label set.The
GraphSchema/EntityType/RelationType/PropertyTypenames and theschema=kwarg still work via PEP-562 module aliases and property aliases; each emits a singleDeprecationWarningon first use.Test plan
RUN_INTEGRATION=1).test_ontology_store,test_attribute_prompt,test_models,test_graph_extraction,test_pipeline,test_cypher_generation,test_deprecations.examples/08_ontology_lifecycle.py) verified against live FalkorDB: ingest with declared ontology, inspect, save/load JSON, extend viaset_ontology, hit the modification / contradiction guards, accept subset re-declaration, verify legacy aliases still work.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests