Cognee: Build Memory and Knowledge Graphs for Your AI Agents

# Cognee: Bangun Memory dan Knowledge Graph untuk AI Agent-mu Halo temen-temen, di tutorial kali ini aku mau ngajak kalian kenalan sama salah satu tools yang menurutku bakal makin penting seiring mak...

By Ruby Abdullah · · tutorial
cogneeknowledge-graphai-agentsragpython

Cognee: Build Memory and Knowledge Graphs for Your AI Agents

Hey folks, in this tutorial I want to introduce you to a tool that I think is going to get more and more important as people build more AI agents. It is called Cognee. If you have ever built a chatbot or an agent with an LLM, you have surely run into the classic problem: the model is great at conversation but has terrible memory. Every new session starts from scratch, it does not remember previous conversations, and it forgets the documents you fed it. Usually we solve this with RAG (Retrieval Augmented Generation) backed by a vector database. But plain vector search has a limitation: it only finds things that are semantically similar, and it does not understand the relationships between pieces of information.

This is exactly where Cognee comes in. Cognee is an open-source framework that turns your documents and conversations into a combined knowledge graph plus vector store that you can query. So you do not just get chunks of text that look similar, you also get the structure of relationships between entities. Think of it as a second brain for your agent that can remember facts and also understand how those facts connect to each other.

What I love is that Cognee follows an approach they call ECL: Extract, Cognify, Load. You extract data from various sources, cognify it (turn it into a graph plus embeddings), then load it into a queryable store. Everything runs on top of asyncio, so it is designed for async Python from the ground up. In this tutorial I will walk through installation, configuring an LLM provider, the core loop of add, cognify, and search, plus the concepts of data points and ontology, using it for RAG, and visualizing the graph. Let us get started.

Introduction: What Cognee Is and Why You Need It

Before we jump into code, let me explain the core concept so you do not get lost during the hands-on part. Cognee is essentially a memory engine for AI agents. But unlike simply storing chat history in an array, Cognee builds a structured representation of your knowledge.

There are three mental components you need to understand. First, there is the vector store, where your text is stored as embeddings for semantic search. Second, there is the graph store, where entities and the relationships between them are stored as nodes and edges. Third, there is the relational store, for metadata and document tracking. Cognee manages all three automatically for you, so you do not need to worry about setting up Neo4j, Qdrant, or Postgres manually at the start. By default everything uses local embedded storage.

Why is the graph important? Let me give you a simple example. Say you ingest a document about a company. Vector search can find the paragraph that mentions "founder" if you ask about the founder. But if you ask "who works at the company founded by person A", vector search struggles because that requires connecting several facts. A graph can answer questions like that because the relationship is stored explicitly as an edge. It is the combination of both that makes Cognee retrieval more powerful than plain RAG.

Cognee fits great for cases like agents that need long-term memory, RAG systems that need multi-hop reasoning, internal company knowledge bases, or personal assistants that remember your preferences over time. Alright, let us install it first.

Installation

Installing Cognee is super easy, just one pip command. I recommend you create a virtual environment first to keep things tidy.

python -m venv venv

source venv/bin/activate # on Windows: venv\Scripts\activate

pip install cognee

Cognee needs Python 3.10 or higher, folks, so make sure your Python version is recent enough. You can check with:

python --version

After installation, the next most important step is configuring the LLM provider. Cognee needs an LLM for two main things: first to extract entities and relationships during the cognify process, and second to generate embeddings (if you use embeddings from the same provider). By default Cognee uses OpenAI, so the fastest way is to set the API key via an environment variable.

export LLMAPIKEY="sk-your-openai-api-key"

If you want to be more explicit or use a different provider, Cognee provides a config module you can call from Python. I prefer this approach because it is clearer and easier to track in your codebase.

import cognee

Configure the LLM provider

cognee.config.setllmprovider("openai")

cognee.config.setllmmodel("gpt-4o-mini")

cognee.config.setllmapikey("sk-your-openai-api-key")

If you use a different provider, for example a local model through Ollama or an OpenAI-compatible endpoint, you can set the endpoint too. This is super useful if you want to run everything offline or use a self-hosted model.

import cognee

cognee.config.setllmprovider("ollama")

cognee.config.setllmmodel("llama3.1:8b")

cognee.config.setllmendpoint("http://localhost:11434/v1")

cognee.config.setllmapikey("ollama") # placeholder, Ollama does not need a real key

Besides the LLM, you can also control where data is stored via environment variables or config. By default Cognee creates a folder in the system directory to store the database. If you want to control the location, just set the DATAROOTDIRECTORY and SYSTEMROOTDIRECTORY variables. For now, do not worry about it, the defaults work fine. Now let us move to the most exciting part, the basic usage.

