feat: KGRag - Knowledge Graph-Enhanced RAG with Mellea

[ydzhu98]

feat: add KG-RAG pipeline — Knowledge Graph-Enhanced RAG with Mellea

Adds a complete Knowledge Graph-enhanced Retrieval-Augmented Generation (KG-RAG) library and end-to-end example to mellea-contribs, inspired by Bidirection and rewritten from the ground up to follow mellea-contribs design patterns.

---

What this PR adds

Core library: mellea_contribs/kg/

The library is structured around a four-layer architecture that cleanly separates concerns from user-facing orchestration down to the database:

- Layer 1 — Application Orchestration 
 orchestrate_qa_retrieval()  ·  orchestrate_kg_update() 
 KGPreprocessor  ·  KGEmbedder 

- Layer 2 — Components & Query Building 
 CypherQuery  ·  GraphResult  ·  GraphTraversal 

- Layer 3 — LLM-Guided Logic  (@generative functions) 
  QA (8):     break_down_question · extract_topic_entities
             align_topic_entities · prune_relations 
             prune_triplets · evaluate_knowledge_sufficiency 
             validate_consensus · generate_direct_answer 
 Update (5): extract_entities_and_relations · align_entity  
             decide_entity_merge · align_relation · decide_merge

- Layer 4 — Backend Abstraction  
 GraphBackend (abstract)  ·  Neo4jBackend  ·  MockGraphBackend

Layer 1 provides two high-level entry points. orchestrate_qa_retrieval() implements Think-on-Graph multi-hop reasoning: it breaks the question into independent solving routes, aligns topic entities against the KG, traverses and prunes relation paths at each hop, and reaches a consensus answer across routes. orchestrate_kg_update() drives incremental KG construction: it extracts entities and relations from a document, aligns each against existing KG nodes, and merges or creates accordingly.

Layer 2 encapsulates query construction (CypherQuery, GraphTraversal) and result formatting (GraphResult) so that Layer 3 functions never touch raw query strings.

Layer 3 contains all LLM decision-making as @generative functions. Each function has a single, well-typed responsibility (e.g., prune_relations() returns RelevantRelations; evaluate_knowledge_sufficiency() returns EvaluationResult). No hand-crafted prompt assembly — Mellea's generative framework handles grounding and output parsing.

Layer 4 defines GraphBackend, a database-agnostic interface. The production Neo4jBackend and the zero-dependency MockGraphBackend share the same API, so every Layer 1–3 call is fully testable without infrastructure.

Module	Purpose
base.py	GraphNode, GraphEdge, GraphPath dataclasses
models.py	Entity and Relation Pydantic models
preprocessor.py	KGPreprocessor abstract base (Layer 1)
embedder.py	KGEmbedder — LiteLLM batch embedding + cosine similarity search
kgrag.py	KGRag — Think-on-Graph multi-hop QA (Layer 1)
graph_dbs/neo4j.py	Production Neo4j backend (Layer 4)
graph_dbs/mock.py	In-memory mock backend (Layer 4)
components/	Query, result, traversal components (Layer 2)
utils/	session_manager, data_utils, progress, eval_utils

End-to-end example: docs/examples/kgrag/

A five-stage pipeline over the CRAG movie benchmark (64K+ movies, 373K+ persons, 1M+ relations):

Step	Script	Description
0	create_tiny_dataset.py	Slice a small test dataset
1	run_kg_preprocess.py	Load predefined movie/person data into Neo4j
2	run_kg_embed.py	Compute and store entity/relation embeddings
3	run_kg_update.py	Extract entities from documents and merge into KG
4	run_qa.py	Answer questions via Think-on-Graph multi-hop retrieval
5	run_eval.py	Score predictions (exact → fuzzy → LLM judge); report CRAG metrics

Domain-specific components (Movie/Person/Award models, preprocessor hints, LLM prompt formatters) live under models/, preprocessor/, and rep/ to illustrate how to extend the library for a new domain.

All scripts load credentials from .env via python-dotenv and support --mock for local testing without Neo4j or an LLM endpoint.

Tests: test/kg/

95 tests covering all four layers: Pydantic models, Layer 3 generative functions, the mock backend, the Neo4j backend (skipped unless NEO4J_URI is set), and all utility modules.

