Cognee Custom Graph Models

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.

pythonCopy code
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

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:

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