K.Boopathi
BlogNotesSecond BrainProjectsHire MeYour reading file0 pts

Congnee

Mar 24, 20268 min read

Cognee Architecture: Comprehensive Guide

Overview

Cognee is an open-source AI memory platform that transforms raw data into persistent knowledge graphs for AI agents. It replaces traditional RAG (Retrieval-Augmented Generation) with an ECL (Extract → Cognify → Load) pipeline combining vector search, graph databases, and LLM-powered entity extraction.

Key Stats:

  • 14,538+ GitHub stars
  • Python 3.9 - 3.12
  • Apache 2.0 License
  • 1,451+ forks

Core Workflow: Four Main Functions

ADD → COGNIFY → SEARCH/MEMIFY

1. add() - Data Ingestion

Purpose: Ingest data (files, URLs, text) into datasets

Flow:

add() 
  → resolve_data_directories() 
  → ingest_data() 
  → save_data_item_to_storage() 
  → Create Dataset + Data records in relational DB

Key Files:

  • cognee/api/v1/add/add.py
  • cognee/tasks/ingestion/ingest_data.py

Supports:

  • Text strings
  • Files (PDF, DOCX, CSV, images, audio, code files)
  • URLs
  • Local file paths

2. cognify() - Knowledge Graph Construction

Purpose: Extract entities/relationships and build knowledge graph

Flow:

cognify()
  → classify_documents() 
  → extract_chunks_from_documents() 
  → extract_graph_from_data() [LLM + Instructor for structured extraction]
  → summarize_text() 
  → add_data_points() [store in graph + vector DBs]

Key Files:

  • cognee/api/v1/cognify/cognify.py
  • cognee/tasks/graph/extract_graph_from_data.py
  • cognee/tasks/storage/add_data_points.py

Outputs:

  • Knowledge graph nodes and edges
  • Vector embeddings for semantic search
  • Document summaries

3. search() - Knowledge Retrieval

Purpose: Query knowledge using various retrieval strategies

Flow:

search(query_text, query_type)
  → route to appropriate retriever
  → filter by permissions
  → return results

13 Search Types Available:

Key Files:

  • cognee/api/v1/search/search.py
  • cognee/modules/retrieval/context_providers/TripletSearchContextProvider.py
  • cognee/modules/search/types/SearchType.py

4. memify() - Graph Enrichment

Purpose: Enrich graph with additional context and rules

Features:

  • Applies learning algorithms to improve knowledge graph
  • Adds contextual rules and patterns
  • Enhances entity relationships

Architecture Layers

┌─────────────────────────────────────────────┐
│  API Layer                                   │
│  (cognee/api/v1/)                          │
│  Routes: add, cognify, search, memify,     │
│  datasets, users, visualize                │
└──────────────┬──────────────────────────────┘
               ↓
┌─────────────────────────────────────────────┐
│  Main Functions                              │
│  (add, cognify, search, memify)             │
└──────────────┬──────────────────────────────┘
               ↓
┌─────────────────────────────────────────────┐
│  Pipeline Orchestrator                       │
│  (cognee/modules/pipelines/)                │
│  Manages task composition and execution      │
└──────────────┬──────────────────────────────┘
               ↓
┌─────────────────────────────────────────────┐
│  Task Execution Layer                        │
│  (cognee/tasks/)                            │
│  Reusable composable units (async)          │
└──────────────┬──────────────────────────────┘
               ↓
┌─────────────────────────────────────────────┐
│  Domain Modules                              │
│  (cognee/modules/)                          │
│  graph, retrieval, ingestion, ontology,     │
│  processing, observability, etc.            │
└──────────────┬──────────────────────────────┘
               ↓
┌─────────────────────────────────────────────┐
│  Infrastructure Adapters                     │
│  (cognee/infrastructure/)                   │
│  LLM, databases, embeddings, files,         │
│  storage backends                           │
└──────────────┬──────────────────────────────┘
               ↓
┌─────────────────────────────────────────────┐
│  External Services                           │
│  OpenAI, Kuzu, LanceDB, Neo4j, etc.        │
└─────────────────────────────────────────────┘

Key Architectural Patterns

1. Pipeline-Based Processing

  • All data flows through task-based pipelines (cognee/modules/pipelines/)
  • Tasks are composable units that can run sequentially or in parallel
  • Example tasks: classify_documents, extract_graph_from_data, add_data_points
from cognee.modules.pipelines.tasks.Task import Task
 
async def my_custom_task(data):
    processed_data = process(data)
    return processed_data
 
task = Task(my_custom_task)

2. Interface-Based Database Adapters

Supports multiple backends through adapter interfaces:

Graph Databases:

  • Kuzu (default) - Fast, embedded graph DB
  • Neo4j - Popular enterprise graph DB
  • Neptune - AWS managed service

Vector Databases:

  • LanceDB (default) - Columnar vector DB
  • ChromaDB - Vector database
  • PGVector - PostgreSQL extension
  • Qdrant, Weaviate, Milvus - Other options

