Custom Graph Model

A minimal guide to creating custom graph models and loading them into Cognee with remember(). Reach for one when you want extraction to return a specific set of entity types and relationships instead of a free-form knowledge graph.

Typed Edge[...] fields and FromIdentity() are Python SDK only. A graph model sent over HTTP (POST /cognify, POST /remember) travels as a JSON schema and is rebuilt from it, and that schema carries neither the Edge generics nor the annotation markers — the typed edges are silently lost. Pass the model class directly through the SDK to use them.

Before You Start

Complete Quickstart to understand basic operations
Ensure you have LLM Providers configured
Read DataPoints for the conceptual overview of nodes, edges, and metadata
Have some structured data you want to model

Code in Action

import asyncio
import os
from typing import Annotated, Literal

from cognee import forget, remember, visualize_graph
from cognee.low_level import DataPoint, Edge, FromIdentity

CUSTOM_PROMPT = (
    "Extract every person, the role they hold, and every group with its members. "
    "Extract friendships, family links (married_to or sibling_of), who reports to whom, "
    "and other named relationships between people."
)


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
    # An edge can also live on the node that owns it. Endpoints of the same type have to
    # be named as strings here, because Person is not bound inside its own body yet.
    # Put an edge here when one side clearly owns it, as each person has one manager.
    reports_to: list[Edge["Person", "Person"]] = []
    metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]}


class Group(DataPoint):
    name: str
    members: list[Person] | None = None
    metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]}


class PeopleGraph(DataPoint):
    # Edges on the root suit a relationship with no obvious owner. Each one shows a way
    # of naming: fixed by the field, chosen from a Literal, or free-form from the LLM.
    # When building Edge values by hand, set source= explicitly here: an omitted source
    # falls back to the declaring node, and this root is not a Person.
    people: list[Person]
    groups: list[Group] = []
    friends_with: list[Edge[Person, Person]] = []
    family_links: list[Edge[Person, Person, Literal["married_to", "sibling_of"]]] = []
    other_links: list[Edge[Person, Person, str]] = []


async def main():
    await forget(everything=True)

    text = (
        "Maya and Owen are engineers on the Search team and are friends. "
        "Priya is a manager and Maya's sibling. Owen mentors Maya. "
        "Maya and Owen both report to Priya."
    )

    await remember(
        text,
        graph_model=PeopleGraph,
        custom_prompt=CUSTOM_PROMPT,
        self_improvement=False,
    )

    graph_path = os.path.join(os.path.dirname(__file__), ".artifacts", "custom_graph.html")
    await visualize_graph(graph_path)


if __name__ == "__main__":
    asyncio.run(main())

What Just Happened

Step 1: Define Your Entity Classes and Relationships

from typing import Annotated, Literal

from cognee import forget, remember, visualize_graph
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
    # An edge can also live on the node that owns it. Endpoints of the same type have to
    # be named as strings here, because Person is not bound inside its own body yet.
    # Put an edge here when one side clearly owns it, as each person has one manager.
    reports_to: list[Edge["Person", "Person"]] = []
    metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]}


class Group(DataPoint):
    name: str
    members: list[Person] | None = None
    metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]}

Create Pydantic models that inherit from DataPoint to represent your node types. A field holding another DataPoint becomes an edge named after the field — members: list[Person] extracts Person nodes and the members edges from each Group — while Annotated[Role, FromIdentity()] asks the LLM for a role’s identity string instead of a nested object, and list[Edge["Person", "Person"]] asks for flat reports_to relationship rows. The metadata dict sets which fields are embedded for retrieval (index_fields) and which derive a deterministic node id (identity_fields), so repeated mentions of Maya collapse into one node instead of duplicating.

Step 2: Define Your Top-Level Graph Container

class PeopleGraph(DataPoint):
    # Edges on the root suit a relationship with no obvious owner. Each one shows a way
    # of naming: fixed by the field, chosen from a Literal, or free-form from the LLM.
    # When building Edge values by hand, set source= explicitly here: an omitted source
    # falls back to the declaring node, and this root is not a Person.
    people: list[Person]
    groups: list[Group] = []
    friends_with: list[Edge[Person, Person]] = []
    family_links: list[Edge[Person, Person, Literal["married_to", "sibling_of"]]] = []
    other_links: list[Edge[Person, Person, str]] = []

Wrap your top-level entities in a container model; this is what you pass as graph_model. Declaring an edge on the root fits a relationship neither side owns, such as a friendship, while reports_to sits on Person because the reporting line belongs to the person rather than to the graph as a whole. The third Edge type parameter decides how the relationship is named.

Step 3: Remember Your Data with the Custom Model

CUSTOM_PROMPT = (
    "Extract every person, the role they hold, and every group with its members. "
    "Extract friendships, family links (married_to or sibling_of), who reports to whom, "
    "and other named relationships between people."
)

text = (
    "Maya and Owen are engineers on the Search team and are friends. "
    "Priya is a manager and Maya's sibling. Owen mentors Maya. "
    "Maya and Owen both report to Priya."
)

await remember(
    text,
    graph_model=PeopleGraph,
    custom_prompt=CUSTOM_PROMPT,
    self_improvement=False,
)

This ingests the text and builds the graph in one call. The custom graph_model acts as the extraction schema while CUSTOM_PROMPT tells the LLM what to look for, and the prompt names the married_to and sibling_of options explicitly so the LLM fills family_links with the values that Literal allows.

Step 4: Visualize Your Data

graph_path = os.path.join(os.path.dirname(__file__), ".artifacts", "custom_graph.html")
await visualize_graph(graph_path)

This renders the generated graph to custom_graph.html so you can verify nodes, relationships, and overall schema behavior.

