GraphRAG - Part 1: Data Pipeline

Last update: February 17, 2026

This repo is part of a larger GraphRAG project, in which I demonstrate how the GraphRAG pattern works end-to-end.

Part 1 is a Knowledge Graph Construction & RAG Preparation Pipeline built with Dagster, which orchestrates data ingestion from Wikipedia and Wikidata, extracts structured knowledge using LLM-based information extraction (Google Gemini), constructs a knowledge graph in Neo4j, and prepares vector embeddings in Pinecone for a Retrieval-Augmented Generation system. The app is containerized and ready for cloud deployment.

Domain: Electronic Music

This pipeline is specifically tuned for the Electronic Music domain. It captures the rich, interconnected history of electronic artists, from early pioneers to contemporary producers. The dataset encompasses a wide range of sub-genres—including Techno, House, Ambient, IDM, and Drum & Bass—modeling the complex relationships between artists, their releases, and the evolving taxonomy of electronic musical styles.

About the Project

This data pipeline combines structured data extraction from knowledge bases with LLM-powered information extraction from unstructured text to build a comprehensive music knowledge graph:

Structured Data Sources
- Wikidata (Artist & Genre discovery via SPARQL, entity resolution via Action API)
- Wikipedia (Full-text articles for artists and genres)
LLM-Based Information Extraction
- Google Gemini extracts entities (persons, groups, albums, songs, labels, genres, etc.) and typed relationships from Wikipedia article chunks
- Extracted mentions are aggregated, deduplicated, and resolved against Wikidata for canonical identifiers
- spaCy (transformer-based) provides vector similarity for entity disambiguation

The goal is to transform unstructured Wikipedia articles into a structured knowledge graph and vector-ready text chunks, enabling a hybrid search approach:

Deterministic Search: Graph traversal and fulltext queries via Neo4j.
Semantic Search: Vector similarity over embedded text chunks and community summaries via Pinecone.

We leverage Polars for high-performance lazy data transformation, msgspec for fast serialization, Pydantic for LLM output validation, and Dagster Resources for clean infrastructure management.

Tech Stack

Orchestration: Dagster (Assets, Resources, Partitions, Asset Checks, I/O Managers)
Databases: Neo4j (Graph — Aura cloud), Pinecone (Vector)
Data Engineering: Polars (lazy evaluation), Msgspec (serialization), Tiktoken (tokenization), Ftfy (text cleaning)
Data Validation & Config: Pydantic, Pydantic Settings
AI & ML: Google Gemini (LLM — entity extraction & community summaries), Sentence Transformers (Snowflake Arctic Embed S — graph embeddings, 384d), Pinecone Inference (llama-text-embed-v2 — chunk embeddings, 1024d), spaCy (en_core_web_trf — entity disambiguation)
NLP & Text Processing: LangChain Text Splitters (chunking), Tiktoken (cl100k_base tokenizer)
Graph Algorithms: igraph, leidenalg (hierarchical community detection)
Networking: curl-cffi (async HTTP with browser impersonation)
Logging & Utilities: Structlog, Tqdm
Language: Python 3.13+
Tooling: uv (package manager), Ruff (linter/formatter), Ty (type checker), Bandit (security), Pytest (testing)
Deployment: Docker, Docker Compose

Architecture Overview

This project implements a strict separation of concerns following Dagster's philosophy, dividing the codebase into distinct layers with clear responsibilities.

Architecture Visualization

