Skip to content

getting-started.md

Latest commit

 

History

History
251 lines (172 loc) · 7.22 KB

getting-started.md

File metadata and controls

251 lines (172 loc) · 7.22 KB
Raw

Getting Started with GraphRAG SDK

A step-by-step tutorial for building a knowledge graph from documents and querying it with natural language.

1. Prerequisites

  • Python 3.10+
  • FalkorDB (easiest via Docker -- see below)
  • An LLM API key from one of the supported providers:
    • Azure OpenAI
    • OpenAI
    • Anthropic

2. Installation

Install the SDK with all optional dependencies:

pip install "graphrag-sdk[all]"

For a local editable install from a cloned repo:

pip install -e "./graphrag_sdk[all]"

3. Start FalkorDB

Run FalkorDB as a Docker container:

docker run -d --name falkordb -p 6379:6379 falkordb/falkordb

Verify it is running:

docker ps | grep falkordb

4. Configure Environment

Set the environment variables for your LLM provider. The example below uses Azure OpenAI:

export AZURE_OPENAI_API_KEY="your-key"
export AZURE_OPENAI_ENDPOINT="https://your-resource.openai.azure.com/"
export AZURE_OPENAI_API_VERSION="2024-12-01-preview"

If you use a .env file, load it yourself before importing the SDK (e.g., via python-dotenv or export commands). The SDK reads environment variables but does not auto-load .env files.

5. Define a Schema

A GraphSchema tells the extraction pipeline which entity and relationship types to look for in your documents.

from graphrag_sdk import GraphSchema, EntityType, RelationType

schema = GraphSchema(
    entities=[
        EntityType(label="Person", description="A human being"),
        EntityType(label="Organization", description="A company or institution"),
        EntityType(label="Place", description="A geographic location"),
    ],
    relations=[
        RelationType(label="WORKS_AT", description="Employment relationship"),
        RelationType(label="LOCATED_IN", description="Geographic location"),
        RelationType(label="KNOWS", description="Personal acquaintance"),
    ],
)

You can add as many entity and relationship types as your domain requires. Descriptions help the LLM decide when to extract each type.

6. Initialize GraphRAG

Create a GraphRAG instance by providing a connection, LLM, embedder, and schema:

from graphrag_sdk import GraphRAG, ConnectionConfig, LiteLLM, LiteLLMEmbedder

rag = GraphRAG(
    connection=ConnectionConfig(host="localhost", graph_name="my_graph"),
    llm=LiteLLM(model="azure/gpt-4.1"),
    embedder=LiteLLMEmbedder(model="azure/text-embedding-3-large", dimensions=256),
    schema=schema,
    embedding_dimension=256,  # must match your embedding model's output dimension
)

The default is 256 (matched-Matryoshka dimensions of text-embedding-3-large). If your embedding model produces a different dimensionality (e.g., 1024, 1536, or 3072), set embedding_dimension accordingly.

ConnectionConfig accepts additional parameters such as port, username, password, and query_timeout_ms. See docs/configuration.md for the full list.

Alternative providers: The SDK also exports OpenRouterLLM and OpenRouterEmbedder for use with OpenRouter. See docs/configuration.md for details.

7. Ingest a Document

From a file path

result = await rag.ingest("path/to/document.txt")
print(f"Created {result.nodes_created} nodes, {result.relationships_created} relationships")
print(f"Indexed {result.chunks_indexed} chunks")

From raw text

result = await rag.ingest("acme_doc", text="Acme Corp was founded in 1985 by Jane Doe in Austin, Texas.")
print(f"Created {result.nodes_created} nodes, {result.relationships_created} relationships")
print(f"Indexed {result.chunks_indexed} chunks")

The ingestion pipeline runs a 9-step process: Load, Chunk, Lexical Graph, Extract (includes quality filtering), Prune, Resolve, Write, then Mentions and Chunk Indexing in parallel.

8. Query the Knowledge Graph

Retrieve context only

Use retrieve() when you want to inspect the context or use your own LLM:

context = await rag.retrieve("Who works at Acme Corp?")
for item in context.items:
    print(f"[{item.score:.2f}] {item.content[:100]}...")

Generate an answer

Use completion() for the full RAG pipeline — retrieval + answer generation:

result = await rag.completion("Who works at Acme Corp?")
print(result.answer)

With context inspection

Pass return_context=True to see which chunks and entities the retriever used to build the answer:

result = await rag.completion("Who works at Acme Corp?", return_context=True)
for item in result.retriever_result.items:
    print(f"[{item.metadata.get('section', '')}] {item.content[:100]}...")

Multi-turn conversations

completion() supports native multi-turn conversations. Messages are passed directly to the LLM's chat API as structured messages:

from graphrag_sdk import ChatMessage

result = await rag.completion(
    "What happened to her after that?",
    history=[
        ChatMessage(role="user", content="Who works at Acme Corp?"),
        ChatMessage(role="assistant", content="Jane Doe works at Acme Corp."),
    ],
)
print(result.answer)

You can also pass history as plain dicts:

result = await rag.completion(
    "Tell me more.",
    history=[
        {"role": "user", "content": "Who founded Acme?"},
        {"role": "assistant", "content": "Jane Doe founded Acme in 1985."},
    ],
)

Supported roles: "system", "user", "assistant". Invalid roles raise ValueError.

9. Inspect Graph Contents

Use get_statistics() to see a summary of what the graph contains:

stats = await rag.graph_store.get_statistics()
print(f"Nodes: {stats['node_count']}, Edges: {stats['edge_count']}")

You can also run raw Cypher queries against the graph:

results = await rag.graph_store.query_raw("MATCH (p:Person)-[:WORKS_AT]->(o:Organization) RETURN p.name, o.name LIMIT 10")
for row in results.result_set:
    print(row)

10. Finalize

After all documents have been ingested, run finalize() to deduplicate entities, backfill embeddings, and create indexes:

results = await rag.finalize()
print(f"Deduplicated: {results['entities_deduplicated']}")
print(f"Embedded: {results['entities_embedded']} entities, {results['relationships_embedded']} relationships")

This step is important for query accuracy. It merges duplicate entities (e.g., "J. Doe" and "Jane Doe") and ensures all entities have vector embeddings for semantic search.

11. Next Steps

  • docs/configuration.md -- Tuning connection settings, chunking parameters, and retrieval options.
  • docs/strategies.md -- Custom extraction and resolution strategies.
  • docs/benchmark.md -- Reproducing benchmark results on the GraphRAG-Bench Novel corpus (20 novels, 2,010 questions).

Synchronous API

If you are not in an async context, use the synchronous convenience methods:

result = rag.ingest_sync("path/to/document.txt")
context = rag.retrieve_sync("Who works at Acme Corp?")
result = rag.completion_sync("Who works at Acme Corp?")
results = rag.finalize_sync()

These wrap the async methods in asyncio.run() and are useful for scripts and notebooks.