Build a Knowledge Graph with Cognee

Build an AI memory system that connects facts across documents using knowledge graphs.

Introduction

30 Second Summary

Every time you close a chat with an AI assistant, the conversation vanishes. The AI has no memory of what you discussed, what you asked, or what it learned about your world. Knowledge graphs change this by giving AI systems a structured, persistent memory that connects facts together.

In this project, you will build a knowledge graph memory system using Cognee that ingests documents about a fictional company, connects the information into a graph, and answers questions that require reasoning across multiple sources.

What You'll Build

A personal knowledge assistant that ingests a collection of documents and queries them using multiple retrieval strategies, demonstrating how knowledge graphs connect information that traditional search misses.

By the end of this project, you'll have:

  • A configured Cognee environment that uses OpenAI for embeddings and LLM calls with zero infrastructure setup.
  • A full ECL pipeline (Extract, Cognify, Load) that builds a knowledge graph from multiple documents and enriches it with memory algorithms.
  • A search mode comparison showing how graph-based retrieval produces richer answers than traditional RAG by traversing relationships across documents.
  • Secret Mission: Build a multi-source research assistant that answers cross-document reasoning questions by connecting facts across research papers, meeting notes, and technical specs.

Are there any prerequisites?

You should have experience writing Python and be familiar with async/await syntax. You also need an OpenAI API key with credits available for embeddings and LLM calls.

Set Up Your Cognee Environment

Before you can build a knowledge graph memory system, you need a clean workspace with the right libraries installed.

In this step, you will create a dedicated project folder, set up an isolated Python environment, install Cognee, and configure your OpenAI API key.

In this step, get ready to:
  • Install Cognee in a virtual environment.
  • Configure your OpenAI API key for Cognee.
  • Open your project in VS Code.
Create your virtual environment and install Cognee

A virtual environment gives your project its own isolated set of Python packages. This prevents conflicts between projects that need different library versions.

  • Open your terminal by pressing Cmd+Space (macOS) or the Windows key (Windows) and typing Terminal.
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir cognee-memory-project
cd cognee-memory-project
  • Confirm you are inside the new folder by running:
pwd

You should see a path ending in cognee-memory-project.

  • Create and activate your virtual environment by running the commands for your operating system:

macOS

python3 -m venv cognee-env
source cognee-env/bin/activate

Windows

python -m venv cognee-env
cognee-env\Scripts\activate

Linux

python3 -m venv cognee-env
source cognee-env/bin/activate

You should see (cognee-env) appear at the beginning of your terminal prompt. This confirms the virtual environment is active.

Don't see (cognee-env) in your prompt?

Make sure you ran the activation command (source on macOS/Linux or the Scripts\activate path on Windows). If you closed and reopened your terminal, you need to activate the environment again.

help me activate my virtual environment

  • Install the Cognee package by running:
pip install cognee

What is Cognee installing?

This command downloads Cognee and its dependencies, including local databases for storage. Cognee uses SQLite for relational data, LanceDB for vector storage, and Kuzu for graph queries.

No external servers are required. Everything runs on your machine.

  • Verify Cognee installed correctly by running:
python -c "import cognee; print('Cognee installed successfully')"

You should see Cognee installed successfully printed in your terminal.

Seeing ModuleNotFoundError?

This means the install did not complete or you are running Python outside your virtual environment. Check that (cognee-env) appears in your prompt before running the verify command.

If the install itself failed, try upgrading pip first with pip install --upgrade pip and then run pip install cognee again.

help me fix my Cognee installation

Open your project in VS Code
  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Visual Studio Code and press Enter to open it.
  • In VS Code, click File in the top menu bar.
  • Click Open Folder.
  • Navigate to your Desktop, select the cognee-memory-project folder, and click Open.

You should see the cognee-env folder listed in the Explorer sidebar on the left.

Configure your environment variables

Cognee reads your API credentials from a .env file in your project root. This single file provides the key that Cognee uses for both LLM calls and embedding generation.

  • In the Explorer sidebar, right-click in the empty space below cognee-env.
  • Select New File....
  • Type .env as the file name and press Enter.
  • Paste the following content into the new file:
LLM_API_KEY=your-openai-api-key-here
  • Replace your-openai-api-key-here with your actual OpenAI API key.
  • Save the file with Cmd+S (macOS) or Ctrl+S (Windows).

What does LLM_API_KEY configure?

Cognee uses this single key to power two things: LLM calls (defaulting to openai/gpt-5-mini) and embedding generation (defaulting to openai/text-embedding-3-large).