┌──────────────────────────────────────────────────────────────────────────────┐
│                         DAGSTER DEFINITIONS LAYER                            │
│                    "The How" — Infrastructure & Configuration                │
│  ┌─────────────────┐  ┌─────────────────┐  ┌─────────────────────────────┐   │
│  │  definitions.py │  │   resources.py  │  │     io_managers.py          │   │
│  │  (Entry Point)  │  │ (Factories/DI)  │  │ (Parquet/JSONL Persistence) │   │
│  └────────┬────────┘  └────────┬────────┘  └────────────┬────────────────┘   │
│           │                    │                        │                    │
│  ┌────────┴────────┐  ┌────────┴────────┐  ┌────────────┴──────────┐         │
│  │   checks.py     │  │  partitions.py  │  │    settings.py        │         │
│  │ (Quality Gates) │  │  (By Decade)    │  │ (pydantic-settings)   │         │
│  └─────────────────┘  └─────────────────┘  └───────────────────────┘         │
└──────────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────────┐
│                          ASSET LAYER (defs/assets/)                         │
│                   "The What" — Business Logic & Transformation              │
│                                                                             │
│  ┌─────────────────────────────────────────────────────────────────────┐     │
│  │  STAGE 1: Artist & Genre Discovery                                  │     │
│  │  ┌─────────────────┐   ┌─────────────────┐   ┌─────────────────┐   │     │
│  │  │ artist_index.py │──▶│  artists.py     │──▶│  genres.py      │   │     │
│  │  │ (Wikidata)      │   │  (Enrichment)   │   │  (Extraction)   │   │     │
│  │  └─────────────────┘   └────────┬────────┘   └────────┬────────┘   │     │
│  └─────────────────────────────────┼──────────────────────┼───────────┘     │
│                                    │                      │                 │
│  ┌─────────────────────────────────┼──────────────────────┼───────────┐     │
│  │  STAGE 2: Content Acquisition   │                      │           │     │
│  │                                 ▼                      ▼           │     │
│  │  ┌─────────────────────────────────────────────────────────────┐   │     │
│  │  │          wikipedia_articles.py (Fetch & Cache)              │   │     │
│  │  └──────────────────────────┬──────────────────────────────────┘   │     │
│  └─────────────────────────────┼─────────────────────────────────────┘     │
│                                │                                           │
│  ┌─────────────────────────────┼─────────────────────────────────────┐     │
│  │  STAGE 3: Chunking & Embedding                                    │     │
│  │                             ▼                                     │     │
│  │  ┌──────────────┐   ┌──────────────────┐                          │     │
│  │  │  chunks.py   │──▶│  vector_db.py    │──▶ Pinecone (chunks)     │     │
│  │  │  (Tiktoken)  │   │  (Pinecone Inf.) │                          │     │
│  │  └──────┬───────┘   └──────────────────┘                          │     │
│  └─────────┼─────────────────────────────────────────────────────────┘     │
│            │                                                               │
│  ┌─────────┼─────────────────────────────────────────────────────────┐     │
│  │  STAGE 4: LLM Information Extraction                              │     │
│  │         ▼                                                         │     │
│  │  ┌──────────────┐   ┌─────────────────────┐   ┌────────────────┐  │     │
│  │  │ mentions.py  │──▶│ aggregated_         │──▶│ wikidata_      │  │     │
│  │  │ (Gemini)     │   │ mentions.py         │   │ candidates.py  │  │     │
│  │  └──────────────┘   └─────────────────────┘   └───────┬────────┘  │     │
│  │                                                        │          │     │
│  │  ┌───────────────────┐   ┌──────────────────┐         │          │     │
│  │  │ graph_elements.py │◀──│ linked_entities  │◀────────┘          │     │
│  │  │ (Global Merge)    │   │ .py (spaCy)      │                    │     │
│  │  └───────┬───────────┘   └──────────────────┘                    │     │
│  └──────────┼────────────────────────────────────────────────────────┘     │
│             │                                                              │
│  ┌──────────┼────────────────────────────────────────────────────────┐     │
│  │  STAGE 5: Knowledge Graph Construction                            │     │
│  │          ▼                                                        │     │
│  │  ┌─────────────────────┐   ┌──────────────────┐                   │     │
│  │  │ graph_summaries.py  │──▶│  graph_db.py     │──▶ Neo4j Aura    │     │
│  │  │ (Entity Summaries)  │   │  (Ingest + Index) │                  │     │
│  │  └─────────────────────┘   └────────┬─────────┘                   │     │
│  └─────────────────────────────────────┼─────────────────────────────┘     │
│                                        │                                   │
│  ┌─────────────────────────────────────┼─────────────────────────────┐     │
│  │  STAGE 6: Community Detection (GraphRAG)                          │     │
│  │                                     ▼                             │     │
│  │  ┌───────────────────────┐   ┌──────────────────────┐             │     │
│  │  │ detected_communities  │──▶│ community_summaries  │──▶ Pinecone │     │
│  │  │ .py (Leiden)          │   │ .py (Gemini)         │  (summaries)│     │
│  │  └───────────────────────┘   └──────────────────────┘             │     │
│  └───────────────────────────────────────────────────────────────────┘     │
└───────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌───────────────────────────────────────────────────────────────────────────┐
│                         UTILITIES LAYER (utils/)                          │
│                   Domain-Agnostic, Reusable Components                    │
│  ┌────────────────────────────────────────────────────────────────────┐   │
│  │                    NETWORK & I/O PRIMITIVES                        │   │
│  │  ┌─────────────────────┐  ┌────────────────────────────────────┐   │   │
│  │  │  network_helpers.py │  │      io_helpers.py                 │   │   │
│  │  │  (HTTP retries,     │  │  (JSON/text files,                 │   │   │
│  │  │   concurrency)      │  │   cache key generation)            │   │   │
│  │  └─────────────────────┘  └────────────────────────────────────┘   │   │
│  └────────────────────────────────────────────────────────────────────┘   │
│  ┌────────────────────────────────────────────────────────────────────┐   │
│  │                      DOMAIN ADAPTERS (API Clients)                 │   │
│  │  ┌─────────────────┐  ┌─────────────────┐  ┌─────────────────┐    │   │
│  │  │ wikidata_       │  │ wikipedia_      │  │ neo4j_          │    │   │
│  │  │ helpers.py      │  │ helpers.py      │  │ helpers.py      │    │   │
│  │  └─────────────────┘  └─────────────────┘  └─────────────────┘    │   │
│  │  ┌─────────────────┐  ┌─────────────────┐                         │   │
│  │  │ pinecone_       │  │ llm_            │                         │   │
│  │  │ helpers.py      │  │ helpers.py      │                         │   │
│  │  └─────────────────┘  └─────────────────┘                         │   │
│  └────────────────────────────────────────────────────────────────────┘   │
│  ┌─────────────────────────────────────────────────────────────────────┐  │
│  │                     DATA TRANSFORMATION                             │  │
│  │  ┌───────────────────────────────────────────────────────────────┐  │  │
│  │  │  data_transformation_helpers.py                               │  │  │
│  │  │  (Text normalization, chunking, deduplication, Unicode fixing) │  │  │
│  │  └───────────────────────────────────────────────────────────────┘  │  │
│  └─────────────────────────────────────────────────────────────────────┘  │
└───────────────────────────────────────────────────────────────────────────┘

