> ## Documentation Index > Fetch the complete documentation index at: https://docs.cognee.ai/llms.txt > Use this file to discover all available pages before exploring further. # Custom Graph Model > Step-by-step guide to creating custom graph models and using remember with them 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](/getting-started/quickstart) to understand basic operations * Ensure you have [LLM Providers](/setup-configuration/llm-providers) configured * Read [DataPoints](/core-concepts/building-blocks/datapoints) for the conceptual overview of nodes, edges, and metadata * Have some structured data you want to model ## Code in Action ```python theme={null} 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 ```python theme={null} 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 ```python theme={null} 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 ```python theme={null} 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 ```python theme={null} 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 A custom graph model constrains extraction *before* the LLM runs, while an [ontology](/core-concepts/further-concepts/ontologies) grounds names *after* it. See "How does an ontology relate to a custom graph?" under [Additional details and examples](/core-concepts/further-concepts/ontologies#additional-details-and-examples) for the trade-offs and for how to combine them. 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. 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. 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 : `); 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. 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](/core-concepts/building-blocks/datapoints#transparent-containers-nodes-that-group-rather-than-describe). ## 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](/guides/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](/guides/custom-prompts) or [custom tasks](/guides/custom-tasks-pipelines) to refine extraction * Validate pipeline results with `visualize_graph` before promoting changes to production The node, edge, and metadata model these classes are built on. Every parameter `remember()` accepts, including `graph_model`. Build DataPoints yourself and insert them with `add_data_points`, skipping LLM extraction. Runnable guide scripts, including this one, in the cognee repo. This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.