--- name: cognee-custom-graph-models description: Use when defining the shape of cognee's knowledge graph with graph_model= — writing DataPoint node classes, choosing identity and index fields so nodes merge and are searchable, declaring typed Edge fields and FromIdentity references, building a model from a JSON schema, or debugging duplicated nodes, missing edges, or InvalidReferenceTypeError. --- # Custom graph models By default cognee extracts a generic `KnowledgeGraph` of entities and relationships. Pass your own model with `graph_model=` and the LLM fills your node and edge types instead. ```python from typing import Annotated, Literal import cognee from cognee.low_level import DataPoint, Edge, FromIdentity class Role(DataPoint): name: str metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]} class Person(DataPoint): name: str is_a: Annotated[Role, FromIdentity()] | None = None # reference by name reports_to: list[Edge["Person", "Person"]] = [] # edge owned by Person metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]} class PeopleGraph(DataPoint): # the root the LLM fills people: list[Person] friends_with: list[Edge[Person, Person]] = [] family: list[Edge[Person, Person, Literal["married_to", "sibling_of"]]] = [] metadata: dict = {"index_fields": [], "transparent": True} await cognee.remember(text, graph_model=PeopleGraph, custom_prompt="Extract every person...") ``` Full example: `examples/guides/custom_graph_model.py`. ## Use it ### Nodes: DataPoint classes Every node type subclasses `DataPoint` (`from cognee.low_level import DataPoint`). Its fields become: - **Node properties:** scalars, strings, dicts, and anything that is not a DataPoint. - **Edges named after the field:** a field holding a DataPoint or a list of DataPoints. `members: list[Person]` becomes `members` edges. `dict[str, Person]`, sets and plain tuples are stored as properties, not edges. ### Identity and search: `metadata` | Key | What it does | |---|---| | `identity_fields` | The node id is derived from these field values (normalized: lowercased, spaces to `_`, apostrophes removed). The same entity from two chunks or two runs becomes **one node**. | | `index_fields` | Each field gets a vector collection named `_`, so recall can find the node. | | `transparent` | The node is not stored; its children take its place. Use it for a root container like `PeopleGraph`. | **Without `identity_fields` every node gets a random id, so the same person is duplicated in every chunk and every run.** Set it on every node type that represents a real-world entity. **Write `metadata` explicitly**, as in the examples above. There is also an annotation shortcut (`from cognee.infrastructure.engine import Dedup, Embeddable`; `name: Annotated[str, Embeddable(), Dedup()]`), but today only half of it works: - `Dedup()` works: ids are derived from the marked fields. - `Embeddable()` does not index. The markers update the class-level default, but each instance still carries `{"index_fields": []}`, and indexing reads the instance, so no vector collection is created and recall cannot find the node. Markers are also ignored entirely when the class declares `metadata` itself. ### Typed edges: `list[Edge[Source, Target, Name]]` The LLM answers edges as flat rows of identity strings (`source`, `target`), and cognee resolves them to the extracted nodes. The third parameter controls the relationship name: | Declaration | Relationship name | |---|---| | `list[Edge[Person, Person]]` | The field name (`friends_with`) | | `list[Edge[Person, Person, Literal["a", "b"]]]` | The LLM picks one value | | `list[Edge[Person, Person, str]]` | Free-form from the LLM, normalized | - **Where to declare:** on the root model for relationships with no obvious owner, or on the owning node. On the owner, endpoints of the owner's own type must be strings (`Edge["Person", "Person"]`), because the class is not defined yet inside its own body. - **Always a list:** `Edge[...]` or `Edge[...] | None` on its own raises. - **Both endpoint types need exactly one `identity_fields` entry.** ### References: `Annotated[Target, FromIdentity()]` Instead of a nested object, the LLM answers the identity string of a node (`is_a: "engineer"`), and cognee links to that node. Supported spellings: `Target`, `Target | None`, `list[Target]`, `list[Target] | None`. Anything else raises `InvalidReferenceTypeError`. The target needs exactly one identity field, and its other required fields need defaults. ### Edge values you build by hand `Edge(source=..., target=..., relationship_type=..., weight=..., properties={...})`. An omitted `source` falls back to the node declaring the field; on a parametrized field that node must be the declared `Source` type, or it raises. On a root container, always pass `source=`. The tuple form `(Edge(weight=0.8), target_node)` attaches edge properties to a plain DataPoint field (`examples/guides/custom_data_models.py`). ### From JSON instead of Python - `cognee.low_level.graph_model_from_spec(spec)`: a small entity/relation spec (names, fields, `one`/`many` relations) compiled to DataPoint classes, with identity and index on `name` by default. Example: `examples/guides/graph_model_from_json.py`. - `cognee.low_level.graph_schema_to_graph_model(json_schema)`: a JSON Schema (needs a top-level `title`; only internal `#` refs). - HTTP: `POST /api/v1/remember` takes a `graph_model` form field (JSON schema string); `POST /api/v1/cognify` takes a `graphModel` JSON object. `POST /api/v1/llm/infer-schema` proposes a schema from sample text. Neither JSON path can express typed `Edge` fields or `FromIdentity`; use Python classes for those. ## Pitfalls - **Duplicated nodes** → missing `identity_fields`. - **Node never shows up in recall** → no `index_fields`, or the recalling process never imported the model class (graph completion searches the collections of DataPoint classes loaded in that process). - **Edge rows silently missing** → an unresolved row is dropped with a warning. Endpoints resolve by **exact type** against nodes extracted from the same chunk: a subclass instance does not match where its base is declared, and a node from another chunk or an earlier run is not a candidate. - **`InvalidReferenceTypeError`, "declares Edge in a shape the LLM extraction cannot fill"** → an unsupported `FromIdentity` or `Edge` spelling. These are raised when the model is converted during extraction, not at class definition, so they appear mid-pipeline. - **String endpoints resolve only to the owning model itself or a module-level class.** A string naming another class defined inside a function raises `InvalidReferenceTypeError`, so define models at module level. - **Field names that collide with DataPoint's own fields** (`id`, `type`, `version`, `metadata`, `created_at`, `belongs_to_set`, …) are stripped from what the LLM sees. Rename them. - **A subclass that overrides `metadata`** replaces the parent's entirely. Dropping `identity_fields` only logs a warning. A subclass also hashes ids under its own class name. - **Write a `custom_prompt`.** Without one the generic knowledge-graph prompt is used; your schema reaches the LLM only as structured output. - **Not with GLiNER.** `extractor="gliner_demo"` (alias `"gliner"`) raises with a custom `graph_model`, and so does the default `GRAPH_EXTRACTOR=auto` when no LLM key is configured (it resolves to `gliner_demo`). - **Remote mode drops it.** After `cognee.serve(url)`, `remember()` and `cognify()` do not forward `graph_model`; the server builds a generic graph. - A custom model skips the generic path's ontology resolution, per-graph node dedup, and `functional_relationships`. Summaries still run. ## How it works `extract_content_graph` converts a DataPoint model into a plain Pydantic model for the LLM: infrastructure fields and `metadata` are stripped, typed edge fields become row lists (`FriendsWithEdge` with `source`/`target` strings), and `FromIdentity` fields become strings. The answer is converted back into DataPoint instances with ids from `identity_fields`, edge rows are resolved against the nodes in that answer, and the result is attached to the chunk (`chunk.contains`) and stored by `add_data_points`. Ownership is recorded per document, so `forget(data_id=...)` removes a custom-model document's nodes while shared nodes survive. - LLM boundary, both directions: `cognee/shared/llm_graph_model.py` - DataPoint, metadata, ids: `cognee/infrastructure/engine/models/DataPoint.py` - Markers: `cognee/infrastructure/engine/models/FieldAnnotations.py` - Edge: `cognee/infrastructure/engine/models/Edge.py` - Property vs edge decision: `cognee/modules/graph/utils/field_edges.py` - Custom-model branch of extraction: `cognee/tasks/graph/extract_graph_from_data.py` - JSON schema / spec paths: `cognee/shared/graph_model_utils.py`, `cognee/modules/graph_models/` - Vector collections: `cognee/tasks/storage/index_data_points.py` ## Extending it - Tests for the LLM round trip: `cognee/tests/unit/modules/graph/test_content_graph_to_data_point.py`. Edge typing: `cognee/tests/unit/interfaces/graph/test_typed_edge_model.py`, `test_typed_edges_graph.py`. Identity: `cognee/tests/unit/infrastructure/engine/test_identity_fields.py`. Property vs edge: `cognee/tests/unit/modules/graph/test_field_edges.py`. - Deletion of custom-model nodes: `cognee/tests/test_delete_custom_graph.py`. - A new `Edge` or `FromIdentity` spelling must be handled in both directions in `llm_graph_model.py`, and rejected with `InvalidReferenceTypeError` when unsupported, never silently accepted.