DEV Community

Cover image for Building a RAPTOR-Inspired Hierarchical Retrieval System with PostgreSQL, pgvector, and Apache AGE
kai wen ng
kai wen ng

Posted on Edited on

Building a RAPTOR-Inspired Hierarchical Retrieval System with PostgreSQL, pgvector, and Apache AGE

I've been working on improving my personal knowledge system, particularly how documents are chunked, organized, and retrieved.


Traditional RAG pipelines typically follow a straightforward process:

Document → Chunking → Embedding → Vector Database → Retrieval

While effective, flat chunking has limitations. Individual chunks often lack broader context, and relationships between different sections of a document are not explicitly represented.

To address this, I started implementing a hierarchical chunking strategy inspired by RAPTOR (Recursive Abstractive Processing for Tree-Organized Retrieval).

Why RAPTOR?

RAPTOR introduces hierarchical abstractions of document content rather than relying exclusively on individual chunks.

Some interesting advantages:

  1. Hierarchical information: Retrieve information at different levels of abstraction, from individual passages to higher-level summaries.
  2. Context reconstruction: Related chunks can be grouped together, while their original document positions allow neighboring passages to be reconstructed when necessary.
  3. Overlapping clusters: A single chunk can belong to multiple semantic groups, allowing its information to contribute to different summaries.

However, these benefits introduce additional computational overhead.

Compared to conventional vector-based retrieval, we now need dimensionality reduction, clustering, recursive hierarchy construction, LLM summarization, and additional embeddings.

Here's how I approached the implementation.

1. Initial Document Chunking

Regardless of how sophisticated the retrieval strategy becomes, we still need an initial chunking process.

I used PyMuPDF4LLM to extract and chunk document content, with the following configuration:

  • Minimum chunk size: 200 tokens
  • Maximum chunk size: 800 tokens

The goal is to produce chunks containing enough meaningful context without becoming unnecessarily large.

Each chunk is stored in PostgreSQL, alongside its metadata, original position, and vector embedding.

This forms the foundation of the hierarchy.

2. Semantic Clustering with Gaussian Mixture Models

Instead of assigning chunks to rigid, mutually exclusive clusters, I chose a Gaussian Mixture Model (GMM).

Unlike hard clustering methods such as K-Means, GMM provides a probability distribution over clusters for each embedding.

For example, a chunk could have the following membership probabilities:

Cluster Probability
Cluster A 0.65
Cluster B 0.30
Cluster C 0.05

With a probability threshold of 0.1, the chunk belongs to both Cluster A and Cluster B.

This is particularly useful for technical documents, where one passage might discuss multiple related concepts.

My implementation uses:

model = GaussianMixture(
    n_components=5,
    covariance_type="full",
    reg_covar=1e-6,
    n_init=3,
    max_iter=300,
    random_state=224,
)

model.fit(embeddings)

probabilities = model.predict_proba(embeddings)
assignments = probabilities > 0.1
Enter fullscreen mode Exit fullscreen mode

I also ensure that every chunk belongs to at least one cluster by assigning it to its highest-probability cluster if no probability exceeds the threshold.

This allows overlapping semantic relationships without leaving chunks unassigned.

Reducing Computational Costs with UMAP

One obvious problem with clustering thousands of high-dimensional embeddings is computational cost.

This becomes especially expensive when using GMM with full covariance matrices.

To reduce the dimensionality before clustering, I introduced UMAP (Uniform Manifold Approximation and Projection).

reducer = UMAP(
    n_neighbors=30,
    n_components=100,
    metric="cosine",
    min_dist=0.0,
    random_state=224,
    n_jobs=1,
)

reduced_embeddings = reducer.fit_transform(embeddings)
Enter fullscreen mode Exit fullscreen mode

UMAP attempts to preserve important neighborhood relationships while reducing dimensionality.

This adds a preprocessing step, but reduces the dimensionality of the subsequent GMM fitting problem.

