​
Customize data ingestion using Pydantic

Difficulty: Medium

​
Overview

Cognee let’s you organize and model your user’s data for LLMs to use. In this way you can choose how to load only the data you need. Let’s say you need all persons mentioned in a novel. We enable you to:
  1. Specify which persons you want extracted
  2. Load them into the cognee data store
  3. Retrieve them with natural language query
In other words, custom models let you define the structure of your graph, or impose certain structures on it. You can control what nodes get created and how they’re connected. This is useful when dealing with structured or semi-structured data, though they do require some manual coding. Let’s try it out!

​
Let’s model your data based on your preferences

Why is this important? Let’s visualize our data before and after. On this image you can see that purple color nodes are exactly the nodes that represent people mentioned in the novel. After Let’s create the graph ourselves.

​
Step 1: Clone Required Repositories

​
Clone Main Repository

First, clone the main Cognee repository: 
git clone https://github.com/topoteretes/cognee.git

​
Clone Starter Repository

Clone the getting started repository with examples: 
git clone https://github.com/topoteretes/cognee-starter.git
These repositories contain all the necessary code and examples for custom data modeling.

​
Step 2: Install Dependencies

cd cognee

​
Install with UV

Install Cognee with all development dependencies: 
uv sync --dev --all-extras --reinstall
This ensures you have all the necessary packages for custom data model development.

​
Step 3: Create Your Custom Model Script

​
Use Example from Starter Repository

Create a Python script called example_ontology.py and copy the content from the following file: Custom Model Example This example demonstrates how to define custom Pydantic models for specific data extraction.

​
Understand the Model Structure

The custom model defines exactly which entities you want extracted and how they should be structured in your knowledge graph.

​
Step 4: Execute Your Script

​
Run the Custom Model Script

Activate the virtual environment and execute your script using Python: 
source .venv/bin/activate && python example_ontology.py
Make sure that the script has access to the data in the cognee-starter repository.

​
Monitor Execution

The script will process your data and create entities according to your custom model definitions.

​
Step 5: Inspect Your Knowledge Graph

​
Generate Visualization

The script will create an HTML file in the cognee directory (.artifacts/graph_visualization.html) that you can inspect and check the graph. You can also run a small HTTP server that will render your semantic layer: 
import webbrowser
import os
from cognee.api.v1.visualize.visualize import visualize_graph

await visualize_graph()
home_dir = os.path.expanduser("~")
html_file = os.path.join(home_dir, "graph_visualization.html")
webbrowser.open(f"file://{html_file}")
# display(html_file) in notebook

​
Analyze Results

In the generated visualization, you’ll see:
  • Purple nodes representing the people entities you defined
  • Structured relationships based on your custom model
  • Clean, organized data extraction focused on your specific needs

​
Advanced Customization

​
Define More Complex Models

You can extend your custom models to include additional properties and relationships: 
from pydantic import BaseModel
from typing import List, Optional

class Person(BaseModel):
    name: str
    role: Optional[str] = None
    relationships: List[str] = []
    attributes: Optional[dict] = None

​
Handle Different Data Types

Custom models can be adapted for various content types:
  • Literary texts (characters, themes, settings)
  • Business documents (people, organizations, projects)
  • Technical documentation (components, processes, dependencies)

​
Troubleshooting

​
Common Issues

Model not extracting expected entities:
  • Verify your model definitions match the content structure
  • Check that field names are descriptive and relevant
  • Ensure your text contains the entities you’re trying to extract
Script execution errors:
  • Confirm all dependencies are installed correctly
  • Check file paths and data accessibility
  • Verify your Python environment is properly configured

​
Next Steps

Now that you’ve created custom data models, you can:
  1. Expand your models with more complex entity types
  2. Integrate multiple models for comprehensive data extraction
  3. Build domain-specific applications using your structured data
  4. Create automated pipelines for ongoing data processing

​
Join the Conversation!

Have questions? Join our community now to connect with professionals, share insights, and get your questions answered!