GraphRAG: Enhancing LLMs with knowledge graphs for superior retrieval

pip install graphrag weave

mkdir -p ./ragtest/input curl https://www.gutenberg.org/cache/epub/24022/pg24022.txt -o ./ragtest/input/book.txt

graphrag init --root ./ragtest

llm:
  api_key: ${GRAPHRAG_API_KEY} # set this in the generated .env file
  type: openai_chat # or azure_openai_chat
  model: gpt-4o-mini
  model_supports_json: true # recommended if this is available for your model.
  
parallelization:
  stagger: 1.0
  num_threads: 1

async_mode: threaded # or asyncio

embeddings:
  async_mode: threaded # or asyncio
  vector_store: 
    type: lancedb
    db_uri: 'output/lancedb'
    container_name: default
    overwrite: true
  llm:   
    api_key: ${GRAPHRAG_API_KEY}
    type: openai_embedding # or azure_openai_embedding
    model: text-embedding-3-small


from pathlib import Path
from graphrag.logger.types import LoggerType
from graphrag.cli.index import index_cli
import weave

# Initialize Weave for logging/telemetry
weave.init("graph_rag_index_pipeline")

def main():
    # Hardcoded parameters
    root_dir = Path("./")
    config_path = Path("./settings.yaml")
    verbose = True
    memprofile = False
    resume = None  # Resume an existing run; set to None for fresh runs
    logger = LoggerType.RICH  # Logger type, e.g., RICH or SIMPLE
    dry_run = False  # If True, pipeline will validate but not execute
    cache = True  # Enable or disable LLM cache
    skip_validation = False  # Skip validation steps
    output_dir = Path("./output")  # Output directory for indexing results

    print("Starting the indexing pipeline with Weave telemetry...")
    try:
        index_cli(
            root_dir=root_dir,
            config_filepath=config_path,
            verbose=verbose,
            resume=resume,
            memprofile=memprofile,
            cache=cache,
            logger=logger,
            dry_run=dry_run,
            skip_validation=skip_validation,
            output_dir=output_dir,
        )
        print("Indexing pipeline completed successfully.")
    except Exception as e:
        print(f"Error during indexing: {e}")
        weave.log("index_pipeline_error", {"error": str(e)})

if __name__ == "__main__":
    main()


from pathlib import Path
from graphrag.cli.query import run_global_search
import weave; weave.init("graph_rag_global_search")

def main():
    # Hardcoded parameters
    root_dir = Path("./")
    data_dir = Path("./output")
    config_path = Path("./settings.yaml")
    query = "What is the role of Scrooge in the narrative?"
    community_level = 2
    dynamic_community_selection = False
    response_type = "Multiple Paragraphs"
    streaming = False

    print(f"Running global search with query: '{query}'...")
    try:
        run_global_search(
            config_filepath=config_path,  # Adjust if needed to include configs
            data_dir=data_dir,
            root_dir=root_dir,
            community_level=community_level,
            dynamic_community_selection=dynamic_community_selection,
            response_type=response_type,
            streaming=streaming,
            query=query,
        )
        print("Query completed successfully.")
    except Exception as e:
        print(f"Error while running global search: {e}")

if __name__ == "__main__":
    main()

from pathlib import Path
from graphrag.cli.query import run_local_search
import weave


weave.init("graph_rag_local_search")

def main():
    # Hardcoded parameters
    root_dir = Path("./")
    data_dir = Path("./output")
    config_path = Path("./settings.yaml")
    query = "How does Scrooge's reaction to the Ghost of Christmas Yet to Come illustrate his fear of the future?"
    community_level = 2
    response_type = "Multiple Paragraphs"
    streaming = False

    print(f"Running local search with query: '{query}'...")
    try:
        run_local_search(
            config_filepath=config_path,
            data_dir=data_dir,
            root_dir=root_dir,
            community_level=community_level,
            response_type=response_type,
            streaming=streaming,
            query=query,
        )
        print("Query completed successfully.")
    except Exception as e:
        print(f"Error while running local search: {e}")

if __name__ == "__main__":
    main()