This is the heart of Cognee. There are three functions you will use over and over: cognee.add() to ingest data, cognee.cognify() to build the graph, and cognee.search() to retrieve. Because Cognee runs on asyncio, all these functions are coroutines, so you have to await them and run them inside an event loop. Let me show you the most minimal example first so you get a feel for the flow.

import asyncio

import cognee

async def main():

# 1. Add raw data to Cognee

text = """

Ruby Abdullah is an AI engineer from Indonesia.

He founded rubythalib.ai, an AI education platform.

The platform focuses on teaching machine learning and LLMs.

"""

await cognee.add(text)

# 2. Process the data into a knowledge graph + embeddings

await cognee.cognify()

# 3. Query the knowledge you built

results = await cognee.search("What is rubythalib.ai?")

for result in results:

print(result)

if name == "main":

asyncio.run(main())

Let us break down what happens. When you call cognee.add(text), Cognee stores that raw text into the relational store as a document. At this stage there is no graph yet, no embeddings yet, it is just saved as raw data. You can call add many times to pile up documents before cognify.

Now, the magic happens in cognee.cognify(). This function takes all the data you have added, then runs the pipeline: it chunks the text, calls the LLM to extract entities and relationships, builds embeddings for each chunk and entity, then stores everything into the graph store and vector store. This process is the most time-consuming and the most token-hungry, so it is normal for it to take a while on large documents. Once cognify finishes, your knowledge is structured and ready to query.

Finally, cognee.search() is what you use to ask questions. Cognee has several search types that behave differently, and this is really important to understand. I will cover that in a moment. What matters is that search leverages both the graph and the vector store to give relevant answers.

Adding Multiple Data Sources

You are not limited to strings. Cognee can accept a list of strings, a path to a file, even a folder. This helps a lot when you have many documents.

import asyncio

import cognee

async def ingestmultiple():

documents = [

"Python is a popular programming language for data science.",

"PyTorch and TensorFlow are the main deep learning frameworks.",

"Hugging Face provides thousands of pre-trained models.",

]

# Add a list of documents at once

await cognee.add(documents)

# You can also add a file from a path

# await cognee.add("file:///home/rubya/docs/report.pdf")

await cognee.cognify()

print("Done building the graph!")

if name == "main":

asyncio.run(ingestmultiple())

Understanding Search Types

This is the part that often confuses people, so let me explain it slowly. Cognee provides several SearchType options you pick via the querytype parameter. Each one has a different character.

import asyncio

import cognee

from cognee import SearchType

async def searchexamples():

# GRAPHCOMPLETION: natural language answer based on the graph (RAG-like)

answer = await cognee.search(

"Who founded rubythalib.ai?",

querytype=SearchType.GRAPHCOMPLETION,

)

print("Graph completion:", answer)

# RAGCOMPLETION: answer based on classic vector chunks

rag = await cognee.search(

"Tell me about the AI education platform",

querytype=SearchType.RAGCOMPLETION,

)

print("RAG completion:", rag)

# CHUNKS: return the relevant raw text chunks

chunks = await cognee.search(

"machine learning",

querytype=SearchType.CHUNKS,

)

print("Chunks:", chunks)

# INSIGHTS: return relationships between entities from the graph

insights = await cognee.search(

"rubythalib.ai",

querytype=SearchType.INSIGHTS,

)

print("Insights:", insights)

if name == "main":

asyncio.run(searchexamples())

So here is the difference, folks. GRAPHCOMPLETION is the one I use most when I want a natural answer that leverages the graph structure. It pulls context from the graph then hands it to the LLM to compose an answer. RAGCOMPLETION is similar but the context comes only from vector chunks, so this is the equivalent of traditional RAG. CHUNKS gives you raw text pieces without LLM processing, great if you want to control the prompt yourself. And INSIGHTS is special, it gives you relationships from the graph in the form of entity triplets, great for exploring or debugging your graph.

Advanced Usage: Data Points, Ontology, RAG, and Visualization

Alright now we level up. After understanding the basic loop, there are several advanced features that make Cognee far more powerful for serious use cases.

Data Points: Giving Structure to Knowledge

By default, Cognee automatically extracts entities from your text. But sometimes you already have structured data and want to feed it into the graph with a schema you define yourself. This is where the DataPoint concept comes in. A DataPoint is basically a Pydantic model that defines a node in your graph, complete with its fields and relationships.

import asyncio

import cognee

from cognee.lowlevel import DataPoint

class Instructor(DataPoint):

name: str

expertise: str

# metadata tells Cognee which fields to embed

metadata: dict = {"indexfields": ["name", "expertise"]}

class Course(DataPoint):

title: str

description: str

taughtby: Instructor # this becomes a relationship (edge) in the graph

metadata: dict = {"indexfields": ["title", "description"]}

async def addstructured():

ruby = Instructor(

name="Ruby Abdullah",

expertise="AI Engineering and LLMs",

)

