GraphRAG Toolkit

# For document ingestion and lexical graph construction:
pip install graphrag-lexical-graph

# For querying existing knowledge graphs:
pip install graphrag-toolkit-byokg-rag

from graphrag_toolkit.storage import GraphStoreFactory, VectorStoreFactory

# Amazon Neptune example:
graph_store = GraphStoreFactory.for_graph_store(
    'neptune-db://your-cluster.cluster-xxxx.us-east-1.neptune.amazonaws.com:8182'
)

# Amazon OpenSearch Serverless example:
vector_store = VectorStoreFactory.for_vector_store(
    'aoss://your-collection.us-east-1.aoss.amazonaws.com'
)

from graphrag_toolkit.indexing import LexicalGraphIndex
from llama_index.core import SimpleDirectoryReader

# Load your documents
documents = SimpleDirectoryReader('./docs').load_data()

# Build the graph
index = LexicalGraphIndex(
    graph_store=graph_store,
    vector_store=vector_store
)
index.extract_and_build(documents, show_progress=True)

from graphrag_toolkit.retrieval import GraphRAGQueryEngine

query_engine = GraphRAGQueryEngine(
    graph_store=graph_store,
    vector_store=vector_store
)

response = query_engine.query(
    "What are the key relationships between the entities in these documents?"
)
print(response)

{
  "mcpServers": {
    "graphrag-toolkit": {
      "command": "python",
      "args": ["-m", "graphrag_toolkit.mcp"],
      "env": {
        "NEPTUNE_ENDPOINT": "neptune-db://your-cluster.cluster-xxxx.us-east-1.neptune.amazonaws.com:8182",
        "AOSS_ENDPOINT": "aoss://your-collection.us-east-1.aoss.amazonaws.com",
        "AWS_DEFAULT_REGION": "us-east-1"
      }
    }
  }
}