pip install llama-index-graph-rag-cognee

import os
import asyncio
from llama_index.core import Document
from llama_index.graph_rag.cognee import CogneeGraphRAG


async def main():
    # Initialize CogneeGraphRAG
    cognee_rag = CogneeGraphRAG(
        llm_api_key=os.environ["OPENAI_API_KEY"],
        llm_provider="openai",
        llm_model="gpt-4o-mini",
        graph_db_provider="kuzu",  # or "neo4j", "networkx"
        vector_db_provider="lancedb",
        relational_db_provider="sqlite",
        relational_db_name="cognee_db",
    )

    # Create documents
    documents = [
        Document(
            text="Apple Inc. is a technology company founded by Steve Jobs."
        ),
        Document(
            text="Steve Jobs was the CEO of Apple and known for innovation."
        ),
        Document(text="The iPhone was released by Apple in 2007."),
    ]

    # Add documents to the knowledge graph
    await cognee_rag.add(documents, dataset_name="apple_knowledge")

    # Process data into the knowledge graph
    await cognee_rag.process_data("apple_knowledge")

    # Search the knowledge graph
    results = await cognee_rag.search("Who founded Apple?")
    print("Search Results:", results)

    # Generate and visualize the knowledge graph
    viz_path = await cognee_rag.visualize_graph(
        open_browser=True,  # Automatically open in browser
        output_file_path=".",  # Save to current directory
    )
    print(f"Visualization saved to: {viz_path}")


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

import os
import pandas as pd
import asyncio
from llama_index.core import Document
from llama_index.graph_rag.cognee import CogneeGraphRAG


async def advanced_example():
    # Load news data
    news = pd.read_csv(
        "https://raw.githubusercontent.com/tomasonjo/blog-datasets/main/news_articles.csv"
    )[:5]

    documents = [
        Document(text=f"{row['title']}: {row['text']}")
        for i, row in news.iterrows()
    ]

    # Initialize with enterprise databases
    cognee_rag = CogneeGraphRAG(
        llm_api_key=os.environ["OPENAI_API_KEY"],
        llm_provider="openai",
        llm_model="gpt-4o-mini",
        graph_db_provider="neo4j",  # Enterprise graph database
        vector_db_provider="qdrant",  # Scalable vector database
        relational_db_provider="postgresql",
        relational_db_name="cognee_production",
    )

    # Build knowledge graph
    await cognee_rag.add(documents, "news_dataset")
    await cognee_rag.process_data("news_dataset")

    # Multiple search approaches
    print("=== Graph-based Search ===")
    graph_results = await cognee_rag.search(
        "Tell me about the people mentioned"
    )
    for result in graph_results:
        print(f"📊 {result}")

    print("\n=== RAG-based Search ===")
    rag_results = await cognee_rag.rag_search(
        "Tell me about the people mentioned"
    )
    for result in rag_results:
        print(f"🔍 {result}")

    print("\n=== Related Nodes ===")
    related = await cognee_rag.get_related_nodes("person")
    for node in related:
        print(f"🔗 {node}")

    # Generate visualization
    await cognee_rag.visualize_graph(
        open_browser=False, output_file_path="/path/to/your/output/directory"
    )


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