Separation of Concerns

The architecture strictly separates "The What" (business logic) from "The How" (infrastructure), following Dagster best practices.

1. The `assets/` Folder: "The What" (Business Logic)

This folder contains the Data Definition Graph.

Aspect	Description
Role	Defines what data exists, how it is computed, and its dependencies
Content	Pure transformation logic. Takes data in (as parameters) and returns data out (as return values)
Change Frequency	High. This is where you edit code when business requirements change
Dagster Rule	Assets are unaware of where they run or where data is stored. They "ask" for a resource and "return" a dataframe

Key Design Patterns:

Assets return pl.LazyFrame, list[Model], or MaterializeResult — never write files directly
LLM extraction prompts and entity type definitions live in assets, not utils
Validation logic and business rules are implemented within assets
Assets delegate parsing/filtering to helper modules for reusability

2. The `defs/` Root Files: "The How" (Infrastructure)

These files define the Execution Environment.

File	Role	Change Frequency
`resources.py`	Connection factories (Neo4j, Pinecone, Wikidata, Wikipedia, Gemini). Handles secrets, timeouts, connection pooling	Low
`io_managers.py`	Bridge between Python memory and File System/Cloud Storage. Handles serialization (JSONL vs Parquet) and path organization	Low
`partitions.py`	Slicing strategy. Defines the "shape" of pipeline execution (by decade)	Low
`checks.py`	Quality contracts. Defines rules the data must obey after materialization	Medium

Key Design Patterns:

Explicit Resource Factories: Resources expose get_client() / get_driver() context managers rather than implicit lifecycle hooks
Secrets via EnvVar: resources.py uses EnvVar("NEO4J_PASSWORD") for secrets, not settings.py
Streaming I/O: io_managers.py uses sink_parquet() for LazyFrames (O(1) memory)

3. The `utils/` Folder: "The How-To" (Reusable Logic)

Contains domain-agnostic, reusable helpers that can be used across different projects.