Everything else runs locally with no external servers required.

Not sure where to find your OpenAI API key?

Log in to your OpenAI account and navigate to the API keys page. Copy an existing key or create a new one. Make sure your account has credits available.

help me find my OpenAI API key

Confirm that .env now appears in the Explorer sidebar alongside the cognee-env folder.

✔️ Awesome, I've got everything!

Great. Double check that your .env file is saved.

ⓧ I'd like to double check the full code

Here is the complete .env file you should have:

LLM_API_KEY=your-openai-api-key-here

Your environment is ready. Next up, you will store your first memory in Cognee and see how it retrieves knowledge using graph-based reasoning.

Store and Recall Your First Memory

Your environment is ready and Cognee is installed. Now it's time to give your knowledge graph memory system its first piece of knowledge and ask it a question.

Cognee's high-level API makes this surprisingly simple with two functions: remember() stores facts by running the full ingestion pipeline, and recall() retrieves answers using the best retrieval strategy.

In this step, get ready to:
  • Store a single fact using Cognee's remember() function and retrieve it with recall().
  • Query across multiple related facts to see cross-fact reasoning in action.
Use the high-level API to store and retrieve knowledge

Cognee's remember() function handles the entire ingestion pipeline in one call. It chunks your text, extracts entities, builds a knowledge graph, and enriches connections. The recall() function then auto-routes your query to the best retrieval strategy.

  • In VS Code, create a new file called first_memory.py inside your cognee-memory-project folder.
  • Paste in the following code:
import asyncio
import cognee


async def main():
    # Reset to a clean state
    await cognee.forget(everything=True)

    # Store a single fact in memory
    await cognee.remember(
        "Cognee is an open-source AI memory platform that builds knowledge graphs from documents."
    )

    # Recall from memory
    results = await cognee.recall(query_text="What is Cognee?")
    print("Single fact recall:")
    for result in results:
        print(result.text)


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

What does this code do?

  • cognee.forget(everything=True) permanently deletes all stored data and system state, giving you a guaranteed clean slate.
  • cognee.remember() ingests the text and runs the full pipeline (chunking, entity extraction, graph building, enrichment) in one call.
  • cognee.recall(query_text=...) searches your knowledge graph and returns contextual answers as a list of RecallResponse objects. Each has a .text attribute containing the answer.
  • The async/await pattern is required because Cognee's operations involve network calls to the OpenAI API and database writes that run asynchronously.
  • Save first_memory.py.
  • Run the script by entering this command in your terminal:
python first_memory.py

You should see Single fact recall: followed by a response describing what Cognee is.

Why does it take a moment?

Behind the scenes, remember() is calling the OpenAI API to extract entities, generate embeddings, and build graph relationships. The first run takes longer because it sets up local databases (SQLite, LanceDB, and Kuzu).

