nano-graphrag

repository·main·Indexed 26 days ago

https://github.com/gusye1234/nano-graphrag

A lightweight, fast, and asynchronous implementation of GraphRAG designed as a portable alternative to the official Microsoft implementation. It supports Global Search, Local Search, and Naive RAG modes. The library features built-in support for OpenAI, Amazon Bedrock, Neo4j, and networkx, with extensible interfaces for custom LLMs, embedding functions, and storage backends.

Tokens
5.1K
Snippets
13
Records
28
Agent score
86%

What's inside nano-graphrag

  1. Benchmark comparison: nano-graphrag vs Microsoft GraphRAG

    main

    The project provides benchmark data comparing nano-graphrag against the official Microsoft GraphRAG implementation (specifically commit 61b5eea34783c58074b3c53f1689ad8a5ba6b6ee).

    Key Performance Metrics:

    • Indexing Speed: nano-graphrag is faster, indexing the benchmark corpus (A Christmas Carol) in less than 4 minutes, compared to over 5 minutes for GraphRAG.
    • Concurrency: nano-graphrag uses a maximum of 16 concurrent Async API requests, while GraphRAG uses up to 25.
    • Model Parity: Both implementations use gpt-4o and OpenAI Embeddings for the benchmark tests.

    Note: Benchmarks were conducted without caching on the same device and network connection.

  2. Configure Azure OpenAI and Amazon Bedrock

    main

    Azure OpenAI

    Refer to .env.example.azure to set your Azure credentials. Enable it by passing using_azure_openai=True to the GraphRAG constructor.

    Amazon Bedrock

    Ensure AWS credentials are set (e.g., via aws configure). Enable it by passing the following to the GraphRAG constructor:

    • using_amazon_bedrock=True
    • best_model_id: The ID for the high-capability model.
    • cheap_model_id: The ID for the efficient model.
  3. Install nano-graphrag

    main

    You can install nano-graphrag either from PyPI or from source. Installing from source is recommended if you wish to hack or modify the implementation.

    From PyPI:

    pip install nano-graphrag

    From Source:

    # clone this repo first
    cd nano-graphrag
    pip install -e .
    pip install nano-graphrag
  4. Set up Neo4j for GraphRAG storage

    main

    To use Neo4j as the storage backend for nano-graphrag, follow these steps:

    1. Install Neo4j: Use version 5.x.
    2. Install Neo4j GDS: Install the Graph Data Science (GDS) plugin.
    3. Start the server: Ensure the Neo4j server is running.
    4. Configure Credentials: Obtain your NEO4J_URL, NEO4J_USER, and NEO4J_PASSWORD.
      • Default URL: neo4j://localhost:7687
      • Default User: neo4j
      • Default Password: neo4j
  5. Quick Start with nano-graphrag

    main

    To use nano-graphrag, you must first set your OPENAI_API_KEY in your environment.

    1. Set Environment Variable:
    export OPENAI_API_KEY="sk-..."
    1. Basic Usage: Initialize GraphRAG with a working_dir. This directory will store the context, allowing it to be reloaded automatically in future sessions. Use .insert() to add text and .query() to perform searches.

    2. Search Modes:

    • Global Search: Default mode for high-level queries.
    • Local Search: More scalable and often better for specific queries. Use QueryParam(mode="local") to enable.
    • Naive RAG: Standard RAG without graph features. Enable via enable_naive_rag=True in the constructor and use QueryParam(mode="naive") during query.
    from nano_graphrag import GraphRAG, QueryParam
    
    graph_func = GraphRAG(working_dir="./dickens")
    
    with open("./book.txt") as f:
        graph_func.insert(f.read())
    
    # Perform global graphrag search
    print(graph_func.query("What are the top themes in this story?"))
    
    # Perform local graphrag search
    print(graph_func.query("What are the top themes in this story?", param=QueryParam(mode="local")))
  6. Customize storage components

    main

    You can replace the default storage backends by passing custom classes to the GraphRAG constructor:

    • Key-Value Storage: Use key_string_value_json_storage_cls to replace the default disk file storage (implements base.BaseKVStorage).
    • Vector Storage: Use vector_db_storage_cls to replace the default nano-vectordb (implements base.BaseVectorStorage).
    • Graph Storage: Use graph_storage_cls to replace the default networkx backend (implements base.BaseGraphStorage).
  7. Fine-tune Entity Relationship Extraction with DSPy

    main

    You can optimize the TypedEntityRelationshipExtractor from nano_graphrag using the DSPy framework to improve extraction performance. The process involves evaluating a baseline model, then using optimizers like BootstrapFewShotWithRandomSearch or MIPROv2 to refine prompt instructions and few-shot examples based on specific metrics.

    Workflow

    1. Load Data: Prepare training, validation, and development datasets (typically as pickled objects).
    2. Baseline Evaluation: Use dspy.evaluate.Evaluate with metrics like entity_recall_metric and relationships_similarity_metric to establish a baseline score for the TypedEntityRelationshipExtractor.
    3. Optimization:
      • Use BootstrapFewShotWithRandomSearch for simple bootstrapping.
      • Use MIPROv2 for advanced optimization that generates candidate instructions and few-shot examples using a larger 'prompt model' to guide a smaller 'task model'.
    4. Save Model: Once optimized, save the compiled model using the .save() method.
    from nano_graphrag.entity_extraction.module import TypedEntityRelationshipExtractor
    from nano_graphrag.entity_extraction.metric import relationships_similarity_metric, entity_recall_metric
    import dspy
    
    # Initialize the extractor
    model = TypedEntityRelationshipExtractor()
    
    # Example: Using BootstrapFewShotWithRandomSearch
    optimizer = dspy.teleprompt.BootstrapFewShotWithRandomSearch(
        metric=relationships_similarity_metric, 
        num_threads=os.cpu_count(),
        num_candidate_programs=10,
        max_labeled_demos=5,
        max_bootstrapped_demos=2,
    )
    rs_model = optimizer.compile(model, trainset=trainset, valset=valset)
    
    # Save the optimized model
    rs_model.save("path_to_save_model.json")
  8. Fix Leiden.EmptyNetworkError:EmptyNetworkError

    main

    The Leiden.EmptyNetworkError:EmptyNetworkError occurs when nano-graphrag attempts to compute communities on an empty network. This is typically caused by the LLM failing to extract any entities or relations.

    To resolve this:

    1. Verify LLM Output Format: Ensure the LLM response matches the expected format for entity extraction, for example: ("entity"<|>"Cruz"<|>"person"<|>"Cruz is associated with a vision of control and order...")
    2. Use System Instructions: If your LLM struggles with formatting, provide a system prompt to enforce the desired structure. You can use this as the default for your LLM calling function:
      {
          "role": "system",
          "content": "You are an intelligent assistant and will follow the instructions given to you to fulfill the goal. The answer should be in the format as in the given example."
      }
    3. Upgrade LLM: Try using a larger or more capable LLM model.
  9. Fix 'Didn't extract any entities' warning in Ollama

    main

    If you see warnings like WARNING:nano-graphrag:Didn't extract any entities, maybe your LLM is not working, it may be because the Ollama num_ctx (context window) is too small (default is often 2048) to handle the entity extraction prompt.

    You can fix this by creating a new Ollama model with a larger context window (e.g., 32000). For example, to update qwen2:

    1. Export the existing model configuration to a file: ollama show --modelfile qwen2 > Modelfile
    2. Edit the Modelfile and add a PARAMETER num_ctx line below the FROM instruction: PARAMETER num_ctx 32000
    3. Create the new model: ollama create -f Modelfile qwen2:ctx32k
    4. Use the new model name qwen2:ctx32k in your nano-graphrag configuration.
    ollama show --modelfile qwen2 > Modelfile
    # Add 'PARAMETER num_ctx 32000' to Modelfile
    ollama create -f Modelfile qwen2:ctx32k
  10. Use Neo4jStorage with GraphRAG

    main

    You can pass a Neo4j instance to GraphRAG by specifying Neo4jStorage as the graph_storage_cls and providing a configuration dictionary via addon_params. The configuration requires neo4j_url and neo4j_auth (a tuple of username and password).

    from nano_graphrag import GraphRAG
    from nano_graphrag._storage import Neo4jStorage
    import os
    
    neo4j_config = {
      "neo4j_url": os.environ.get("NEO4J_URL", "neo4j://localhost:7687"),
      "neo4j_auth": (
          os.environ.get("NEO4J_USER", "neo4j"),
          os.environ.get("NEO4J_PASSWORD", "neo4j"),
      )
    }
    GraphRAG(
      graph_storage_cls=Neo4jStorage,
      addon_params=neo4j_config,
    )