Module	Responsibility	Key Functions
`network_helpers.py`	HTTP requests with exponential backoff, concurrency control, async generators	`make_async_request_with_retries()`, `run_tasks_concurrently()`, `yield_batches_concurrently()`
`io_helpers.py`	JSON/text file I/O, cache key generation, async file operations	`async_read_json_file()`, `async_write_json_file()`, `generate_cache_key()`
`data_transformation_helpers.py`	Text normalization, Unicode fixing (ftfy), chunking, sentence splitting	`chunk_text()`, `split_sentences()`, `normalize_and_clean_text()`, `strip_json_fences()`
`wikidata_helpers.py`	Wikidata SPARQL & Action API adapter with caching	`run_extraction_pipeline()`, `async_fetch_wikidata_entities_batch()`
`wikipedia_helpers.py`	Wikipedia API adapter with section parsing	`async_fetch_wikipedia_article()`, `parse_wikipedia_sections()`
`neo4j_helpers.py`	Generic Cypher execution with retry logic, graph construction, community detection	`execute_cypher()`, `clear_database()`, `build_igraph()`, `run_leiden_multilevel()`
`pinecone_helpers.py`	Pinecone vector operations, embedding via Inference API	`clear_index()`, `generate_embeddings_pinecone()`
`llm_helpers.py`	LLM inference (Gemini), tokenization (Tiktoken)	`generate_text_gemini_async()`, `load_tokenizer_only()`, `get_device()`

Design Rules for Utils:

No global config: Utils never import settings.py directly — configuration is passed as arguments
Dependency injection: API keys, paths, URLs, timeouts passed as function parameters
No domain logic: Schema definitions and entity type enums belong in assets, not utils
100% reusable: All helpers can be used across different projects without modification

The "Use" Hierarchy