course = Course(

title="LLM Fundamentals for Beginners",

description="Learn the basics of large language models from scratch.",

taughtby=ruby,

)

# Add the data point directly to the graph engine

from cognee.lowlevel import setup

from cognee.tasks.storage import adddatapoints

await setup()

await adddatapoints([course])

print("Structured data point added successfully!")

if name == "main":

asyncio.run(addstructured())

The cool thing about this approach is that because taughtby is typed as Instructor (which is also a DataPoint), Cognee automatically creates an edge between the Course node and the Instructor node in the graph. The indexfields in metadata tell Cognee which fields need to be embedded for vector search. So you get full control over the graph structure without relying on LLM extraction that is sometimes inconsistent.

Ontology: Giving Domain Rules

An ontology is like a blueprint that defines what entity types and relationships are valid in your domain. If you have a specific domain, say medical or legal, an ontology helps Cognee extract entities more accurately and consistently. Cognee can accept an ontology file in OWL/RDF format (XML).

import asyncio

import cognee

async def cognifywithontology():

await cognee.add("Aspirin is a drug used to relieve pain.")

# Point cognify to your domain ontology file

await cognee.cognify(ontologyfilepath="ontologies/medical.owl")

results = await cognee.search("Which drug relieves pain?")

print(results)

if name == "main":

asyncio.run(cognifywithontology())

With an ontology, entity extraction becomes more targeted. For example Cognee will know that "Aspirin" is an instance of the "Drug" class and not just a generic entity. This keeps your graph cleaner and makes queries more precise. You are not required to use an ontology at the start, but once your domain gets more complex, this becomes a lifesaver.

Using Cognee for Smarter RAG

Now this is the most sought-after use case. How do you build a RAG pipeline with Cognee. The advantage over vanilla RAG is that your retrieval leverages the graph, so it can answer questions that require connecting several facts. Let me show you an example of building a mini knowledge base then querying it for RAG.

import asyncio

import cognee

from cognee import SearchType

async def buildragsystem():

# Reset first to keep it clean (optional, be careful in production)

await cognee.prune.prunedata()

await cognee.prune.prunesystem(metadata=True)

knowledge = [

"rubythalib.ai offers AI Engineering and Data Science classes.",

"The AI Engineering class is taught by Ruby Abdullah.",

"Ruby Abdullah has experience building production LLM systems.",

"Data Science covers statistics, Python, and machine learning.",

]

await cognee.add(knowledge)

await cognee.cognify()

# Multi-hop query: needs to connect who the instructor is + their experience

question = "What experience does the AI Engineering instructor have?"

answer = await cognee.search(

question,

querytype=SearchType.GRAPHCOMPLETION,

)

print("Question:", question)

print("Answer:", answer)

if name == "main":

asyncio.run(buildragsystem())

Notice the question: "What experience does the AI Engineering instructor have?". To answer this, the system first has to know who teaches the AI Engineering class (Ruby Abdullah), then look up Ruby Abdullah's experience. That is two hops in the graph. A regular RAG that only relies on similarity search would struggle, but Cognee can handle it because the relationship is explicit. This is what I mean when I say the graph makes retrieval smarter.

One important note, I used cognee.prune.prunedata() and prunesystem() at the start to reset. Be careful, this deletes all data. It is really useful during development so each run starts clean, but never let it get called in production.

Separating Data with Datasets

If you have many distinct contexts, for example per-user or per-project data, you can separate them using datasets. This helps isolation so one user's retrieval does not mix with another user's.

import asyncio

import cognee

async def usedatasets():

await cognee.add(

"User A preference: likes detailed explanations with code examples.",

datasetname="usera",

)

await cognee.add(

"User B preference: likes short and to-the-point explanations.",

datasetname="userb",

)

await cognee.cognify(datasets=["usera", "userb"])

# Search only within a specific dataset

results = await cognee.search(

"What is the user's preference?",

datasets=["usera"],

)

print(results)

if name == "main":

asyncio.run(usedatasets())

Visualizing the Knowledge Graph

This is my favorite feature for debugging. Cognee can render your knowledge graph into an interactive HTML file you can open in a browser. So you can visually see the nodes and edges that formed, and make sure the extraction makes sense.

import asyncio

import cognee

async def visualize():

await cognee.add(

"Ruby Abdullah teaches at rubythalib.ai. "

"The platform teaches machine learning and LLMs."

)

await cognee.cognify()

# Render the graph to an interactive HTML file

await cognee.visualizegraph("./knowledgegraph.html")

print("Graph saved to knowledgegraph.html, open it in a browser!")

if name == "main":

asyncio.run(visualize())

Once you open knowledgegraph.html, you will see nodes like "Ruby Abdullah", "rubythalib.ai", "machine learning" connected by edges that show their relationships. Trust me, seeing the graph visually helps a ton to understand why your search returns certain results. If there is an entity that should connect but does not, you can immediately notice it and fix the data.