Seeing an API error?

  • Check that your .env file is in the cognee-memory-project folder (the same folder you're running the script from). Cognee loads environment variables from the working directory.
  • Verify that your LLM_API_KEY value is a valid OpenAI API key with credits available.
  • Make sure there are no extra spaces or quotes around your key in the .env file.

Help me debug my Cognee API error

Store multiple related facts and query across them

The real power of a knowledge graph becomes visible when you store multiple related facts and ask a question that requires connecting them. Let's extend the script to store three facts about different people, then query the relationship between them.

  • In first_memory.py, find this for loop inside the main() function:
    for result in results:
        print(result.text)
  • Add the following code below that for loop, above the if __name__ == "__main__": line:
    print("\n" + "=" * 60 + "\n")

    # Reset and store multiple related facts
    await cognee.forget(everything=True)

    await cognee.remember(
        [
            "Alice moved to Paris in 2010. She works as a software engineer.",
            "Bob lives in New York. He is a data scientist.",
            "Alice and Bob met at a conference in 2015.",
        ],
        self_improvement=False,
    )

    # Query across multiple memories
    results = await cognee.recall(query_text="What is the connection between Alice and Bob?")
    print("Multi-fact recall:")
    for result in results:
        print(result.text)

What does this code do?

  • cognee.forget(everything=True) resets state again so the new facts start on a clean graph.
  • cognee.remember() accepts a list of strings, treating each as a separate document to ingest into the knowledge graph.
  • self_improvement=False skips the enrichment step. This is useful when batch-loading data before a single enrichment pass later.
  • The recall() query asks about the connection between Alice and Bob. No single fact contains the full answer. Cognee must traverse its graph to connect information from all three facts.
  • Save first_memory.py.
  • Run the script again:
python first_memory.py

You should see both sections of output. The Multi-fact recall: section shows a contextual answer that connects Alice and Bob's relationship, their locations, and how they met.

Only seeing the single fact output?

  • Make sure your new code is inside the main() function (indented with 4 spaces) and above the if __name__ block.
  • Check that the if __name__ == "__main__": line is still at the bottom of the file with no indentation.

Help me fix my first_memory.py script

✔️ Awesome, I've got everything!

Great. Double check you've saved first_memory.py and confirmed both output sections appear in your terminal.

ⓧ I'd like to double check the full code

import asyncio
import cognee


async def main():
    # Reset to a clean state
    await cognee.forget(everything=True)

    # Store a single fact in memory
    await cognee.remember(
        "Cognee is an open-source AI memory platform that builds knowledge graphs from documents."
    )

    # Recall from memory
    results = await cognee.recall(query_text="What is Cognee?")
    print("Single fact recall:")
    for result in results:
        print(result.text)

    print("\n" + "=" * 60 + "\n")

    # Reset and store multiple related facts
    await cognee.forget(everything=True)

    await cognee.remember(
        [
            "Alice moved to Paris in 2010. She works as a software engineer.",
            "Bob lives in New York. He is a data scientist.",
            "Alice and Bob met at a conference in 2015.",
        ],
        self_improvement=False,
    )

    # Query across multiple memories
    results = await cognee.recall(query_text="What is the connection between Alice and Bob?")
    print("Multi-fact recall:")
    for result in results:
        print(result.text)


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

Your first memory is stored and retrievable. Next up, you'll go deeper with the ECL pipeline to see exactly how Cognee builds and enriches the knowledge graph under the hood.

Build the ECL Pipeline

In the previous step, you used remember() and recall() to store and retrieve knowledge with a single function call. That high-level API is convenient, but it hides the pipeline stages happening underneath.

Now you will use Cognee's lower-level ECL pipeline to control each stage individually. ECL stands for Extract, Cognify, and Load. You will add raw documents, build a knowledge graph from them, enrich the graph, and then query it with graph-aware search.

In this step, get ready to:
  • Ingest multiple company documents into a named dataset with cognee.add().
  • Build and enrich a knowledge graph with cognee.cognify() and cognee.memify().
  • Query the graph with SearchType.GRAPH_COMPLETION to find cross-document connections.
Ingest multiple documents with cognee.add()

The add() function ingests raw text into Cognee without building the graph. This gives you control over when graph construction happens.

  • Create a new file called ecl_pipeline.py in your cognee-memory-project folder.
  • Add the following code to set up the imports, define four company knowledge documents, and ingest them into a dataset:
import asyncio
import cognee
from cognee import SearchType


async def main():
    # Reset to clean state
    await cognee.forget(everything=True)

    # Define company knowledge documents
    documents = [
        "Project Alpha is a machine learning initiative led by Sarah Chen. It uses transformer models to analyze customer feedback and identify product improvement opportunities. The project started in January 2024 and is expected to deliver results by Q3 2024.",
        "Sarah Chen joined the company in 2020 as a senior data scientist. She has expertise in natural language processing and has published papers on sentiment analysis. She mentors junior team members including James Liu.",
        "James Liu is a junior data scientist who joined in 2023. He works on Project Alpha under Sarah Chen's mentorship. His focus is on data preprocessing and feature engineering for the transformer models.",
        "Project Beta is a computer vision project led by Marcus Johnson. It aims to automate quality inspection on the manufacturing floor using convolutional neural networks. The project depends on data pipelines built by the Project Alpha team.",
    ]

    # Step 1: Add data to Cognee
    for doc in documents:
        await cognee.add(doc, dataset_name="company_knowledge")

    print("Data ingested successfully.")

What does this code do?

cognee.forget(everything=True) wipes any previously stored data so you start fresh. This prevents conflicts with data from first_memory.py.

The documents list contains four text strings about a fictional company. Each describes a different person or project, with deliberate connections between them (Sarah mentors James, Project Beta depends on Project Alpha).

The for loop calls cognee.add() on each document, storing it in a dataset named company_knowledge. At this point the data is ingested but no graph has been built yet.

Unlike remember() which handles the entire pipeline in one call, add() only ingests. This separation lets you batch-load data before triggering a single graph-building pass.

  • Save ecl_pipeline.py.
  • Confirm the file appears in your VS Code Explorer sidebar under cognee-memory-project.
Build and enrich the knowledge graph

With your documents ingested, you can now build the knowledge graph and enrich it. These are two distinct stages that transform raw text into a structured, queryable graph.

  • In ecl_pipeline.py, add the following code directly below the print("Data ingested successfully.") line:
    # Step 2: Build the knowledge graph
    await cognee.cognify()
    print("Knowledge graph built.")

    # Step 3: Enrich the graph with memory algorithms
    await cognee.memify()
    print("Graph enriched with memory algorithms.")

What do cognify and memify do?

cognee.cognify() runs a six-stage pipeline: classify documents, check permissions, extract chunks, use the LLM to mine entities and relationships, generate summaries, and embed everything into the graph.

cognee.memify() enriches the graph after construction. It prunes stale nodes, strengthens frequently used connections, reweights edges based on usage signals, and adds derived facts. This improves graph quality for retrieval.

  • Save ecl_pipeline.py.
  • You won't see a change yet. The next section adds the search query that exercises the graph you just built.
Query with graph-aware search

The knowledge graph is built. Now you will query it with a question that requires connecting information across multiple documents.

  • In ecl_pipeline.py, add the following code directly below the print("Graph enriched with memory algorithms.") line:
    # Step 4: Search the knowledge graph
    print("\n" + "=" * 60)
    print("Cross-document query: How are Project Alpha and Project Beta connected?")
    print("=" * 60)

    results = await cognee.search(
        query_text="How are Project Alpha and Project Beta connected?",
        query_type=SearchType.GRAPH_COMPLETION,
    )

    for result in results:
        print(result)


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

What does GRAPH_COMPLETION do?

SearchType.GRAPH_COMPLETION tells Cognee to traverse the knowledge graph, collect relevant entities and relationships, and then use the LLM to generate a coherent answer grounded in graph context.

The query asks how Project Alpha and Project Beta are connected. No single document contains the full answer. The graph traces the relationship through shared data pipelines mentioned in the Project Beta document.

  • Save ecl_pipeline.py.
  • Run the script in your terminal:
python ecl_pipeline.py

This may take 30-60 seconds

The cognify() step makes multiple LLM calls to extract entities and relationships from your four documents. Wait for the script to complete without interrupting it.

  • You should see progress messages ("Data ingested successfully.", "Knowledge graph built.", "Graph enriched with memory algorithms.") followed by a cross-document answer explaining how Project Alpha and Project Beta are connected through shared team members and data pipelines.

Script hangs or throws an error?

  • If the script hangs for more than 2 minutes, check that your LLM_API_KEY in .env is a valid OpenAI API key with credits available. An invalid key causes silent retries.
  • If you see a connection or authentication error, verify your .env file is in the same directory you are running the script from. Cognee loads it from the current working directory.
  • If the search result is empty or generic, the LLM may not have extracted entities properly. Try running the script again. The graph extraction depends on LLM output quality.

Help me troubleshoot my Cognee ECL pipeline script. I'm running ecl_pipeline.py and seeing an error or hang during cognify.

✔️ Awesome, I've got everything!

Great. Double check you have saved ecl_pipeline.py and it ran successfully with a cross-document answer.

ⓧ I'd like to double check the full code

Here is the complete ecl_pipeline.py file:

import asyncio
import cognee
from cognee import SearchType


async def main():
    # Reset to clean state
    await cognee.forget(everything=True)

    # Define company knowledge documents
    documents = [
        "Project Alpha is a machine learning initiative led by Sarah Chen. It uses transformer models to analyze customer feedback and identify product improvement opportunities. The project started in January 2024 and is expected to deliver results by Q3 2024.",
        "Sarah Chen joined the company in 2020 as a senior data scientist. She has expertise in natural language processing and has published papers on sentiment analysis. She mentors junior team members including James Liu.",
        "James Liu is a junior data scientist who joined in 2023. He works on Project Alpha under Sarah Chen's mentorship. His focus is on data preprocessing and feature engineering for the transformer models.",
        "Project Beta is a computer vision project led by Marcus Johnson. It aims to automate quality inspection on the manufacturing floor using convolutional neural networks. The project depends on data pipelines built by the Project Alpha team.",
    ]

    # Step 1: Add data to Cognee
    for doc in documents:
        await cognee.add(doc, dataset_name="company_knowledge")

    print("Data ingested successfully.")

    # Step 2: Build the knowledge graph
    await cognee.cognify()
    print("Knowledge graph built.")

    # Step 3: Enrich the graph with memory algorithms
    await cognee.memify()
    print("Graph enriched with memory algorithms.")

    # Step 4: Search the knowledge graph
    print("\n" + "=" * 60)
    print("Cross-document query: How are Project Alpha and Project Beta connected?")
    print("=" * 60)

    results = await cognee.search(
        query_text="How are Project Alpha and Project Beta connected?",
        query_type=SearchType.GRAPH_COMPLETION,
    )

    for result in results:
        print(result)


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

You now have a working ECL pipeline that gives you full control over each stage of knowledge graph construction. Next up, you will compare different search modes to see how graph-aware retrieval outperforms traditional RAG.

Compare Search Modes

In the previous step, you built a knowledge graph from multiple company documents and queried it with GRAPH_COMPLETION. The result showed cross-document reasoning in action.

How do you know graph-based retrieval is actually better than traditional RAG? The best way to find out is to run the exact same query through different search strategies and compare the results side by side.

In this step, get ready to:
  • Create a comparison script that queries the same knowledge graph three ways.
  • Run the script and observe how each search mode produces different results.
Query with GRAPH_COMPLETION, RAG_COMPLETION, and CHUNKS

Cognee offers multiple search modes that each approach retrieval differently. GRAPH_COMPLETION traverses the knowledge graph and uses the LLM to reason over connected entities. RAG_COMPLETION retrieves the most similar text chunks and generates an answer from them. CHUNKS returns raw text segments matched semantically, with no synthesis.

  • Create a new file called compare_search.py in your cognee-memory-project folder and paste the following code:
import asyncio
import cognee
from cognee import SearchType


async def main():
    # This script assumes you have already run ecl_pipeline.py
    # and have data in your knowledge graph.

    query = "What is Sarah Chen's role and how does her work impact other projects?"

    # Mode 1: GRAPH_COMPLETION - Uses full graph context + LLM reasoning
    print("=" * 60)
    print("GRAPH_COMPLETION (Graph-aware Q&A)")
    print("=" * 60)
    results = await cognee.search(
        query_text=query,
        query_type=SearchType.GRAPH_COMPLETION,
    )
    for result in results:
        print(result)

    print("\n")

    # Mode 2: RAG_COMPLETION - Traditional RAG using document chunks
    print("=" * 60)
    print("RAG_COMPLETION (Traditional RAG)")
    print("=" * 60)
    results = await cognee.search(
        query_text=query,
        query_type=SearchType.RAG_COMPLETION,
    )
    for result in results:
        print(result)

    print("\n")

    # Mode 3: CHUNKS - Raw text segments matching semantically
    print("=" * 60)
    print("CHUNKS (Raw semantic matches)")
    print("=" * 60)
    results = await cognee.search(
        query_text=query,
        query_type=SearchType.CHUNKS,
    )
    for result in results:
        print(result)


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

What does this code do?

This script runs a single query against the knowledge graph you built in the previous step, but uses three different retrieval strategies.

  • query is a question that requires connecting information across multiple documents (Sarah's role, her mentorship, and her project's impact on Project Beta).
  • SearchType.GRAPH_COMPLETION walks the graph's entity relationships and feeds the connected context to the LLM for reasoning.
  • SearchType.RAG_COMPLETION retrieves the closest text chunks by embedding similarity and generates an answer from those chunks alone.
  • SearchType.CHUNKS returns raw text segments ranked by semantic similarity with no LLM synthesis at all.

Why does this script skip cognee.forget()?

Unlike earlier scripts, this one does not reset state. It reads from the knowledge graph that ecl_pipeline.py already built. The data persists locally in Cognee's default databases (SQLite, LanceDB, and Kuzu), so you can query it from a separate script without rebuilding.

  • Save compare_search.py.
  • Run the script by entering this command in your terminal:
python compare_search.py
  • You should see three clearly labeled sections of output in your terminal, one for each search mode.

Seeing an empty result or error?

  • Make sure you ran ecl_pipeline.py successfully before this script. The comparison script reads from the existing knowledge graph.
  • Check that your virtual environment is still activated. You should see (cognee-env) at the start of your terminal prompt.
  • Verify your .env file is in the same directory you are running from. Cognee loads it at import time.

Why is my compare_search.py returning empty or erroring out?

✔️ Awesome, I've got everything!

Great. Double check you have saved compare_search.py and it ran without errors.

ⓧ I'd like to double check the full code

import asyncio
import cognee
from cognee import SearchType


async def main():
    # This script assumes you have already run ecl_pipeline.py
    # and have data in your knowledge graph.

    query = "What is Sarah Chen's role and how does her work impact other projects?"

    # Mode 1: GRAPH_COMPLETION - Uses full graph context + LLM reasoning
    print("=" * 60)
    print("GRAPH_COMPLETION (Graph-aware Q&A)")
    print("=" * 60)
    results = await cognee.search(
        query_text=query,
        query_type=SearchType.GRAPH_COMPLETION,
    )
    for result in results:
        print(result)

    print("\n")

    # Mode 2: RAG_COMPLETION - Traditional RAG using document chunks
    print("=" * 60)
    print("RAG_COMPLETION (Traditional RAG)")
    print("=" * 60)
    results = await cognee.search(
        query_text=query,
        query_type=SearchType.RAG_COMPLETION,
    )
    for result in results:
        print(result)

    print("\n")

    # Mode 3: CHUNKS - Raw text segments matching semantically
    print("=" * 60)
    print("CHUNKS (Raw semantic matches)")
    print("=" * 60)
    results = await cognee.search(
        query_text=query,
        query_type=SearchType.CHUNKS,
    )
    for result in results:
        print(result)


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

Look at the three output sections your script produced. The differences reveal why each search mode exists.

What should I see in the output?

  • GRAPH_COMPLETION produces the richest answer. It traces Sarah's role as lead data scientist, her mentorship of James Liu, and how Project Alpha's data pipelines feed into Project Beta. The graph lets the LLM reason across multiple hops of connected entities.
  • RAG_COMPLETION returns a coherent answer, but grounded only in the chunks that are most semantically similar to the query. It may miss the connection to Project Beta because that fact lives in a separate document that does not directly mention Sarah.
  • CHUNKS returns raw text segments with no LLM synthesis. You see the original document passages ranked by similarity, not a composed answer.

Why does GRAPH_COMPLETION find connections the others miss?

Traditional RAG matches your query against text chunks using embedding similarity. If the answer spans multiple documents, it can only find it if those documents happen to be semantically close to the query.

GRAPH_COMPLETION works differently. During cognify(), Cognee extracted entities (Sarah Chen, Project Alpha, Project Beta) and relationships (leads, mentors, depends on) into a graph. At query time, it traverses those relationships to gather context that spans documents. The LLM then reasons over this connected context.

You now have concrete evidence of when graph-based retrieval outperforms traditional RAG. Next up, you will organize your knowledge base into separate node sets so you can scope searches to specific subsets of data.

Organize Knowledge with Node Sets

In Step 4, you compared three search modes and saw how graph-based retrieval connects information across documents.

Now imagine your knowledge graph contains hundreds of documents from different departments. Searching everything at once returns noisy results. In this step, you will organize your knowledge into tagged categories and query specific subsets using node sets.

In this step, get ready to:
  • Ingest documents tagged with node set categories.
  • Run filtered searches scoped to specific node sets.
  • Compare filtered vs. unfiltered results to see how scoping narrows answers.
Ingest categorized documents using node_set tags

When you call cognee.add(), you can pass a node_set parameter to tag documents with a category label. During cognify(), these tags become NodeSet nodes in the graph. At search time, you can filter results to only the subgraph connected to a specific node set.

  • Create a new file called organized_search.py in your cognee-memory-project folder.
  • Paste the following code to define two categories of documents and ingest them with node set tags:
import asyncio
import cognee
from cognee import SearchType
from cognee.modules.engine.models.node_set import NodeSet


async def main():
    # Reset to clean state
    await cognee.forget(everything=True)

    # Team member documents
    team_docs = [
        "Sarah Chen is the lead data scientist with expertise in NLP and transformer models. She joined in 2020.",
        "James Liu is a junior data scientist specializing in data preprocessing. He joined in 2023.",
        "Marcus Johnson leads the computer vision team and has 10 years of experience in manufacturing AI.",
    ]

    # Project documents
    project_docs = [
        "Project Alpha uses transformer models to analyze customer feedback. It started in January 2024 and targets Q3 2024 delivery.",
        "Project Beta automates quality inspection using convolutional neural networks on the manufacturing floor.",
        "Project Gamma is a recommendation engine that combines insights from both Alpha and Beta to suggest process improvements.",
    ]

    # Add team documents with 'team' node set
    for doc in team_docs:
        await cognee.add(doc, node_set=["team"])

    # Add project documents with 'projects' node set
    for doc in project_docs:
        await cognee.add(doc, node_set=["projects"])

    # Build and enrich the knowledge graph
    await cognee.cognify()
    await cognee.memify()
    print("Knowledge graph built with node sets.\n")


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

What does this code do?

This script separates documents into two logical groups. Team profiles go into the team node set, and project descriptions go into the projects node set.

The node_set=["team"] parameter on cognee.add() tells Cognee to associate each document with that category. When the graph is built with cognify(), these become dedicated nodes you can filter against at search time.

  • Save organized_search.py.

Why use node sets instead of separate databases?

Node sets let you keep all your knowledge in one graph while still querying specific slices. This is essential for production systems where you need access control, multi-tenant data separation, or domain-specific retrieval without duplicating infrastructure.

Perform filtered vs. unfiltered searches

With your categorized data ingested, you will now run the same query three ways: across all data, filtered to team profiles only, and filtered to projects only.

  • In organized_search.py, find this line inside your main() function:
    print("Knowledge graph built with node sets.\n")
  • Add the following search operations directly below that line, still inside the main() function:
    # Search across ALL data (no filter)
    print("=" * 60)
    print("UNFILTERED SEARCH (all node sets)")
    print("=" * 60)
    results = await cognee.search(
        query_text="Who works on machine learning?",
        query_type=SearchType.GRAPH_COMPLETION,
    )
    for result in results:
        print(result)

    print("\n")

    # Search ONLY within the 'team' node set
    print("=" * 60)
    print("FILTERED SEARCH (team node set only)")
    print("=" * 60)
    results = await cognee.search(
        query_text="Who works on machine learning?",
        query_type=SearchType.GRAPH_COMPLETION,
        node_type=NodeSet,
        node_name=["team"],
    )
    for result in results:
        print(result)

    print("\n")

    # Search ONLY within the 'projects' node set
    print("=" * 60)
    print("FILTERED SEARCH (projects node set only)")
    print("=" * 60)
    results = await cognee.search(
        query_text="What projects involve machine learning?",
        query_type=SearchType.GRAPH_COMPLETION,
        node_type=NodeSet,
        node_name=["projects"],
    )
    for result in results:
        print(result)

What does this code do?

The first search has no node set filter, so Cognee searches across all ingested data. The second and third searches use node_type=NodeSet and node_name=[...] to scope results to only the subgraph connected to that specific tag.

The NodeSet class is imported from cognee.modules.engine.models.node_set. This tells Cognee's search to project only the matching portion of the graph before running the query.

  • Save organized_search.py.

Seeing an import error for NodeSet?

Make sure your import line reads from cognee.modules.engine.models.node_set import NodeSet exactly. This is a nested module path and any typo will cause an ImportError.

Help me fix my NodeSet import error.

Compare filtered results

Now run the script to see how node set filtering narrows your search results.

  • Run the script by executing this command in your terminal:
python organized_search.py
  • Look at the three outputs in your terminal. The unfiltered search draws from both team profiles and project descriptions. The team-filtered search only references people and their skills. The projects-filtered search only references project descriptions.

Why do filtered results look different?

When you pass node_name=["team"], Cognee only traverses the portion of the knowledge graph connected to the team node set node. Entities and relationships from project documents are invisible to this search, even though they exist in the same graph.

This pattern is how production systems implement scoped retrieval. Different user roles, departments, or data domains each get their own node set, and searches stay within boundaries without needing separate databases.

All three searches returning the same results?

Confirm that you used node_set=["team"] and node_set=["projects"] in your cognee.add() calls. If you omitted the node_set parameter, all documents land in the default set and filtering has no effect.

Help me figure out why my node set filtering isn't working.

✔️ Awesome, I've got everything!

Great. Double check you have saved organized_search.py and that all three search outputs show different scoped results.

ⓧ I'd like to double check the full code

import asyncio
import cognee
from cognee import SearchType
from cognee.modules.engine.models.node_set import NodeSet


async def main():
    # Reset to clean state
    await cognee.forget(everything=True)

    # Team member documents
    team_docs = [
        "Sarah Chen is the lead data scientist with expertise in NLP and transformer models. She joined in 2020.",
        "James Liu is a junior data scientist specializing in data preprocessing. He joined in 2023.",
        "Marcus Johnson leads the computer vision team and has 10 years of experience in manufacturing AI.",
    ]

    # Project documents
    project_docs = [
        "Project Alpha uses transformer models to analyze customer feedback. It started in January 2024 and targets Q3 2024 delivery.",
        "Project Beta automates quality inspection using convolutional neural networks on the manufacturing floor.",
        "Project Gamma is a recommendation engine that combines insights from both Alpha and Beta to suggest process improvements.",
    ]

    # Add team documents with 'team' node set
    for doc in team_docs:
        await cognee.add(doc, node_set=["team"])

    # Add project documents with 'projects' node set
    for doc in project_docs:
        await cognee.add(doc, node_set=["projects"])

    # Build and enrich the knowledge graph
    await cognee.cognify()
    await cognee.memify()
    print("Knowledge graph built with node sets.\n")

    # Search across ALL data (no filter)
    print("=" * 60)
    print("UNFILTERED SEARCH (all node sets)")
    print("=" * 60)
    results = await cognee.search(
        query_text="Who works on machine learning?",
        query_type=SearchType.GRAPH_COMPLETION,
    )
    for result in results:
        print(result)

    print("\n")

    # Search ONLY within the 'team' node set
    print("=" * 60)
    print("FILTERED SEARCH (team node set only)")
    print("=" * 60)
    results = await cognee.search(
        query_text="Who works on machine learning?",
        query_type=SearchType.GRAPH_COMPLETION,
        node_type=NodeSet,
        node_name=["team"],
    )
    for result in results:
        print(result)

    print("\n")

    # Search ONLY within the 'projects' node set
    print("=" * 60)
    print("FILTERED SEARCH (projects node set only)")
    print("=" * 60)
    results = await cognee.search(
        query_text="What projects involve machine learning?",
        query_type=SearchType.GRAPH_COMPLETION,
        node_type=NodeSet,
        node_name=["projects"],
    )
    for result in results:
        print(result)


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

You now have a fully organized knowledge graph with scoped retrieval. Next up, you will push these skills further in the Secret Mission by building a multi-source research assistant.

Secret mission

Build a Multi-Source Research Assistant

Can you combine everything you've learned into a realistic research scenario? Build an assistant that answers questions requiring connections across research papers, meeting notes, and technical specs. No single source holds the full answer.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, so there are no ongoing cloud costs. The only charges were the OpenAI API calls you already made for embeddings and LLM reasoning, which are pay-per-use and already consumed.

Decide whether to keep your project files on disk, pause by deactivating the virtual environment, or delete everything entirely.

Resources you used:

  • The cognee-memory-project folder on your Desktop (contains all Python scripts, the .env file, and the cognee-env virtual environment).
  • Cognee's local databases (SQLite, LanceDB, Kuzu) stored inside the .cognee_system folder within your virtual environment's Cognee package directory.
  • OpenAI API credits consumed during the project (already used, not recoverable).

Keep everything running

No action needed. Choose this if you plan to keep experimenting with Cognee, adding more documents, or building on top of your knowledge graph.

  • Your project files stay on your Desktop.
  • Cognee's local databases persist between sessions, so you can run cognee.search() without rebuilding the graph each time.
  • No ongoing costs are incurred unless you run the scripts again (which triggers new OpenAI API calls).

Pause - I'll come back to this later

Deactivate the virtual environment to free up your terminal session. All files and databases stay on disk for when you return.

  • Deactivate the virtual environment by running:
deactivate
  • When you come back, navigate to your project folder and reactivate the environment with source cognee-env/bin/activate (macOS/Linux) or cognee-env\Scripts\activate (Windows).

Delete - I don't want to use this again

Remove all project files, the virtual environment, and Cognee's local databases. This frees up disk space completely.

  • First, deactivate the virtual environment if it is still active:
deactivate
  • Delete the entire project folder by running the appropriate command for your operating system.

macOS/Linux:

rm -rf ~/Desktop/cognee-memory-project

Windows (PowerShell):

Remove-Item -Recurse -Force ~\Desktop\cognee-memory-project
  • Confirm the folder is gone by running:
ls ~/Desktop

You should no longer see cognee-memory-project in the output.

What about Cognee's local databases?

Cognee stores its databases (SQLite, LanceDB, Kuzu) inside a .cognee_system folder within the Cognee package directory, which lives inside your cognee-env virtual environment. Deleting the project folder removes the virtual environment and all of Cognee's local data along with it.

Nice Work!

Nice Work!

You just built a knowledge graph memory system that ingests documents, extracts entities and relationships, and answers questions using graph-based reasoning. That is a production-grade pattern for giving AI agents persistent, structured memory.

You've learned how to:

  • Build a knowledge graph memory system using Cognee, from the high-level remember/recall API to the lower-level ECL pipeline (add, cognify, memify, search).
  • Compare three search modes (GRAPH_COMPLETION, RAG_COMPLETION, CHUNKS) and demonstrate why graph-based retrieval finds connections that traditional RAG misses.
  • Organize knowledge using node sets for filtered, scoped retrieval across categorized documents.
  • Secret Mission: Built a multi-source research assistant demonstrating cross-document reasoning across research papers, meeting notes, and technical specifications.

Ready to quiz yourself?