Relational Databases:

  • SQLite (default) - Local, file-based
  • PostgreSQL - Production-grade

Key Interfaces:

  • cognee/infrastructure/databases/graph/graph_db_interface.py
  • cognee/infrastructure/databases/vector/vector_db_interface.py

3. Multi-Tenant Access Control

  • Hierarchy: User → Dataset → Data
  • Enable with: ENABLE_BACKEND_ACCESS_CONTROL=True
  • Isolated databases per user+dataset (Kuzu, LanceDB, SQLite, Postgres)
  • Permissions: Read, write, delete, share per dataset

Core Data Models

Engine Models (cognee/infrastructure/engine/models/)

class DataPoint:
    """Base class for all graph nodes"""
    id: str
    version: int
    metadata: dict
    
class Edge:
    """Graph relationships"""
    source: str
    target: str
    relationship_type: str
    
class Triplet:
    """(Subject, Predicate, Object) representation"""
    subject: str
    predicate: str
    object: str

Graph Models (cognee/shared/data_models.py)

class KnowledgeGraph:
    """Container for nodes and edges"""
    nodes: List[Node]
    edges: List[Edge]
    
class Node:
    """Entity in the graph"""
    id: str
    name: str
    type: str
    description: str
    
class Edge:
    """Relationship between nodes"""
    source_node_id: str
    target_node_id: str
    relationship_name: str

Key Infrastructure Components

LLM Gateway (cognee/infrastructure/llm/LLMGateway.py)

Unified interface for multiple LLM providers:

  • OpenAI (default)
  • Anthropic Claude
  • Google Gemini
  • Ollama (local models)
  • Mistral
  • AWS Bedrock
  • Azure OpenAI
  • Custom/OpenRouter/vLLM

Uses Instructor for structured output extraction (converts LLM responses into Pydantic models).

Embedding Engines

Factory pattern for embeddings:

cognee/infrastructure/databases/vector/embeddings/get_embedding_engine.py

Supports:

  • OpenAI embeddings
  • Ollama embeddings
  • HuggingFace embeddings
  • Local models

Document Loaders

Support for multiple document formats:

cognee/infrastructure/files/

Formats:

  • PDF, DOCX, CSV, TXT
  • Images (OCR-capable)
  • Audio (transcription)
  • Code files
  • Web content

Directory Structure

cognee/
├── api/v1/                      # FastAPI routes
│   ├── add/                      # Data ingestion
│   ├── cognify/                  # Graph construction
│   ├── search/                   # Query interface
│   ├── memify/                   # Graph enrichment
│   ├── datasets/                 # Dataset management
│   ├── users/                    # Authentication
│   ├── visualize/                # Graph visualization
│   └── ...
├── cli/                          # Command-line interface
├── infrastructure/               # Infrastructure adapters
│   ├── databases/                # Graph, vector, relational DBs
│   ├── llm/                      # LLM providers & gateway
│   ├── files/                    # Document loaders
│   └── storage/                  # File storage backends
├── modules/                      # Domain logic
│   ├── graph/                    # Graph operations
│   ├── retrieval/                # Retrieval strategies
│   ├── ontology/                 # OWL ontology support
│   ├── pipelines/                # Pipeline orchestration
│   ├── search/                   # Search types & implementations
│   └── ...
├── tasks/                        # Reusable async tasks
│   ├── ingestion/                # Data ingestion tasks
│   ├── graph/                    # Graph extraction tasks
│   ├── storage/                  # Storage tasks
│   └── ...
├── shared/                       # Cross-cutting utilities
│   ├── logging_utils.py          # Logging
│   ├── data_models.py            # Shared data structures
│   └── settings.py               # Configuration
└── tests/                        # Test suite
    ├── unit/                     # Unit tests
    ├── integration/              # Integration tests
    ├── cli_tests/                # CLI tests
    └── tasks/                    # Task tests

Configuration & Setup

# Create virtual environment
uv venv && source .venv/bin/activate
 
# Install Cognee
uv pip install -e .
 
# Set environment variables
LLM_API_KEY="your_openai_api_key"
LLM_MODEL="openai/gpt-4o-mini"

Default Databases (No Setup Needed):

  • Relational: SQLite
  • Vector: LanceDB
  • Graph: Kuzu

Switch to Different Providers

PostgreSQL + PGVector:

DB_PROVIDER=postgres
VECTOR_DB_PROVIDER=pgvector
GRAPH_DATABASE_PROVIDER=kuzu

Neo4j:

GRAPH_DATABASE_PROVIDER=neo4j
GRAPH_DATABASE_URL=bolt://localhost:7687

Python SDK Entry Points (cognee/__init__.py)

import cognee
import asyncio
 
