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.
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.
| 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 <ClassName>_<field>, 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.
anthropics/skills180k