3. Recursively Constructing the Hierarchy

Once the initial clusters are generated, I recursively subdivide clusters that exceed a predefined context budget.

The process is relatively straightforward:

  1. Group the original document chunks using GMM.
  2. Calculate the total content length of each cluster.
  3. If a cluster exceeds the threshold, run GMM again on its members.
  4. Create child clusters and repeat until the groups fit within the budget.

Currently, I use 8,000 characters as the maximum cluster content length.

This is a temporary heuristic to keep the summarization workload manageable for my local qwen3.5:4b model. Eventually, I plan to replace it with token-aware budgeting.

The resulting hierarchy looks something like this:

Book
├── Cluster A (Level 0)
│   ├── Cluster A1 (Level 1)
│   │   ├── Chunk 1
│   │   └── Chunk 2
│   └── Cluster A2 (Level 1)
│       ├── Chunk 3
│       └── Chunk 4
│
└── Cluster B (Level 0)
    ├── Chunk 2
    ├── Chunk 5
    └── Chunk 6
Enter fullscreen mode Exit fullscreen mode

Notice that Chunk 2 belongs to multiple clusters.

This is intentional. Overlapping memberships allow information to participate in multiple semantic groupings.

One important distinction: my implementation uses top-down recursive splitting followed by bottom-up summarization. The original RAPTOR approach builds its hierarchy bottom-up through iterative clustering and summarization.

Therefore, this is a RAPTOR-inspired variation rather than an exact reproduction.

4. Combining PostgreSQL, pgvector, and Apache AGE

I wanted to combine the strengths of relational storage, vector similarity search, and graph traversal without introducing an entirely separate database system.

My current stack uses three PostgreSQL components:

Technology Responsibility
PostgreSQL Store original chunks, summaries, metadata, and document positions
pgvector Store embeddings and support approximate nearest-neighbor search
Apache AGE Represent and traverse hierarchical relationships

The graph structure uses relationships such as:

(Book)-[:CONTAINS]->(Cluster)
(Cluster)-[:CONTAINS]->(Cluster)
(Cluster)-[:SUMMARIZES]->(Chunk)
Enter fullscreen mode Exit fullscreen mode

Each graph node references its corresponding record in PostgreSQL through a UUID.

This means the graph doesn't need to duplicate the original text.

For example, I can use Apache AGE to traverse the descendants of a particular cluster, retrieve their UUIDs, and join those identifiers with the PostgreSQL chunks table.

Bridging Graph Traversal with Relational Joins

One interesting advantage of using Apache AGE is that it runs directly inside PostgreSQL.

This allows me to combine Cypher graph traversal with traditional SQL JOIN operations in a single query.

During hierarchy construction, each cluster is assigned a UUID and persisted in PostgreSQL. Its relationships with parent clusters and chunks are then recorded in Apache AGE.

When generating summaries, I need to reconstruct the relevant text from those relationships.

Instead of duplicating the original text inside graph nodes, I retrieve the corresponding chunk UUIDs through Cypher and join them back to the PostgreSQL chunks table.

Here's the core idea from my SQLAlchemy repository:

SELECT
    (result.source_id::text)::UUID AS cluster_id,
    ch.meta_data,
    ch.raw_text AS content
FROM cypher('raptor_graph', $$
    MATCH (book:Book {id: 'BOOK_UUID'})
        -[*1..]->(c:Cluster {level: 1})
        -[r]->(d)
    RETURN DISTINCT
        c.id,
        d.id
$$) AS result(
    source_id agtype,
    target_id agtype
)
JOIN public.chunks ch
    ON ch.id = (result.target_id::text)::UUID
ORDER BY cluster_id, ch.position;
Enter fullscreen mode Exit fullscreen mode

The execution flow is:

Cluster ID → Cypher traversal → Chunk UUIDs → PostgreSQL JOIN → Original text → LLM summarization

