Custom Graph Model

A minimal guide to creating custom graph models and loading them into Cognee with remember(). Before you start:

Complete Quickstart to understand basic operations
Ensure you have LLM Providers configured
Have some structured data you want to model

What Custom Graph Models Do

Define a schema (Pydantic models inheriting from DataPoint) that constrains how Cognee extracts entities and relationships from text
Guide remember() to produce graph nodes and edges that match the defined domain model
Improve consistency of labels, properties, and relationship structure across ingested documents
Control search quality by choosing which fields are indexed via metadata

Code in Action

Step 1: Define Your Entity Classes

class Activity(DataPoint):
    name: str
    metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]}

class Person(DataPoint):
    name: str
    likes: List[Activity] | None = None
    metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]}

Create Pydantic models that inherit from DataPoint to represent the node types in your custom graph. The nested likes: List[Activity] field tells Cognee to extract Activity nodes and the likes edges connecting each Person to them. The metadata dict controls how each node is indexed and identified:

index_fields — the fields Cognee embeds and indexes for retrieval (here, name).
identity_fields — the fields used to derive a deterministic node id (namespaced by class name). Without them a node gets a random id, so repeated mentions never merge; with identity_fields=["name"], the same Alice or a shared activity like board games collapses into a single node instead of creating duplicates.

Step 2: Define Your Top-Level Graph Container

class PeopleGraph(DataPoint):
    people: List[Person]

Wrap your top-level entities in a container model. This is the model you pass as graph_model, and it tells the LLM to return a list of Person entities (each with their nested activities).

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, only the fields you declare on each subclass become part of the extraction schema. Inherited DataPoint infrastructure fields (like id) are stripped out, so they do not expand the LLM’s response schema. Adding a field (e.g. age: int on Person) tells the LLM to extract that value; nested DataPoint fields (like likes: List[Activity]) tell it to extract those related entities and the edges between them.
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.

Step 3: Remember Your Data with the Custom Model

CUSTOM_PROMPT = (
    "Extract all people mentioned in the text. "
    "For each person, extract ALL activities they like, including shared activities."
)

text = (
    "Alice likes biking and swimming. Bob likes playing basketball. "
    "Alice and Bob are friends. "
    "Charlie likes skiing. "
    "Alice and Bob like playing board games together."
)

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, and CUSTOM_PROMPT tells the LLM what to look for — here, every person and all the activities they like, including activities shared between people.

Step 4: Visualize Your Data

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

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

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

Additional examples

Additional examples about custom data models are available on our github.

Full Example

Latest guide

import asyncio
import os
from typing import List

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

CUSTOM_PROMPT = (
    "Extract all people mentioned in the text. "
    "For each person, extract ALL activities they like, including shared activities."
)


class Activity(DataPoint):
    name: str
    metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]}


class Person(DataPoint):
    name: str
    likes: List[Activity] | None = None
    metadata: dict = {"index_fields": ["name"], "identity_fields": ["name"]}


class PeopleGraph(DataPoint):
    people: List[Person]


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

    text = (
        "Alice likes biking and swimming. Bob likes playing basketball. "
        "Alice and Bob are friends. "
        "Charlie likes skiing. "
        "Alice and Bob like playing board games together."
    )

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

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


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

Legacy guide

import os  
import asyncio  
from typing import Any, List  
from pydantic import SkipValidation  
  
from cognee import add, cognify, prune, visualize_graph  
from cognee.low_level import DataPoint  
  
CUSTOM_PROMPT = (  
    "Extract a simple graph containing Programming Language and Fields that it is used in."  
) 
  
# Define a custom graph model for programming languages.  
class FieldType(DataPoint):  
    name: str = "Field"  
  
class Field(DataPoint):  
    name: str  
    is_type: FieldType  
    metadata: dict = {"index_fields": ["name"]}  
  
class ProgrammingLanguageType(DataPoint):  
    name: str = "Programming Language"  
   
class ProgrammingLanguage(DataPoint):  
    name: str  
    used_in: List[Field] = None  
    is_type: ProgrammingLanguageType  
    metadata: dict = {"index_fields": ["name"]}  
   
async def visualize_data():  
    graph_file_path = os.path.join(  
        os.path.dirname(__file__), ".artifacts", "custom_graph_model_entity_schema_definition.html"  
    )  
    await visualize_graph(graph_file_path)  
  
async def main():  
    # Prune data and system metadata before running, only if we want "fresh" state.  
    await prune.prune_data()  
    await prune.prune_system(metadata=True)  
  
    text = "The Python programming language is widely used in data analysis, web development, and machine learning."  
  
    await add(text)  
    await cognify(graph_model=ProgrammingLanguage, custom_prompt=CUSTOM_PROMPT)  
  
    await visualize_data()  
   
if __name__ == "__main__":  
    asyncio.run(main())

This example shows the complete workflow with metadata for indexing. In practice, you can create complex nested models with multiple relationships and sophisticated data structures.

Low-Level LLM

Learn about direct LLM interaction

Core Concepts

Understand knowledge graph fundamentals

API Reference

Explore API endpoints

Getting Started

Core Concepts

Setup Configuration

Guides

Examples

CLI

Rust SDK

TypeScript SDK

OSS

What Custom Graph Models Do

Code in Action

Step 1: Define Your Entity Classes

Step 2: Define Your Top-Level Graph Container

Step 3: Remember Your Data with the Custom Model

Step 4: Visualize Your Data

Use in Custom Tasks and Pipelines

Additional examples

Full Example

Low-Level LLM

Core Concepts

API Reference

​What Custom Graph Models Do

​Code in Action

​Step 1: Define Your Entity Classes

​Step 2: Define Your Top-Level Graph Container

​Step 3: Remember Your Data with the Custom Model

​Step 4: Visualize Your Data

​Use in Custom Tasks and Pipelines

​Additional examples

​Full Example

Low-Level LLM

Core Concepts

API Reference

What Custom Graph Models Do

Code in Action

Step 1: Define Your Entity Classes

Step 2: Define Your Top-Level Graph Container

Step 3: Remember Your Data with the Custom Model

Step 4: Visualize Your Data

Use in Custom Tasks and Pipelines

Additional examples

Full Example