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.

Identity and search: metadata