Best Practices

After using Cognee in a few projects, there are several things I learned that I want to share so you do not fall into the same holes.

First, about cognify cost and performance. Remember that cognify() calls the LLM for each chunk, so the more data, the more tokens used. For development, use a cheap model like gpt-4o-mini first. Only when you need better extraction quality should you upgrade to a smarter model. Do not just cognify thousands of documents with an expensive model without testing on a small sample first.

Second, batch your add calls before cognify. Instead of calling add then cognify back and forth for each document, add all documents first then cognify once. This is more efficient because Cognee can process in batch and cross-document relationships get detected better.

Third, pick the right SearchType. I often see people use GRAPHCOMPLETION for everything when sometimes CHUNKS is enough. If you only need text chunks to process in your own system, CHUNKS is faster and less token-hungry. Use GRAPHCOMPLETION only when you actually need a natural answer that leverages graph reasoning.

Fourth, leverage DataPoint for structured data. If you already have clean data from a database, do not turn it into loose text and let the LLM re-extract it. That wastes tokens and the result can be inconsistent. Just build a DataPoint with a clear schema, so your graph is precise and deterministic.

Fifth, be careful with async context. Because all Cognee functions are coroutines, make sure you always run them inside a proper event loop. If you integrate into a web framework like FastAPI, you can just await inside the route handler because FastAPI is already async. But in a plain script, wrap it in asyncio.run(). Do not mix sync and async code carelessly because it can create errors that are hard to debug.

Sixth, for production, consider using a proper backend. By default Cognee uses embedded storage that is fine for prototypes. But for production with large data, you can configure Cognee to use a real vector database like Qdrant or a graph database like Neo4j. Cognee supports this via config, so migrating from prototype to production is relatively smooth.

import cognee

Example production backend configuration

cognee.config.setvectordbprovider("qdrant")

cognee.config.setgraphdbprovider("neo4j")

Seventh, monitor and validate your graph periodically. Entity extraction with an LLM is not always perfect. Every now and then render the graph with visualizegraph() and check whether the entities and relationships make sense. If you find a lot of noise or wrong relationships, that is a signal to improve your input data quality or add an ontology.

Eighth, manage your API keys securely. Do not hardcode your API key in the code like the examples above, that is only for illustration. In a real project, use environment variables or a secret manager. I usually put it in a .env file then load it with python-dotenv.

import os

from dotenv import loaddotenv

import cognee

loaddotenv()

cognee.config.setllmapikey(os.getenv("LLMAPIKEY"))

Conclusion

Alright folks, we have wandered pretty far into Cognee. We started from the core concept that Cognee turns documents and conversations into a combined knowledge graph plus vector store, then we went into the core loop of add, cognify, and search. We also covered the search types that behave differently, the DataPoint concept for structured data, ontology for specific domains, using Cognee for RAG that can do multi-hop reasoning, separating data with datasets, and visualizing the graph for debugging.

What makes Cognee interesting in my opinion is that it bridges the gap between traditional RAG that only relies on vector similarity and the need for more complex reasoning. With a graph, your agent can not only find similar information but also connect related facts. This is really important if you build an agent you genuinely want to rely on to answer complex questions.

My advice, start small. Try the most minimal add, cognify, search loop, feel how it works, then render the graph so you understand what is happening under the hood. Once you are comfortable, then explore DataPoint and ontology for more precise control. Because everything runs on asyncio, get used to writing clean async code from the start so it is easy to integrate into your application later.

Cognee is still evolving and the community is active, so it is really worth keeping an eye on. If you are building an AI agent that needs long-term memory and reasoning beyond keyword matching, Cognee absolutely deserves a spot in your toolkit. Give it a try, and I hope this tutorial helps you build smarter agents. See you in the next tutorial, folks!

Related Articles

DSPy: Stop Hand-Tuning Prompts, Let the Compiler Optimize Them

DSPy: Berhenti Ngoprek Prompt Manual, Biarkan Compiler yang Optimasi Halo temen-temen, kali ini aku mau ngenalin satu li...

Complete Pinecone Tutorial: Vector Database for AI and Semantic Search

Tutorial Lengkap Pinecone: Vector Database untuk AI dan Semantic Search Pinecone adalah managed vector database yang dir...

LangChain Tutorial: The Most Popular Framework for Building LLM Applications

Tutorial LangChain: Framework Paling Populer untuk Membangun Aplikasi LLM LangChain adalah framework open-source yang di...

ColBERT & RAGatouille Tutorial: Late-Interaction Retrieval for RAG

ColBERT & RAGatouille: Retrieval Late-Interaction untuk RAG yang Lebih Baik Sebagian besar sistem RAG mengandalkan dense...