docs/examples/graph_rag/llama_index_cognee_integration.ipynb
<a href="https://colab.research.google.com/drive/1EpokQ8Y_5jIJ7HdixZms81Oqgh2sp7-E?usp=sharing" target="_parent"></a>
Connecting external knowledge to the LLM efficiently and retrieving it is a key challenge faced by developers. For developers and data scientists, integrating structured and unstructured data into AI workflows often involves multiple tools, complex pipelines, and time-consuming processes.
Enter cognee, a powerful framework for knowledge and memory management, and LlamaIndex, a versatile data integration library. Together, they enable us to transform retrieval-augmented generation (RAG) pipelines, into GraphRAG pipelines, streamlining the path from raw data to actionable insights.
In this post, we’ll explore a demo that leverages cognee and LlamaIndex to create a knowledge graph from a LlamaIndex document, process it into a meaningful structure, and extract useful insights. By the end, you’ll see how these tools can give you new insights into your data by connecting various data sources in one big semantic layer you can analyze.
RAG enhances LLMs by integrating external knowledge sources during inference. It does so by turning the data into a vector representation and storing it in a vector store.
However, building a RAG system presents challenges: handling diverse data formats, data updates, creating a robust metadata layer, and mediocre accuracy
cognee simplifies knowledge and memory management for LLMs, while LlamaIndex facilitates connecting LLMs to structured data sources and enabling agentic use-cases
cognee is inspired by human mind and higer cognitive functions. It mimics ways we construct our mental map of the world and build a semantic understanding of various objects, terms and issues in our everyday lives.
cognee brings this approach to code by allowing developers to create semantic layers that would allow users to store their ontologies which are a formalised depiction of knowledge in graphs.
This lets you use the knowledge you have about a system connect it to LLMs in a modular way, with best data engineering practices, wide choice of vector and graph stores and various LLMs you can use.
Together, they:
Start by importing the required libraries and defining the environment:
!pip install llama-index-graph-rag-cognee==0.1.3
import os
import asyncio
from llama_index.core import Document
from llama_index.graph_rag.cognee import CogneeGraphRAG
if "OPENAI_API_KEY" not in os.environ:
os.environ["OPENAI_API_KEY"] = ""
Ensure you’ve set up your API keys and installed necessary dependencies.
We’ll use a brief profile of an individual as our sample dataset:
documents = [
Document(
text="Jessica Miller, Experienced Sales Manager with a strong track record in driving sales growth and building high-performing teams."
),
Document(
text="David Thompson, Creative Graphic Designer with over 8 years of experience in visual design and branding."
),
]
Instantiate the Cognee framework with configurations for LLM, graph, and database providers:
cogneeRAG = CogneeGraphRAG(
llm_api_key=os.environ["OPENAI_API_KEY"],
llm_provider="openai",
llm_model="gpt-4o-mini",
graph_db_provider="networkx",
vector_db_provider="lancedb",
relational_db_provider="sqlite",
relational_db_name="cognee_db",
)
Load the dataset into the cognee framework:
await cogneeRAG.add(documents, "test")
This step prepares the data for graph-based processing.
Transform the data into a structured knowledge graph:
await cogneeRAG.process_data("test")
The graph now contains nodes and relationships derived from the dataset, creating a powerful structure for exploration.
search_results = await cogneeRAG.search(
"Tell me who are the people mentioned?"
)
print("\n\nAnswer based on knowledge graph:\n")
for result in search_results:
print(f"{result}\n")
search_results = await cogneeRAG.rag_search(
"Tell me who are the people mentioned?"
)
print("\n\nAnswer based on RAG:\n")
for result in search_results:
print(f"{result}\n")
In conclusion, the results demonstrate a significant advantage of the knowledge graph-based approach (Graphrag) over the RAG approach. Graphrag successfully identified all the mentioned individuals across multiple documents, showcasing its ability to aggregate and infer information from a global context. In contrast, the RAG approach was limited to identifying individuals within a single document due to its chunking-based processing constraints. This highlights Graphrag's superior capability in comprehensively resolving queries that span across a broader corpus of interconnected data.
Explore relationships in the knowledge graph:
related_nodes = await cogneeRAG.get_related_nodes("person")
print("\n\nRelated nodes are:\n")
for node in related_nodes:
print(f"{node}\n")
Your agents can now get long-term, short-term memory and memory specific to their domains
Your memory can now automatically optimize itself and allow to respond to questions better
You can use the standard tools out of the box and get things done without much effort
Imagine a graph structure where each node represents a document or entity, and edges indicate relationships.
Here’s the visualized knowledge graph from the simple example above:
Try running it yourself