A few interesting details:

  • Cypher handles the relationships: Locate clusters at a particular hierarchy level and retrieve their connected nodes.
  • SQL handles the data: Convert the returned agtype identifiers into PostgreSQL UUIDs and join them with the original chunk records.
  • SQLAlchemy handles the result: Group the returned text by cluster UUID to construct the context for LLM summarization.

The same database engine can execute both graph traversal and relational operations.

This also avoids duplicating large amounts of text in the graph. The graph stores relationships and identifiers, while PostgreSQL remains the source of truth for document content.

The graph determines which pieces of information are related, while PostgreSQL stores the actual content.

For my personal knowledge system, this is particularly useful because I can combine semantic similarity search through pgvector, hierarchical traversal through Apache AGE, and relational joins through PostgreSQL without operating three separate database services.

5. Recursive Summarization and Embedding

After constructing the hierarchy, the next step is generating summaries.

I process the hierarchy from the deepest level toward the root.

The intended abstraction process is:

Original Chunks
      ↓
Leaf Cluster Summaries
      ↓
Intermediate Cluster Summaries
      ↓
Root Cluster Summaries
Enter fullscreen mode Exit fullscreen mode

The summarization itself runs locally through Ollama:

response = await ollama.AsyncClient().chat(
    model="qwen3.5:4b",
    messages=[
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": context},
    ],
    think=False,
    options={"temperature": 0.5},
)
Enter fullscreen mode Exit fullscreen mode

For a fully recursive summarization pipeline, higher-level summaries should consume their child summaries rather than repeatedly concatenating all original descendant chunks.

This allows information to be progressively compressed as we move up the hierarchy.

Once generated, these summaries are stored in PostgreSQL and embedded again.

Consequently, the retrieval system can search both:

  • Original chunk embeddings, which preserve detailed information.
  • Summary embeddings, which represent broader concepts across multiple chunks.

The hierarchy then provides a path from retrieved summaries back to their underlying source material.

6. Trade-offs and Remaining Challenges

There are several challenges I still need to address.

Computational overhead: UMAP, recursive GMM clustering, LLM summarization, and additional embedding generation make indexing significantly more expensive than conventional RAG.

Overlapping memberships: Allowing chunks to belong to multiple clusters improves representation flexibility, but can also increase storage and summarization costs. Recursive splitting also needs safeguards to ensure child clusters actually become smaller.

Context budgeting: The current 8,000-character threshold is a temporary solution. Token-based limits would provide more predictable control over model input length.

Retrieval evaluation: A more sophisticated hierarchy does not automatically guarantee better retrieval. I still need to benchmark retrieval accuracy, latency, and indexing cost against a conventional flat-chunk baseline.

There are also opportunities to optimize graph construction, reuse shared chunk nodes, and reduce redundant summarization.

Final Thoughts

What started as a simple document chunking experiment has gradually evolved into a hybrid retrieval architecture combining:

PyMuPDF4LLM + UMAP + Gaussian Mixture Models + Ollama + PostgreSQL + pgvector + Apache AGE

The most interesting aspect for me is moving beyond treating documents as isolated vectors.

Instead, information can be represented at multiple levels of abstraction, while preserving relationships between summaries and their underlying source chunks.

It's definitely more computationally expensive than traditional chunking and embedding, but the architecture opens up interesting possibilities for hierarchical retrieval and context reconstruction.

What's Next: Embeddings and Search Ranking

Building the hierarchy is only one part of the retrieval pipeline. The next challenge is determining how to retrieve and rank the most relevant information efficiently.

In my next post, I'll explore embedding models, vector similarity search, approximate nearest-neighbor algorithms such as HNSW and IVF, and search ranking techniques.

I'll also investigate how to combine semantic similarity with hierarchical relationships, and whether reranking can meaningfully improve retrieval quality.

Ultimately, storing information effectively is only half the problem. Finding the right information is just as important.

Top comments (0)