Advanced Usage

Custom graph models and ontologies

A custom graph model constrains extraction before the LLM runs, while an ontology grounds names after it. See “How does an ontology relate to a custom graph?” under Additional details and examples for the trade-offs and for how to combine them.

How graph_model constrains LLM extraction

When you pass graph_model=..., that model is the structured-output schema the LLM must fill in. Internally, Cognee hands your model to the LLM as the response_model for structured extraction, so the LLM can only return entities and relationships that fit the fields you declared — it is not free to invent an arbitrary shape.

Default: without graph_model, Cognee uses its general-purpose KnowledgeGraph schema (free-form nodes and edges).
Custom model: when you pass a DataPoint subclass, the extraction schema is built from every domain field the model carries — including fields it inherits from your own intermediate DataPoint subclasses, not just the ones annotated on the class you pass. If Animal(DataPoint) declares species: str and Dog(Animal) adds breed: str, the LLM is asked for both. Only the fields defined on DataPoint itself (id, version, type, created_at, metadata, and the rest of the infrastructure fields) are stripped out, so they do not expand the LLM’s response schema; the stripping is by field name, so a subclass that overrides metadata — as Person does above — still keeps it out of the schema. Adding a field (e.g. age: int on Person) tells the LLM to extract that value; nested DataPoint fields (like members: list[Person] on Group) tell it to extract those related entities and the edges between them, and Edge[...] fields ask for relationship rows it resolves against those entities.
custom_prompt vs graph_model: they play different roles. graph_model defines the shape (which fields and relationships are allowed), while custom_prompt replaces the system prompt that tells the LLM what to look for. Use them together for predictable, domain-specific extraction.

Naming a typed relationship

The third type parameter of Edge controls how the relationship is named:

Omitted (friends_with) — the edge takes the field’s name.
Literal["married_to", "sibling_of"] (family_links) — the LLM picks one of the listed names.
str (other_links) — the LLM supplies a free-form name, which Cognee normalizes.

How edge rows are resolved

Cognee resolves each row’s source and target against the nodes extracted from the same text, matching on the identity value. When a row names an endpoint that no extracted node matches, that single row is dropped and logged (Skipping unresolved edge on <field>: <row>); the rest of the chunk still stores.Because both features address nodes by their identity value, every class at either end of an edge — and every FromIdentity() target — needs exactly one identity_fields entry, and a FromIdentity() target must be constructible from that field alone. FromIdentity() accepts Target, Target | None, list[Target], and list[Target] | None; any other shape would hand the LLM a whole nested object and mint a duplicate node, so it is rejected outright. Cognee raises InvalidReferenceTypeError when a model breaks either rule, while it builds the extraction schema and before the LLM is called.Those rules cover edges extraction fills in. When you construct Edge values by hand instead, an omitted source falls back to the node the field is declared on — which is what you want on an owner-declared edge like reports_to, but not on a root container: PeopleGraph is not a Person, so a parametrized edge there raises ValueError unless you set source= explicitly. A missing target is always an error. The older unparametrized tuple form records no generics and keeps the permissive fallback.

Keeping the container out of the graph

The container is stored as a node too, so the graph gains a PeopleGraph node with a people edge to each Person. To keep it as an extraction schema only, add metadata: dict = {"index_fields": [], "transparent": True} — Cognee then skips the wrapper and promotes each Person to a top-level node. See Transparent containers.

Declaring a model in JSON

The same model can be written as a plain JSON document — a graph schema spec — and compiled with graph_model_from_spec(), which returns the DataPoint-derived class you pass as graph_model. It is the shape the Cognee UI’s graph-model editor produces, so a model built in the UI can be handed straight to the SDK. See Graph Model from JSON for the spec reference, the validation rules, and a runnable walkthrough.

Use in Custom Tasks and Pipelines

This pattern is useful when you need predictable, domain-specific extraction inside custom workflows.

Reuse the same graph schema across tasks to keep outputs consistent
Run remember(graph_model=...) in workflows where downstream logic expects a specific graph shape
Combine with custom prompts or custom tasks to refine extraction
Validate pipeline results with visualize_graph before promoting changes to production

DataPoints

The node, edge, and metadata model these classes are built on.

remember()

Every parameter remember() accepts, including graph_model.

Custom Data Models

Build DataPoints yourself and insert them with add_data_points, skipping LLM extraction.

More examples

Runnable guide scripts, including this one, in the cognee repo.

Getting Started

Core Concepts

Setup Configuration

Guides

Examples

CLI

OSS

Before You Start

Code in Action

What Just Happened

Step 1: Define Your Entity Classes and Relationships

Step 2: Define Your Top-Level Graph Container

Step 3: Remember Your Data with the Custom Model

Step 4: Visualize Your Data

Advanced Usage

Declaring a model in JSON

Use in Custom Tasks and Pipelines

DataPoints

remember()

Custom Data Models

More examples

​Before You Start

​Code in Action

​What Just Happened

​Step 1: Define Your Entity Classes and Relationships

​Step 2: Define Your Top-Level Graph Container

​Step 3: Remember Your Data with the Custom Model

​Step 4: Visualize Your Data

​Advanced Usage

​Declaring a model in JSON

​Use in Custom Tasks and Pipelines

DataPoints

remember()

Custom Data Models

More examples

Before You Start

Code in Action

What Just Happened

Step 1: Define Your Entity Classes and Relationships

Step 2: Define Your Top-Level Graph Container

Step 3: Remember Your Data with the Custom Model

Step 4: Visualize Your Data

Advanced Usage

Declaring a model in JSON

Use in Custom Tasks and Pipelines