┌─────────────────────────────────────────────────────────────────────┐
│  1. RESOURCES (defs/resources.py)                                   │
│     Provide raw connections (AsyncClient, Neo4j Driver, Pinecone)   │
│                              │                                      │
│                              ▼                                      │
│  2. UTILS (utils/*.py)                                              │
│     Use connections to perform specific actions                     │
│     (fetch_sparql_query, generate_text_gemini_async)                │
│                              │                                      │
│                              ▼                                      │
│  3. ASSETS (defs/assets/*.py)                                       │
│     Orchestrate Utils to achieve business goals                     │
│     (Extract Entities, Build Knowledge Graph, Embed Chunks)         │
│                              │                                      │
│                              ▼                                      │
│  4. I/O MANAGERS (defs/io_managers.py)                              │
│     Persist asset outputs to storage (Parquet, JSONL)               │
│                              │                                      │
│                              ▼                                      │
│  5. CHECKS (defs/checks.py)                                         │
│     Verify final output quality and data trust                      │
└─────────────────────────────────────────────────────────────────────┘

Data Architecture

The pipeline transforms raw data from Wikipedia and Wikidata into two optimized formats: a Knowledge Graph (for structural queries) and Vector-Ready Text Chunks (for semantic search), enriched with LLM-extracted entities and relationships.

Pipeline Flow

graph TD
    subgraph "Stage 1: Artist & Genre Discovery"
        A[Wikidata SPARQL] --> B(artist_index)
        B --> C[artists]
        C --> D[genres]
    end

    subgraph "Stage 2: Content Acquisition"
        C & D --> E[wikipedia_articles]
    end

    subgraph "Stage 3: Chunking & Embedding"
        E -->|Tiktoken| F[chunks]
        F -->|Pinecone Inference| G[vector_db]
        G -->|Upsert| H[(Pinecone — chunks)]
    end

    subgraph "Stage 4: LLM Information Extraction"
        F -->|Gemini| I[mentions]
        I --> J[aggregated_mentions]
        J -->|Wikidata API| K[wikidata_candidates]
        K -->|spaCy vectors| L[linked_entities]
        L --> M[graph_elements]
    end

    subgraph "Stage 5: Knowledge Graph Construction"
        M --> N[graph_summaries]
        N -->|Snowflake Arctic Embed| O[graph_db]
        O -->|Nodes & Relationships| P[(Neo4j Aura)]
    end

    subgraph "Stage 6: Community Detection — GraphRAG"
        P -->|Extract Graph| Q[detected_communities]
        Q -->|Leiden Algorithm| R[community_summaries]
        R -->|Gemini + Pinecone Inference| S[(Pinecone — summaries)]
    end

    style A fill:#e1f5fe
    style P fill:#c8e6c9
    style H fill:#fff3e0
    style S fill:#f3e5f5

Asset Dependency Graph

Asset	Input Dependencies	Output Type	I/O Manager
`artist_index`	None (Wikidata SPARQL)	`pl.LazyFrame`	Parquet
`artists`	`artist_index`	`list[Artist]`	Parquet
`genres`	`artists`	`pl.LazyFrame`	Parquet
`wikipedia_articles`	`artists`, `genres`	`list[Article]`	JSONL
`chunks`	`wikipedia_articles`	`pl.LazyFrame`	Parquet
`vector_db`	`chunks`	`MaterializeResult`	None (sink)
`mentions`	`chunks`	`pl.LazyFrame`	Parquet
`aggregated_mentions`	`mentions`	`pl.LazyFrame`	Parquet
`wikidata_candidates`	`aggregated_mentions`	`pl.LazyFrame`	Parquet
`linked_entities`	`wikidata_candidates`, `aggregated_mentions`	`pl.LazyFrame`	Parquet
`graph_elements`	`linked_entities`	`pl.LazyFrame`	Parquet
`graph_summaries`	`graph_elements`	`pl.LazyFrame`	Parquet
`graph_db`	`graph_summaries`	`MaterializeResult`	None (sink)
`detected_communities`	`graph_db`	`pl.LazyFrame`	Parquet
`community_summaries`	`detected_communities`, `artists`	`MaterializeResult`	None (sink)

Partitioning Strategy

The pipeline uses decade-based partitioning for the initial artist discovery phase:

DECADES_TO_EXTRACT = {
    "1930s": (1930, 1939),
    "1940s": (1940, 1949),
    # ...
    "2020s": (2020, 2029),
}

Benefits:

Parallelizes SPARQL queries across decades (10 concurrent runs)
Provides natural checkpointing (failed decade can be retried independently)
Keeps memory usage constant O(1) regardless of total dataset size

1. Unstructured Data (RAG Preparation)

We process Wikipedia articles to create a high-quality corpus for Retrieval-Augmented Generation (RAG).

Ingestion: Fetches full articles for every valid artist and genre found in the knowledge base.
Tokenization: Uses Tiktoken (cl100k_base) for accurate token counting.
Chunking: Splits text into ~500-token windows with sentence-boundary-aware overlap using a recursive character splitter.
Embedding: Generates 1024-dimensional embeddings via Pinecone Inference (llama-text-embed-v2).
Enrichment: Each chunk is stored with metadata (article name, Wikipedia URL, entity type, community assignment, PageRank score) to enable hybrid filtering during retrieval.

Field	Type	Description
`text`	String	The raw text chunk content.
`metadata`	JSON	Contextual tags: `article_id`, `article_name`, `wikipedia_uri`, `chunk_index`, `total_chunks`, `token_count`, `pagerank`, `entity_type`, `community_id`.

Download the dataset from Hugging Face: [TBD — dataset URL to be added]

2. Knowledge Graph (Neo4j)

We construct a knowledge graph by combining LLM-extracted entities and relationships with Wikidata-resolved identifiers. The graph enables precise multi-hop queries (e.g., "Find all artists who collaborated with members of Kraftwerk").

Entity Types (Node Labels)

Entities are extracted by Gemini and typed according to this schema:

Entity Type	Description	Example
`PERSON`	Individual musician, producer, DJ	"Richie Hawtin"
`GROUP`	Band, collective, duo	"Kraftwerk"
`RECORD_LABEL`	Music label	"Warp Records"
`ALBUM`	Studio album, compilation	"Selected Ambient Works 85-92"
`SONG`	Single, track	"Blue Monday"
`SOCIAL_MOVEMENT`	Cultural or musical movement	"Rave culture"
`GENRE`	Musical genre or sub-genre	"Minimal Techno"
`THEME`	Artistic theme or concept	"Futurism"
`INSTRUMENT`	Musical instrument or tool	"Roland TB-303"
`EVENT`	Festival, concert, happening	"Love Parade"
`CITY`	City or locality	"Berlin"
`COUNTRY`	Country	"Germany"

Relationship Types (Edges)

Relationship	Description
`INCLUDED_IN_ALBUM`	Song → Album
`RELEASED_SINGLE`	Artist → Song
`RELEASED_ALBUM`	Artist → Album
`RELEASED_ON_LABEL`	Artist/Album → Record Label
`ASSOCIATED_WITH_MOVEMENT`	Artist → Social Movement
`INSPIRED_BY_MOVEMENT`	Artist → Social Movement
`MEMBER_OF_GROUP`	Person → Group
`FORMER_MEMBER_OF`	Person → Group
`COLLABORATES_WITH`	Artist ↔ Artist
`HAS_GENRE`	Artist/Album/Song → Genre
`EXPLORES_THEME`	Artist/Song → Theme
`ORIGINATES_FROM`	Artist/Album → City/Country
`USES_INSTRUMENT`	Artist → Instrument
`PERFORMED_AT`	Artist → Event
`LOCATED_IN`	City → Country

Node Properties

All nodes share a base Entity label with these properties:

Property	Type	Description
`id`	String	Unique identifier (UNIQUE constraint)
`name`	String	Entity name (fulltext indexed)
`type`	String	Entity type (dynamic label)
`description`	String	Aggregated descriptions from extraction
`qid`	String	Wikidata QID (when resolved)
`aliases`	List[String]	Alternative names (fulltext indexed)
`mention_count`	Integer	How often the entity was mentioned across chunks
`chunk_ids`	List[String]	Source chunk provenance
`embedding`	Vector(384)	Snowflake Arctic Embed S vector

Relationship Properties

Property	Type	Description
`description`	String	Aggregated relationship descriptions
`mention_count`	Integer	Frequency across chunks
`chunk_ids`	List[String]	Source provenance
`embedding`	Vector(384)	Snowflake Arctic Embed S vector

Indexes

Constraints & Property Indexes:

Index	Type	Target
`entity_id_unique`	UNIQUE constraint	`Entity.id`
`entity_qid_index`	Property index	`Entity.qid`

Fulltext Index:

Index	Label	Properties	Purpose
`entityNameIndex`	Entity	name, aliases	Search entities by name or alias

Vector Indexes (Cosine Similarity, 384 dimensions):

Index	Target	Purpose
`entity_vector_index`	Entity nodes	Semantic search over entity embeddings
`rel_*_vector_index`	Relationships (per type)	Semantic search over relationship embeddings

Fulltext and vector indexes enable queries like:

-- Find entity by name
CALL db.index.fulltext.queryNodes("entityNameIndex", "Kraftwerk") YIELD node, score
RETURN node.name, node.type, node.description, score

-- Find semantically similar entities
CALL db.index.vector.queryNodes("entity_vector_index", 5, $queryVector)
YIELD node, score
RETURN node.name, node.type, score

Dataset Statistics:

Metric	Count
Articles processed	[TBD]
Text chunks	[TBD]
Graph nodes	[TBD]
Graph relationships	[TBD]
Communities detected	[TBD]

3. Vector Database (Pinecone)

The pipeline maintains two Pinecone indexes for different retrieval strategies:

Chunks Index

Index Name: chunks
Embedding Model: llama-text-embed-v2 (Pinecone Inference, 1024 dimensions)
Purpose: Semantic search over article text chunks

┌─────────────────────────────────────────────────────────────────────┐
│  text (str)                                                         │
│  └── Raw text chunk content                                         │
│                                                                     │
│  metadata (dict)                                                    │
│  ├── article_id: str (Wikidata QID)                                 │
│  ├── article_name: str                                              │
│  ├── wikipedia_uri: str                                             │
│  ├── chunk_index: int                                               │
│  ├── total_chunks: int                                              │
│  ├── token_count: int                                               │
│  ├── pagerank: float                                                │
│  ├── entity_type: str (PERSON, GROUP, etc.)                         │
│  └── community_id: int (L2 level assignment)                        │
└─────────────────────────────────────────────────────────────────────┘

Community Summaries Index

Index Name: community-summaries
Embedding Model: llama-text-embed-v2 (Pinecone Inference, 1024 dimensions)
Purpose: Global context retrieval via community-level summaries

┌─────────────────────────────────────────────────────────────────────┐
│  text (str)                                                         │
│  └── LLM-generated community summary                                │
│                                                                     │
│  metadata (dict)                                                    │
│  ├── community_id: int                                              │
│  ├── level: int (0, 1, or 2)                                        │
│  ├── rank: int (member count sum)                                   │
│  └── member_count: int                                              │
└─────────────────────────────────────────────────────────────────────┘

4. Community Detection (GraphRAG Pattern)

The pipeline implements the GraphRAG community detection pattern to enable global context retrieval. Instead of only retrieving individual text chunks, the system can answer high-level questions like "What characterizes German techno artists?" by retrieving community-level summaries.

Hierarchical Community Detection

The Leiden algorithm runs at three resolution levels to create a hierarchy:

Level	Resolution	Description	Example
L0	2.0	Fine-grained	"Berlin Minimal Techno"
L1	0.5	Medium	"German Electronic"
L2	0.1	Coarse	"European Dance Music"

LLM Summary Generation

Community summaries are generated using Google Gemini (gemini-2.5-flash-lite) with async concurrency (30 parallel requests):

Community context:
- Intra-community entities and their types
- Intra-community relationships with descriptions
- Geographic and genre distribution

→ "This community represents the core of European minimal techno, centered
   around the Berlin and Cologne scenes. Artists are characterized by
   stripped-down, hypnotic productions with precise rhythmic structures..."

This enables filtering by community level to retrieve summaries at different granularities, or combining them with individual text chunks for comprehensive answers.

Advanced Design Patterns

Dual Embedding Strategy

The pipeline uses two embedding models optimized for different purposes:

Purpose	Model	Dimensions	Location
Graph embeddings (Neo4j)	Snowflake Arctic Embed S	384	Local (GPU/MPS/CPU)
RAG embeddings (Pinecone)	llama-text-embed-v2	1024	Pinecone Inference API

This separation allows graph-level similarity search (entity disambiguation, relationship discovery) to run locally with a lightweight model, while RAG retrieval benefits from a larger, hosted model.

Lazy Evaluation & Streaming

The pipeline enforces lazy evaluation throughout to ensure O(1) memory usage:

# Assets return LazyFrames, not eager DataFrames
@asset
def artist_index(...) -> pl.LazyFrame:
    return clean_lf  # Never .collect() inside asset

# I/O Managers handle streaming writes
class PolarsParquetIOManager:
    def handle_output(self, context, obj):
        if isinstance(obj, pl.LazyFrame):
            obj.sink_parquet(path)  # Stream directly to disk

LLM Information Extraction Pipeline

Entity and relationship extraction follows a multi-stage pipeline:

Wikipedia chunks
      │
      ▼
┌─────────────────┐
│  mentions.py    │  Gemini extracts entities & relations per chunk
│  (Gemini LLM)   │  Output: ExtractionKnowledgeGraph per chunk
└────────┬────────┘
         │
         ▼
┌────────────────────────┐
│  aggregated_mentions   │  Merge entities across chunks by name
│  (Deduplication)       │  Accumulate mention_count, chunk_ids
└────────┬───────────────┘
         │
         ▼
┌────────────────────────┐
│  wikidata_candidates   │  Search Wikidata for each entity name
│  (Entity Resolution)   │  Returns candidate QIDs with metadata
└────────┬───────────────┘
         │
         ▼
┌────────────────────────┐
│  linked_entities       │  Disambiguate using spaCy word vectors
│  (spaCy Vectors)       │  Match extracted entities to Wikidata QIDs
└────────┬───────────────┘
         │
         ▼
┌────────────────────────┐
│  graph_elements        │  Merge per-document graphs into global
│  (Global Merge)        │  entities and relationships
└────────────────────────┘

Resource Factory Pattern

Resources expose explicit factory methods rather than implicit lifecycle hooks:

class Neo4jResource(ConfigurableResource):
    uri: str
    username: str
    password: str

    @contextmanager
    def get_driver(self, context) -> Generator[Driver, None, None]:
        driver = GraphDatabase.driver(self.uri, auth=(self.username, self.password))
        try:
            driver.verify_connectivity()
            yield driver
        finally:
            driver.close()

# Usage in asset
@asset
def graph_db(neo4j: Neo4jResource, ...):
    with neo4j.get_driver(context) as driver:
        # Use driver

Auto-Generated Polars Schemas

Domain models generate Polars-compatible schemas automatically:

# In models.py
ARTIST_SCHEMA = _generate_polars_schema(Artist)
GENRE_SCHEMA = _generate_polars_schema(Genre)

This ensures type consistency between msgspec Structs and Polars DataFrames without manual schema duplication.

Caching Strategy

Multi-level caching reduces API calls and enables resumable runs:

Cache Level	Location	Format	Purpose
Wikidata entities	`.cache/wikidata/{qid}.json`	JSON	Entity metadata
Wikipedia articles	`.cache/wikipedia/{qid}.txt`	Plain text	Article content
LLM extractions	`.cache/chunks/{hash}.json`	JSON	Gemini extraction results
Graph DB checkpoint	`.cache/graph_db_checkpoint/`	JSON	Resumable graph ingestion
Community summaries	`.cache/community_summaries/`	JSON	LLM-generated summaries

Retry Logic for Cloud Services

Neo4j Aura and Gemini API connections include retry logic with exponential backoff:

def _execute_with_retry(driver, query, max_retries=3, base_delay=2.0):
    for attempt in range(max_retries + 1):
        try:
            with driver.session() as session:
                return session.run(query).single()
        except (ServiceUnavailable, SessionExpired) as e:
            if attempt < max_retries:
                time.sleep(base_delay * (2 ** attempt))
            else:
                raise

Project Structure

graph_rag_part1/
├── src/
│   └── data_pipeline/
│       ├── definitions.py          # Dagster entry point
│       ├── settings.py             # Pydantic settings (paths, timeouts, model config)
│       ├── models.py               # Domain models (msgspec Structs & Pydantic)
│       ├── defs/
│       │   ├── assets/             # Business logic layer
│       │   │   ├── artist_index.py
│       │   │   ├── artists.py
│       │   │   ├── genres.py
│       │   │   ├── wikipedia_articles.py
│       │   │   ├── chunks.py
│       │   │   ├── vector_db.py
│       │   │   ├── mentions.py
│       │   │   ├── aggregated_mentions.py
│       │   │   ├── wikidata_candidates.py
│       │   │   ├── linked_entities.py
│       │   │   ├── graph_elements.py
│       │   │   ├── graph_summaries.py
│       │   │   ├── graph_db.py
│       │   │   ├── detected_communities.py
│       │   │   └── community_summaries.py
│       │   ├── resources.py        # Connection factories (Neo4j, Pinecone, Gemini, etc.)
│       │   ├── io_managers.py      # Parquet/JSONL persistence
│       │   ├── partitions.py       # Decade partitioning
│       │   └── checks.py           # Data quality gates
│       └── utils/                  # Reusable helpers (100% domain-agnostic)
│           ├── network_helpers.py
│           ├── io_helpers.py
│           ├── data_transformation_helpers.py
│           ├── wikidata_helpers.py
│           ├── wikipedia_helpers.py
│           ├── neo4j_helpers.py
│           ├── pinecone_helpers.py
│           └── llm_helpers.py
├── tests/                          # Mirrors src/ structure
│   └── data_pipeline/
│       ├── defs/
│       │   └── assets/
│       │       ├── test_*.py       # One test file per asset
│       │       └── conftest.py     # Shared fixtures
│       └── utils/
│           └── test_*.py           # One test file per helper
├── scripts/                        # Standalone CLI utilities
├── data_volume/                    # Local data & caches
│   ├── .cache/                     # API response caches
│   ├── datasets/                   # Materialized assets (Parquet/JSONL)
│   └── vector_db/                  # Local vector DB storage
├── Dockerfile                      # Python 3.13 + uv
├── docker-compose.yaml             # Development compose config
├── pyproject.toml                  # Dependencies & project metadata
├── uv.lock                         # Locked dependencies
└── .env                            # Secrets (not committed)

Getting Started

Prerequisites

Python 3.13+
uv (Astral's Python package manager)
A Neo4j Aura instance (free tier works)
A Pinecone account (free tier works)
A Google AI Studio API key (for Gemini)

Installation & Setup

Clone the repository:

git clone <repository-url>
cd graph_rag_part1

Install dependencies using uv: This will create a virtual environment and install all required packages from uv.lock.
```
uv sync
```

Configure Environment Variables: Create a .env file in the root directory:

# Neo4j Graph Database (Aura)
NEO4J_URI=neo4j+s://<your-instance-id>.databases.neo4j.io
NEO4J_USERNAME=neo4j
NEO4J_PASSWORD=<your-password>

# Pinecone Vector Database
PINECONE_API_KEY=<your-pinecone-api-key>

# Google Gemini (AI Studio)
GEMINI_API_KEY=<your-gemini-api-key>

# Optional: Dagster Environment (PROD or DEV)
DAGSTER_ENV=DEV

Running the Pipeline

Launch the Dagster development server:
```
uv run dg dev
```
Execute the Pipeline:
- Open http://localhost:3000 in your browser.
- Navigate to Assets → Global Asset Graph.
- Click "Materialize all" to run the full end-to-end pipeline.

Running Tests

uv run pytest

Security and Quality Checks

uv run ruff check .
uv run bandit -r src/

Data Quality Checks

The pipeline includes automated asset checks that validate data after materialization:

Check	Asset	Validation Rule
`check_artist_index_integrity`	`artist_index`	No null IDs/names, no duplicates

Name		Name	Last commit message	Last commit date
scripts		scripts
src		src
tests/data_pipeline		tests/data_pipeline
.dockerignore		.dockerignore
.gitignore		.gitignore
Dockerfile		Dockerfile
README.md		README.md
docker-compose.yaml		docker-compose.yaml
pyproject.toml		pyproject.toml
uv.lock		uv.lock

Latest commit

History

Folders and files

Repository files navigation