Building an Agentic RAG Pipeline: Persistent Graph-Based Memory for AI Agents
Building an Agentic RAG Pipeline with Cognee: Persistent Graph-Based Memory for AI Agents
How to give your AI agents memory that actually remembers
Standard RAG systems find and summarize facts, but they don’t really think. They retrieve chunks of text based on semantic similarity, but they have no understanding of how concepts relate to each other across your documents. In this tutorial, we’ll build an agentic RAG pipeline using Cognee an open-source memory engine that combines vector search with knowledge graphs to give your agents persistent, structural memory.
By the end of this article, you’ll understand how to replace manual chunking and vector store setup with Cognee’s simple ECL (Extract, Cognify, Load) pipeline, and integrate it into an agent framework for intelligent, memory-aware assistants.
Table of Contents
Phase 1: Understanding the Problem
Before we write any code, let’s understand why we need Cognee in the first place.
Why Standard RAG Falls Short for Agents
Traditional RAG (Retrieval-Augmented Generation) works like this: you chunk your documents, create vector embeddings, store them in a vector database, and retrieve the most semantically similar chunks when a user asks a question.
This approach has a fundamental limitation: it has no structural memory.
Ask a standard RAG system “What did Frodo do?” and it might find the right paragraph. But ask it “Who is the uncle of the person who carried the Ring?” and it struggles — because answering that requires connecting two separate facts across your knowledge base. Standard RAG retrieves information in isolation; it cannot perform multi-hop reasoning.
For AI agents, this is a critical problem. Agents need to remember context, link facts together, and build on past interactions to act intelligently. They need memory that understands relationships, not just similarity.
What Cognee Solves
Cognee is an open-source AI memory engine that transforms your documents into a persistent, graph-based memory layer. Instead of just storing text chunks as vectors, Cognee:
The result is that your agent can follow chains of relationships across your entire knowledge base — what we call multi-hop reasoning.
Standard RAG vs Cognee: The Key Difference
With standard RAG, you store vectors and retrieve by similarity. With Cognee, you store a graph of interconnected knowledge and retrieve by both similarity and relationships.
Standard RAG gives you: “Here are the 5 most similar text chunks to your question.”
Cognee gives you: “Here are the relevant facts, and here’s how they connect to each other.”
The ECL Paradigm
Cognee replaces the traditional “chunk, embed, store” pipeline with ECL:
This is what “memory in 6 lines of code” actually looks like, and we’ll explore each step in detail.
Phase 2: Setting Up Cognee as the Knowledge Core
Now let’s get our hands dirty. In this phase, we’ll install Cognee, add documents, and build our first knowledge graph.
Prerequisites
Before starting, make sure you have:
Step 1: Install Cognee
Open your terminal and install Cognee using pip:
pip install cognee
You can also use poetry or uv if you prefer those package managers.
Step 2: Configure Your Environment
Cognee needs access to an LLM for extracting entities and relationships. Set your API key as an environment variable:
export LLM_API_KEY="your-openai-api-key-here"
Alternatively, you can set it directly in Python:
import os
os.environ["LLM_API_KEY"] = "your-openai-api-key-here"
By default, Cognee uses:
All of these run locally with no additional setup required.
Step 3: Add Documents to Cognee
Let’s start with a simple example. We’ll add some text about a fictional company and see how Cognee processes it.
import cognee
import asyncio
async def main():
# First, let's start with a clean slate
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
# Now let's add our documents
company_info = """
TechNova Inc. was founded in 2019 by Sarah Chen and Marcus Williams.
Sarah Chen serves as the CEO and has a background in artificial intelligence.
Marcus Williams is the CTO and previously worked at Google on cloud infrastructure.
The company's flagship product is NovaAI, an enterprise automation platform.
NovaAI uses machine learning to automate repetitive business processes.
In 2023, TechNova raised $50 million in Series B funding led by Sequoia Capital.
The company is headquartered in San Francisco and has 200 employees.
"""
# Add the text to Cognee
await cognee.add(company_info)
print("Documents added successfully!")
if __name__ == "__main__":
asyncio.run(main())
The cognee.add() function accepts various input types: strings, file paths, lists of documents, and even URLs.
Step 4: Build the Knowledge Graph with Cognify
This is where the magic happens. The cognify() function uses your configured LLM to extract entities, relationships, and summaries from your documents, then assembles them into a knowledge graph.
import cognee
import asyncio
async def main():
# Reset and add data (same as before)
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
company_info = """
TechNova Inc. was founded in 2019 by Sarah Chen and Marcus Williams.
Sarah Chen serves as the CEO and has a background in artificial intelligence.
Marcus Williams is the CTO and previously worked at Google on cloud infrastructure.
The company's flagship product is NovaAI, an enterprise automation platform.
NovaAI uses machine learning to automate repetitive business processes.
In 2023, TechNova raised $50 million in Series B funding led by Sequoia Capital.
The company is headquartered in San Francisco and has 200 employees.
"""
await cognee.add(company_info)
# Now build the knowledge graph
await cognee.cognify()
print("Knowledge graph created!")
if __name__ == "__main__":
asyncio.run(main())
When cognify() runs, it processes each chunk of your document through the LLM, which identifies:
All of this gets stored in the graph database with vector embeddings indexed for semantic search.
Step 5: Visualize the Knowledge Graph
Cognee includes a built-in visualization function so you can see what graph it built.
import cognee
import asyncio
from pathlib import Path
async def main():
# Reset, add data, and cognify (same as before)
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
company_info = """
TechNova Inc. was founded in 2019 by Sarah Chen and Marcus Williams.
Sarah Chen serves as the CEO and has a background in artificial intelligence.
Marcus Williams is the CTO and previously worked at Google on cloud infrastructure.
The company's flagship product is NovaAI, an enterprise automation platform.
NovaAI uses machine learning to automate repetitive business processes.
In 2023, TechNova raised $50 million in Series B funding led by Sequoia Capital.
The company is headquartered in San Francisco and has 200 employees.
"""
await cognee.add(company_info)
await cognee.cognify()
# Visualize the graph
graph_path = str(Path("./graph_visualization.html").resolve())
await cognee.visualize_graph(graph_path)
print(f"Graph visualization saved to: {graph_path}")
if __name__ == "__main__":
asyncio.run(main())
Open the generated HTML file in your browser to see an interactive visualization of your knowledge graph.
Step 6: Search the Knowledge Graph
Now we can query our knowledge graph. Cognee’s search combines vector similarity with graph traversal.
import cognee
import asyncio
async def main():
# Assuming we've already run add() and cognify()
results = await cognee.search("Who founded TechNova?")
print("Search Results:")
for result in results:
print(result)
if __name__ == "__main__":
asyncio.run(main())
The search function doesn’t just find similar text — it traverses the knowledge graph to find connected information.
What We’ve Accomplished So Far
Compare what we just did to traditional RAG setup:
Traditional RAG requires: loading documents, chunking them, initializing an embedding model, creating embeddings, setting up a vector database, indexing, and building retrieval functions.
Cognee requires: add(), cognify(), search().
The knowledge graph is built automatically, relationships are extracted by the LLM, and search combines semantic and structural retrieval.
Phase 3: Building Agent Tools with Cognee
Now we need to transform Cognee’s capabilities into tools that an AI agent can use. In agentic RAG, the agent decides when and how to use tools based on the user’s query.
Understanding Tools in Agentic Systems
A tool is simply a function that an agent can call. The agent reads the tool’s name and description, then decides whether to use it based on the current task. For Cognee, we need:
Creating the Cognee Search Tool
We’ll use LangChain’s @tool decorator to make our functions discoverable by agents.
from langchain_core.tools import tool
import cognee
@tool
async def cognee_search_tool(query: str) -> str:
"""
Search the knowledge graph for information.
Use this tool when you need to find facts, relationships,
or context from stored documents.
Args:
query: The question or search query in natural language
Returns:
Relevant information from the knowledge graph
"""
try:
results = await cognee.search(query)
if not results:
return "No relevant information found in memory."
# Format results as a readable string
formatted_results = []
for i, result in enumerate(results, 1):
formatted_results.append(f"{i}. {result}")
return "\n".join(formatted_results)
except Exception as e:
return f"Error searching memory: {str(e)}"
The docstring is critically important here. The agent’s LLM reads this description to understand when to use the tool.
Creating the Cognee Add Tool
This tool allows the agent to store new information during a conversation.
from langchain_core.tools import tool
import cognee
@tool
async def cognee_add_tool(content: str) -> str:
"""
Add new information to the knowledge graph memory.
Use this tool when the user provides important facts, preferences,
or information that should be remembered for future conversations.
Args:
content: The information to store in memory
Returns:
Confirmation that the information was stored
"""
try:
# Add the content to Cognee
await cognee.add(content)
# Process it into the knowledge graph
await cognee.cognify()
return f"Successfully stored in memory: {content[:100]}..."
except Exception as e:
return f"Error storing in memory: {str(e)}"
The Complete Tools Module
Let’s put both tools together in a reusable module:
# cognee_tools.py
from langchain_core.tools import tool
import cognee
@tool
async def cognee_search_tool(query: str) -> str:
"""
Search the knowledge graph for information.
Use this tool when you need to find facts, relationships,
or context from stored documents and past conversations.
Args:
query: The question or search query in natural language
Returns:
Relevant information from the knowledge graph
"""
try:
results = await cognee.search(query)
if not results:
return "No relevant information found in memory."
formatted_results = []
for i, result in enumerate(results, 1):
formatted_results.append(f"{i}. {result}")
return "\n".join(formatted_results)
except Exception as e:
return f"Error searching memory: {str(e)}"
@tool
async def cognee_add_tool(content: str) -> str:
"""
Add new information to the knowledge graph memory.
Use this tool when the user provides important facts, preferences,
or information that should be remembered for future conversations.
Args:
content: The information to store in memory
Returns:
Confirmation that the information was stored
"""
try:
await cognee.add(content)
await cognee.cognify()
return f"Successfully stored in memory: {content[:100]}..."
except Exception as e:
return f"Error storing in memory: {str(e)}"
def get_cognee_tools():
"""Returns the list of Cognee tools for use with an agent."""
return [cognee_search_tool, cognee_add_tool]
Why This Approach Works
These two simple tools give your agent powerful capabilities:
The search tool replaces what would normally be a complex retrieval pipeline. Cognee handles embedding, searching, and graph traversal internally.
The add tool gives your agent the ability to learn. If a user mentions “My name is Ali and I work in AI education,” the agent can store this, and future searches will return this context. This is persistent memory that survives across sessions.
Phase 4: Integrating with an Agent Framework (LangGraph)
Now we’ll bring everything together by building an agent using LangGraph. This agent will use Cognee tools for persistent memory.
Why LangGraph?
LangGraph is a framework for building stateful, multi-step agent applications. It provides:
Combined with Cognee’s memory, we get an agent that can plan, remember, and reason.
Installing Dependencies
First, install the required packages:
pip install langgraph langchain-openai cognee
Setting Up the Agent
We’ll use LangGraph’s create_react_agent function, which creates a ReAct (Reasoning + Acting) agent that can use tools.
# agent.py
import os
import asyncio
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from langchain_core.messages import HumanMessage
import cognee
# Import our Cognee tools
from cognee_tools import get_cognee_tools
# Configure API keys
os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
os.environ["LLM_API_KEY"] = "your-openai-api-key" # For Cognee
async def setup_agent():
"""Create and configure the agent with Cognee memory tools."""
# Initialize the LLM
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
# Get Cognee tools
tools = get_cognee_tools()
# Create the agent
agent = create_react_agent(
model=llm,
tools=tools,
)
return agent
async def main():
# Initialize Cognee (clean start for demo)
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
# Pre-load some knowledge
initial_knowledge = """
The company TechNova Inc. was founded in 2019 by Sarah Chen and Marcus Williams.
Sarah Chen is the CEO with expertise in artificial intelligence.
Marcus Williams is the CTO who previously worked at Google.
Their main product is NovaAI, an enterprise automation platform.
TechNova raised $50 million in Series B funding from Sequoia Capital in 2023.
The company is based in San Francisco with 200 employees.
"""
await cognee.add(initial_knowledge)
await cognee.cognify()
print("Knowledge base initialized!")
# Create the agent
agent = await setup_agent()
# Test conversation
print("\n--- Starting Agent Conversation ---\n")
# First query: Search existing knowledge
response = await agent.ainvoke({
"messages": [HumanMessage(content="Who is the CEO of TechNova?")]
})
print(f"User: Who is the CEO of TechNova?")
print(f"Agent: {response['messages'][-1].content}\n")
# Second query: Multi-hop reasoning
response = await agent.ainvoke({
"messages": [HumanMessage(content="What is the background of the person who founded TechNova and serves as CEO?")]
})
print(f"User: What is the background of the person who founded TechNova and serves as CEO?")
print(f"Agent: {response['messages'][-1].content}\n")
# Third query: Add new memory
response = await agent.ainvoke({
"messages": [HumanMessage(content="Remember this: TechNova just launched a new product called NovaChat, an AI chatbot for customer service.")]
})
print(f"User: Remember this: TechNova just launched a new product called NovaChat...")
print(f"Agent: {response['messages'][-1].content}\n")
# Fourth query: Search for newly added memory
response = await agent.ainvoke({
"messages": [HumanMessage(content="What products does TechNova offer?")]
})
print(f"User: What products does TechNova offer?")
print(f"Agent: {response['messages'][-1].content}\n")
if __name__ == "__main__":
asyncio.run(main())
asyncio.run(main())
Understanding the Agent Flow
Let's trace through what happens when the user asks "Who is the CEO of TechNova?":
For the multi-hop question about the CEO's background, Cognee's graph structure shines. It can traverse: TechNova → FOUNDED_BY → Sarah Chen → HAS_BACKGROUND → artificial intelligence.
Demonstrating Cross-Session Persistence
One of Cognee's key advantages is that memory persists across sessions:
# session_demo.py
import os
import asyncio
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent
from langchain_core.messages import HumanMessage
import cognee
from cognee_tools import get_cognee_tools
os.environ["OPENAI_API_KEY"] = "your-openai-api-key"
os.environ["LLM_API_KEY"] = "your-openai-api-key"
async def create_agent():
llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
tools = get_cognee_tools()
return create_react_agent(model=llm, tools=tools)
async def session_one():
"""First session: Add information"""
print("=== SESSION 1 ===")
agent = await create_agent()
# Store user preferences
response = await agent.ainvoke({
"messages": [HumanMessage(
content="Remember: My name is Ali, I teach AI courses, and I prefer practical examples over theory."
)]
})
print(f"Stored: {response['messages'][-1].content}")
async def session_two():
"""Second session: Retrieve information from previous session"""
print("\n=== SESSION 2 (New Agent Instance) ===")
# Create a completely new agent instance
agent = await create_agent()
# Query information from previous session
response = await agent.ainvoke({
"messages": [HumanMessage(content="What do you know about me?")]
})
print(f"Retrieved: {response['messages'][-1].content}")
async def main():
# Note: We do NOT call prune here to preserve data across sessions
await session_one()
await session_two()
if __name__ == "__main__":
asyncio.run(main())
In Session 2, we create a completely new agent instance with no conversation history. Yet it can still recall information from Session 1 because Cognee stores memory in a persistent graph database.
Using the Official LangGraph Integration
Cognee provides an official integration package that simplifies this setup:
pip install cognee-integration-langgraph
With this package, you can use pre-built sessionized tools:
from langgraph.prebuilt import create_react_agent
from cognee_integration_langgraph import get_sessionized_cognee_tools
from langchain_core.messages import HumanMessage
async def main():
# Get memory tools with session isolation
add_tool, search_tool = get_sessionized_cognee_tools("user-123")
# Create agent with memory
agent = create_react_agent(
"openai:gpt-4o-mini",
tools=[add_tool, search_tool],
)
# Use the agent
response = await agent.ainvoke({
"messages": [HumanMessage(content="Remember: I prefer Python for coding")]
})
print(response["messages"][-1].content)
if __name__ == "__main__":
import asyncio
asyncio.run(main())
The get_sessionized_cognee_tools() function handles session isolation automatically, making it easy to build multi-tenant applications.
Phase 5: Testing and Evaluation
Now let’s evaluate our Cognee-powered agent to understand its capabilities.
Test 1: Basic Retrieval
First, we verify that the agent can retrieve information accurately.
import cognee
import asyncio
async def test_basic_retrieval():
"""Test basic information retrieval"""
print("=== Test 1: Basic Retrieval ===")
# Setup
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
# Add test data
test_data = """
Albert Einstein was born in 1879 in Ulm, Germany.
He developed the theory of relativity.
Einstein received the Nobel Prize in Physics in 1921.
"""
await cognee.add(test_data)
await cognee.cognify()
# Test queries
queries = [
"When was Einstein born?",
"What did Einstein develop?",
"When did Einstein receive the Nobel Prize?",
]
for query in queries:
results = await cognee.search(query)
print(f"\nQuery: {query}")
print(f"Results: {results}")
if __name__ == "__main__":
asyncio.run(test_basic_retrieval())
Test 2: Multi-Hop Reasoning
This is where Cognee’s graph structure provides advantages over standard RAG.
import cognee
import asyncio
async def test_multihop_reasoning():
"""Test multi-hop reasoning capabilities"""
print("=== Test 2: Multi-Hop Reasoning ===")
# Setup
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
# Add interconnected facts
knowledge = """
Marie Curie was a physicist and chemist.
Marie Curie discovered polonium and radium.
Polonium was named after Poland, Marie Curie's homeland.
Marie Curie won two Nobel Prizes.
Pierre Curie was Marie Curie's husband and research partner.
Pierre Curie died in a street accident in 1906.
Irene Joliot-Curie was Marie Curie's daughter.
Irene Joliot-Curie also won a Nobel Prize in Chemistry.
"""
await cognee.add(knowledge)
await cognee.cognify()
# Multi-hop queries
queries = [
"Who discovered the element named after Poland?",
"Did any of Marie Curie's family members win Nobel Prizes?",
"What happened to Marie Curie's research partner?",
]
for query in queries:
results = await cognee.search(query)
print(f"\nQuery: {query}")
print(f"Results: {results}")
if __name__ == "__main__":
asyncio.run(test_multihop_reasoning())
Test 3: Memory Persistence
This test verifies that information survives across sessions.
import cognee
import asyncio
async def test_persistence():
"""Test memory persistence"""
print("=== Test 3: Memory Persistence ===")
# Session 1: Store data
print("\n--- Session 1: Storing data ---")
await cognee.prune.prune_data()
await cognee.prune.prune_system(metadata=True)
await cognee.add("The secret code is ALPHA-BRAVO-7742.")
await cognee.cognify()
print("Data stored in session 1")
# Session 2: Retrieve data (no prune call!)
print("\n--- Session 2: Retrieving data ---")
results = await cognee.search("What is the secret code?")
print(f"Retrieved in session 2: {results}")
# Verify persistence
if results and "ALPHA-BRAVO-7742" in str(results):
print("\n✓ Memory persistence test PASSED")
else:
print("\n✗ Memory persistence test FAILED")
if __name__ == "__main__":
asyncio.run(test_persistence())
Comparing Cognee to Standard RAG
Traditional RAG Pipeline requires:
Cognee Pipeline requires:
Conclusion and Next Steps
In this tutorial, we built an agentic RAG pipeline using Cognee for persistent, graph-based memory. Here’s what we accomplished:
Phase 1 explained why standard RAG falls short for agents that need structural memory and multi-hop reasoning.
Phase 2 showed how to set up Cognee’s knowledge core using the simple ECL paradigm: add, cognify, search.
Phase 3 transformed Cognee’s capabilities into agent tools that can search and update memory.
Phase 4 integrated these tools with LangGraph to create an agent with persistent memory across sessions.
Phase 5 tested our implementation and compared it to traditional RAG approaches.
Key Takeaways
✅ Cognee replaces complex RAG infrastructure with three simple functions
✅ The knowledge graph enables multi-hop reasoning that vector-only approaches cannot achieve
✅ Memory persists automatically across agent sessions without custom state management
✅ The tool-based integration works with any agent framework that supports function calling
How to Improve It Further
There are several ways to extend what we built:
Resources