GraphRAG-SDK

GraphRAG SDK uses two types of providers: LLM (for text generation and extraction) and Embedder (for vector embeddings).

Supported Providers

Recommendations

• Production (Azure): Use LiteLLM with the azure/ prefix.
• Development (OpenAI): Use LiteLLM with OpenAI models.
• Budget-conscious: Use OpenRouterLLM to compare model prices.
• Local models: Implement a custom LLMInterface (e.g., wrapping Ollama or vLLM).
• Local embeddings: Implement a custom Embedder (e.g., wrapping sentence-transformers).

The GraphStore class is the primary write path for all graph data in the SDK. It converts Python objects (GraphNode, GraphRelationship) into Cypher queries.

Key behaviors:

• Idempotency: It uses MERGE instead of CREATE, making re-ingestion safe as it updates existing nodes/relationships rather than creating duplicates.
• Batching: Writes are performed in batches of 500 items using UNWIND for high performance.
• Label Hints: To optimize performance, relationship MATCH queries use label hints (e.g., PART_OF looks for (Document)-[:PART_OF]->(Chunk)). Unknown edge types default to (__Entity__, __Entity__).
• Error Handling: If a batch fails, GraphStore falls back to per-item upserts. For nodes, the first failure raises a DatabaseError. For relationships, failures are logged as warnings and processing continues.
• Data Safety: It includes a None-ID Guard to prevent nodes with empty IDs from being written, and _clean_properties() to sanitize data (e.g., converting dict to JSON strings and dropping None values).

-- Example of the internal Node MERGE pattern
UNWIND $batch AS item
MERGE (n:`Person` {id: item.id})
SET n += item.properties
SET n:__Entity__

When you call rag.ingest(), the SDK runs a 9-step sequential pipeline to transform raw text into a structured knowledge graph.

The Pipeline Flow:

1. Load: Reads text from files (auto-detects .pdf via PdfLoader or others via TextLoader) or direct strings.
2. Chunk: Splits text into overlapping windows (default: 1000 chars size, 100 chars overlap) to ensure entities aren't split across boundaries.
3. Build Lexical Graph (Mandatory): Creates the provenance backbone (Document nodes, Chunk nodes, PART_OF edges, and NEXT_CHUNK edges) so every answer is traceable to its source.
4. Extract Entities & Relationships: Uses an LLM or local models (like GLiNER) to identify entities and their connections. 4b. Quality Filter: Removes malformed data (e.g., nodes with empty IDs).
5. Prune Against Schema: Filters data to match your defined GraphSchema. If no schema is provided, all data passes through.
6. Resolve Duplicates: Merges entities referring to the same real-world object (e.g., using ExactMatchResolution or DescriptionMergeResolution).
7. Write to Graph: Persists data to FalkorDB using batched Cypher queries.
8. Write Mentions (Parallel): Creates MENTIONED_IN edges between entities and chunks for retrieval.
9. Index Chunks (Parallel): Embeds chunk text into vectors for similarity search.

Steps 1-7 are sequential; steps 8-9 run in parallel.

The extraction process converts raw text into structured knowledge using a two-step hybrid approach:

1. Step 1: Entity NER (Named Entity Recognition): A fast, local model (like GLiNER) or an LLM identifies potential entities (name, type, confidence, and character spans).
2. Step 2: LLM Verify + Relationship Extraction: An LLM receives the pre-extracted entities and the original text. It verifies the entities (fixing names, adding missed ones) and extracts factual relationships between them.

Optional Step: Coreference Resolution: Before extraction, a resolver can replace pronouns (e.g., "she") with the canonical entity name (e.g., "Alice") to improve extraction accuracy.

Aggregation: After processing all chunks, the SDK deduplicates entities and relationships across the entire dataset, merging descriptions, source chunk IDs, and character spans.

            Text Chunk
                |
     ┌──────────┴──────────┐
     v                      v
  (Optional)           Step 1: Entity NER
  Coreference          ┌─────────────────┐
  Resolution           │  GLiNERExtractor │  (default, local)
  (resolve pronouns)   │  LLMExtractor    │  (API-based)
  (your own)           │  Custom          │  (your own)
     |                 └────────┬────────┘
  Resolved Text                 |
                        List of entities
                        (name, type, confidence, spans)
                                |
                     Step 2: LLM Verify + Relationships
                     ┌──────────────────────────────────┐
                     │ LLM receives:                     │
                     │   - Pre-extracted entities         │
                     │   - Original text                  │
                     │                                    │
                     │ LLM returns:                       │
                     │   - Verified entities (fixed/added) │
                     │   - Relationships with evidence     │
                     └──────────────┬───────────────────┘
                                    |
                          Aggregate across chunks
                          (dedup entities + relations)
                                    |
                          Convert to GraphData
                          (GraphNode + GraphRelationship)

The VectorStore class handles vector and fulltext operations, including index creation and retrieval.

Indexing Types

• Vector Indexes: Created for nodes (e.g., Chunk) or relationships (e.g., RELATES) using the CREATE VECTOR INDEX syntax.
• Fulltext Indexes: Uses the RediSearch-based API to index text properties on nodes like Chunk or __Entity__.

Key Ingestion Tasks

• Chunk Indexing: Chunks are embedded in batches and written to the database using UNWIND. If batching fails, it falls back to individual processing.
• Entity Embedding Backfill: During finalize(), the SDK identifies __Entity__ nodes where embedding IS NULL, embeds them, and writes the vectors back to the graph.
• Relationship Embedding: RELATES edges containing a fact property but no embedding are batch-embedded and updated.

Search Capabilities

• Vector Search: Supports searching Chunk nodes, __Entity__ nodes, and RELATES edges. For older FalkorDB versions, relationship search falls back to a Cypher-based cosine distance scan.
• Fulltext Search: Allows keyword-based searching on indexed text properties.
• Optimization: The rerank_chunks() function uses stored embeddings to compute similarity locally, avoiding expensive API calls if coverage is $\ge 90%$.