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.
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.