How to Define a Custom Graph Model in Cognee: A Developer’s Guide

To define a custom graph model in Cognee, subclass the DataPoint base class from cognee.infrastructure.engine.models.DataPoint, declare typed attributes for your domain entities, and pass the root model to cognee.cognify() to automatically materialize the knowledge graph.

Cognee is an open-source knowledge graph engine that structures data as interconnected DataPoint objects. When you need to represent domain-specific concepts—such as programming languages, scientific fields, or organizational hierarchies—you create a custom graph model that extends Cognee’s core architecture. This approach preserves Cognee’s built-in versioning, vector indexing, and search capabilities while shaping the graph structure to match your exact ontology.

Understanding the DataPoint Foundation

Every node in a Cognee graph inherits from DataPoint, defined in cognee/infrastructure/engine/models/DataPoint.py. This base class provides essential infrastructure including unique identifiers, timestamps, versioning (lines 24-48), and the get_embeddable_data() method for vector search integration. When you subclass DataPoint, your custom nodes automatically receive these capabilities without additional boilerplate.

The base class also supports a metadata dictionary where you specify indexing behavior. By setting metadata["index_fields"], you indicate which attributes should be embedded for semantic search, enabling Cognee to include your custom entities in vector-based retrieval pipelines.

Step 1: Import DataPoint and Define Domain Classes

Begin your custom model by importing the base class and declaring subclasses that represent your domain entities. Each subclass becomes a node type in the resulting graph, with class attributes defining properties and relationships.

Relationships are established through type annotations: use another DataPoint subclass for one-to-one references, or list[DataPointSubclass] for one-to-many connections.

from cognee.low_level import DataPoint

class FieldType(DataPoint):
    name: str = "Field"
    metadata: dict = {"index_fields": ["name"]}

class Field(DataPoint):
    name: str
    is_type: FieldType
    metadata: dict = {"index_fields": ["name"]}

class ProgrammingLanguage(DataPoint):
    name: str
    used_in: list[Field] = []  # One-to-many relationship

    metadata: dict = {"index_fields": ["name"]}

In cognee/infrastructure/engine/models/DataPoint.py (lines 20-27), the base class uses Pydantic for validation and serialization, ensuring type safety across the graph.

Step 2: Configure Storage and Add Data

Before processing, configure Cognee’s storage directories and ingest your source data. This step is identical to standard Cognee usage, but you will specify your custom model during the graph generation phase.

import pathlib
import cognee

cognee.config.data_root_directory(
    str(pathlib.Path(__file__).parent / ".data_storage/custom")
)
cognee.config.system_root_directory(
    str(pathlib.Path(__file__).parent / ".cognee_system/custom")
)

async def main():
    await cognee.prune.prune_data()
    await cognee.add("Python is a high-level language created by Guido van Rossum.")

Step 3: Generate the Graph with Your Custom Model

Call cognee.cognify() with the graph_model parameter set to your root DataPoint subclass. Cognee introspects the class hierarchy, generates a JSON schema, and materializes the graph structure according to your definitions.

According to the source code in cognee/api/v1/cognify/routers/get_cognify_router.py, the pipeline automatically handles schema conversion and node instantiation when the graph_model argument is provided.

    await cognee.cognify(graph_model=ProgrammingLanguage)
    
    results = await cognee.search(
        "What programming languages are defined?",
        cognee.modules.search.types.SearchType.GRAPH_COMPLETION
    )
    print(results)

Advanced: Dynamic Model Generation from Schemas

For runtime flexibility, Cognee provides conversion utilities in cognee/shared/graph_model_utils.py. These functions enable bidirectional transformation between Pydantic models and JSON schemas, supporting dynamic graph generation without static class definitions.

The module exposes two primary functions:

• graph_model_to_graph_schema() (lines 60-68): Converts a Pydantic model to a JSON schema compatible with the datamodel-code-generator.
• graph_schema_to_graph_model() (lines 16-55): Generates a temporary Python module from a JSON schema, patches it to inherit from DataPoint, and returns a runtime model class.

from cognee.shared.graph_model_utils import (
    graph_model_to_graph_schema,
    graph_schema_to_graph_model,
)
from pydantic import BaseModel

class PlainLanguage(BaseModel):
    name: str
    version: str

# Convert to Cognee-compatible DataPoint subclass

schema = graph_model_to_graph_schema(PlainLanguage)
CogneeLanguage = graph_schema_to_graph_model(schema)

await cognee.cognify(graph_model=CogneeLanguage)

This pattern, implemented in the __main__ block of graph_model_utils.py (lines 87-115), is essential for applications that load domain models from external schema files or user-generated configurations.

How Custom Models Integrate with Search

Once ingested, custom graph models integrate seamlessly with Cognee’s search architecture. The SearchType enumeration in cognee/modules/search/types.py includes GRAPH_COMPLETION, which traverses relationships defined in your custom model to answer complex queries. Because your subclasses inherit DataPoint’s metadata handling, fields marked with index_fields are automatically embedded and available for vector similarity search alongside graph traversal.

Summary

• Subclass DataPoint from cognee/infrastructure/engine/models/DataPoint.py to create domain-specific node types with automatic versioning and serialization.
• Define relationships using typed annotations (single instances for one-to-one, lists for one-to-many) to structure your knowledge graph.
• Set metadata["index_fields"] to specify which attributes should be vectorized for semantic search.
• Pass your root model to cognee.cognify(graph_model=YourModel) to trigger automatic schema generation and graph instantiation.
• Use conversion utilities in cognee/shared/graph_model_utils.py for dynamic model generation from JSON schemas at runtime.
• Query via SearchType.GRAPH_COMPLETION to leverage your custom ontology in retrieval pipelines.

Frequently Asked Questions

What is the minimum required field for a custom DataPoint subclass?

A valid DataPoint subclass must inherit from the base class defined in cognee/infrastructure/engine/models/DataPoint.py and include the required Pydantic fields. While you can define any custom attributes, you should include a metadata dictionary with "index_fields" specified if you want the node to participate in vector search. The base class automatically provides id, type, and timestamp fields.

Can I define nested relationships between custom models?

Yes. Cognee supports arbitrary nesting through type annotations. Define one-to-one relationships by annotating a field with another DataPoint subclass, or one-to-many relationships using list[DataPointSubclass]. The graph engine in cognee/api/v1/cognify/routers/get_cognify_router.py automatically resolves these references when building the knowledge graph during the cognify step.

How does Cognee handle schema generation from my custom model?

When you pass a model to cognee.cognify(), the framework invokes graph_model_to_graph_schema() from cognee/shared/graph_model_utils.py (lines 60-68) to generate a JSON schema. If needed, graph_schema_to_graph_model() (lines 16-55) then uses the datamodel-code-generator to create a runtime-compatible class that patches into cognee.infrastructure.engine.DataPoint. This process ensures your custom model is fully compatible with Cognee’s internal graph representation while preserving type safety.

Is it possible to update a custom graph model after data ingestion?

While you can define new versions of your DataPoint subclasses, Cognee’s graph structure is generated during the cognify pipeline based on the model provided at that time. To change the schema, you typically need to prune the existing graph using cognee.prune.prune_data(), update your model definitions, and re-run cognee.cognify() with the new graph_model parameter. The versioning fields in the base DataPoint class help track these iterations.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}