async def main():
    # 1. Ingest data
    await cognee.add("Cognee turns documents into AI memory.", 
                     dataset_name="my_project")
    
    # 2. Build knowledge graph
    await cognee.cognify(datasets=["my_project"])
    
    # 3. Enrich with memory algorithms
    await cognee.memify()
    
    # 4. Query the graph
    results = await cognee.search("What does Cognee do?", 
                                 query_type="GRAPH_COMPLETION")
    
    # Additional functions
    await cognee.delete(data_id="some_id")                    # Remove data
    await cognee.config()                                      # Configuration
    await cognee.datasets()                                    # Dataset ops
    
asyncio.run(main())

All functions are async - use await or asyncio.run().

Critical Data Flow Examples

ADD Flow (Data Ingestion)

User calls: await cognee.add(data, dataset_name="my_data")
    ↓
resolve_data_directories()
    ↓
ingest_data() 
    ↓
save_data_item_to_storage()
    ↓
Creates:
  - Dataset record
  - Data records in relational DB
  - Files stored in storage backend

COGNIFY Flow (Graph Construction)

User calls: await cognee.cognify(datasets=["my_data"])
    ↓
Pipeline executes tasks sequentially:
  1. classify_documents()
  2. extract_chunks_from_documents()
  3. extract_graph_from_data()
     - LLM extracts entities/relationships
     - Uses Instructor for structured output
  4. summarize_text()
  5. add_data_points()
    ↓
Stores:
  - Nodes/edges in graph DB (Kuzu)
  - Embeddings in vector DB (LanceDB)
  - Summaries in relational DB

SEARCH Flow (Retrieval)

User calls: await cognee.search("What does Cognee do?", query_type="GRAPH_COMPLETION")
    ↓
Router selects retriever based on query_type
    ↓
Retriever executes:
  - For GRAPH_COMPLETION: Graph traversal + LLM completion
  - For CHUNKS: Vector similarity search
  - For CYPHER: Direct graph query
    ↓
Permission filtering applied
    ↓
Results returned to user

Extension Points

Add Custom Task

# 1. Create task in cognee/tasks/
async def my_custom_task(data):
    processed = your_logic(data)
    return processed
 
# 2. Register in pipeline
from cognee.modules.pipelines import Pipeline
await Pipeline().register_task(my_custom_task)

Add Custom Database Backend

# Implement GraphDBInterface or VectorDBInterface
cognee/infrastructure/databases/graph/graph_db_interface.py
cognee/infrastructure/databases/vector/vector_db_interface.py

Add Custom Search Type

# 1. Add to SearchType enum
cognee/modules/search/types/SearchType.py
 
# 2. Implement retriever
cognee/modules/retrieval/context_providers/

Add Custom LLM Provider

Uses litellm library - add configuration:

LLM_PROVIDER="your_provider"
LLM_MODEL="model_name"
LLM_ENDPOINT="api_endpoint"
LLM_API_KEY="your_key"

Key Concepts

Datasets

  • Project-level containers for organizing data
  • Support organization, permissions, isolated workflows
  • Each user can have multiple datasets with different access levels

DataPoints

  • Atomic knowledge units forming graph structure
  • All graph nodes extend DataPoint base class
  • Include versioning and metadata support

Permissions System

  • Multi-tenant architecture (users, roles, ACLs)
  • Read, write, delete, share permissions per dataset
  • Enable: ENABLE_BACKEND_ACCESS_CONTROL=True

Graph Visualization

# Via CLI
cognee-cli -ui  # Launches full stack at http://localhost:3000
 
# Via Python
from cognee.api.v1.visualize import start_visualization_server
await start_visualization_server(port=8080)

Testing Strategy

cognee/tests/
├── unit/             # Individual module tests
├── integration/      # Full pipeline tests
├── cli_tests/        # CLI command tests
└── tasks/            # Task-specific tests

Run Tests:

pytest                                      # All tests
pytest --cov=cognee --cov-report=html     # With coverage
pytest cognee/tests/integration/           # Integration only
pytest cognee/tests/unit/                  # Unit only

Distributed Execution

Cognee supports distributed processing via Modal:

# distributed/entrypoint.py
- Spawns graph_saving_worker (1 instance)
- Spawns data_point_saving_worker (10 instances)
- Processes queued tasks asynchronously

Code Style & Quality

  • Formatter: Ruff
  • Line length: 100 characters
  • Quotes: Double quotes "
  • Pre-commit hooks: Automatic linting/formatting
  • Type hints: Encouraged (mypy checks)
# Before committing
pre-commit run --all-files

Common Usage Pattern

import asyncio
import cognee
 
async def main():
    # Clean slate (optional)
    await cognee.prune.prune_data()
    await cognee.prune.prune_system(metadata=True)
    
    # 1. Add data
    await cognee.add([
        "Alice knows Bob.",
        "NLP is part of AI.",
        "Berlin is Germany's capital."
    ])
    
    # 2. Build knowledge graph
    await cognee.cognify()
    
    # 3. Search
    results = await cognee.search("What is NLP?")
    print(results)
    
    # 4. Visualize
    await cognee.visualize_graph("graph.html")
 
asyncio.run(main())

Graph View