---

Relationship to upstream

Bidirection is a research project implementing bidirectional graph traversal for temporal KG-RAG over movie domain data. This PR takes inspiration from that work and:

• Extracts the generic KG infrastructure into mellea_contribs/kg/ as a reusable library
• Moves domain-specific logic (movie models, preprocessor, representations) into the docs/examples/kgrag/ example directory to keep the library domain-agnostic
• Replaces ad-hoc Neo4j driver calls with the GraphBackend abstraction
• Structures all LLM reasoning as @generative Layer 3 functions following Mellea's framework conventions
• Aligns configuration classes (QAConfig, UpdateConfig, EmbeddingConfig) with mellea's Pydantic config patterns
• Adds the full test suite

---

Prerequisites

1. Start Neo4j

docker run \
    --name neo4j \
    -p 7474:7474 -p 7687:7687 \
    -e NEO4J_AUTH=neo4j/password \
    -e NEO4J_PLUGINS='["apoc"]' \
    neo4j:latest

2. Configure credentials

cd docs/examples/kgrag
cp .env_template .env
# Edit .env: set API_BASE, API_KEY, MODEL_NAME, NEO4J_PASSWORD

3. Acquire the dataset

The pipeline uses two data sources from the CRAG benchmark:

Structured KG databases (used by Step 1 — preprocessing):

# Clone CRAG (requires Git LFS for the database files)
git lfs install
git clone https://github.com/facebookresearch/CRAG.git

# Copy the movie mock API databases into the example dataset directory
cp -r CRAG/mock_api/movie docs/examples/kgrag/dataset/movie

This populates:

File	Size	Description
dataset/movie/movie_db.json	~181 MB	Movie entities (title, cast, awards, …)
dataset/movie/person_db.json	~44 MB	Person entities (actors, directors, …)

JSONL question dataset (used by Steps 3–5 — update, QA, eval):

The crag_movie_dev.jsonl.bz2 file (~140 MB) contains question/answer pairs with associated search results. Contact the CRAG project maintainers for access, then place it at docs/examples/kgrag/dataset/crag_movie_dev.jsonl.bz2.

Once you have the full dataset, create a small slice for quick testing:

cd docs/examples/kgrag/scripts
python create_tiny_dataset.py   # produces dataset/crag_movie_tiny.jsonl.bz2

---

How to run

cd docs/examples/kgrag/scripts

# Full pipeline on the tiny test dataset (~10 docs):
bash run.sh

# Individual steps:
bash run.sh --tiny 3 4 5    # update + QA + eval only
bash run.sh --full 4 5      # QA + eval on full dataset

# No database / no LLM (mock mode, no data files needed):
python run_kg_update.py --dataset ../dataset/crag_movie_tiny.jsonl.bz2 --mock
python run_qa.py --dataset ../dataset/crag_movie_tiny.jsonl.bz2 --mock
python run_eval.py --input ../output/qa_results.jsonl --mock

---

Test done

unit test

pytest test/kg/ -v

manual test

bash run.sh

[mergify]

Merge Protections

Your pull request matches the following merge protections and will not be merged until they are valid.

🟢 Enforce conventional commit

Wonderful, this rule succeeded.

Make sure that we follow https://www.conventionalcommits.org/en/v1.0.0/

• title ~= ^(fix|feat|docs|style|refactor|perf|test|build|ci|chore|revert|release)(?:\(.+\))?:

[ydzhu98]

Overall, the scripts are not fully utilize the logic implemented in the mellea_contribs/kg. As a result, there is a lot of duplicated logic and making the files under scripts huge. They are supposed to be small so that people can easier adapt to other datasets. Please think about how to merge them by either updating the logic implimented in mellea_contribs/kg, or add some additional logic into it.

[ydzhu98]

Let's remove this README.md and include the correpsonding information inside the project readme under examples/kgrag.

[ydzhu98]

base_format_kg_context is not used here. Should we use it or remove the imports?

[ydzhu98]

This file should be removed.

[ydzhu98]

Try to used more predefined models and functions.

[ydzhu98]

Try to used more predefined models and functions.