Representative Centroid Selection for Long-Context RAG

import numpy as np
from sklearn.cluster import KMeans

def get_representative_indices(embeddings, k=5, seed=0):
    """Indices of k real documents: the member of each cluster closest to its centroid."""
    embeddings = np.asarray(embeddings)
    k = min(k, len(embeddings))
    kmeans = KMeans(n_clusters=k, init="k-means++", n_init=10, random_state=seed)
    labels = kmeans.fit_predict(embeddings)
    # cluster_centers_ are synthetic averages, not documents, so pick the
    # closest actual member of each cluster instead.
    picks = []
    for c, center in enumerate(kmeans.cluster_centers_):
        members = np.flatnonzero(labels == c)
        dists = np.linalg.norm(embeddings[members] - center, axis=1)
        picks.append(members[np.argmin(dists)])
    return np.array(picks)

Why it matters: Instead of the top-K most similar (often near-duplicate) passages, you get K passages that cover different regions of the candidate set, which reduces redundancy in a long-context prompt. Run it on the top-N retrieved candidates, not the whole corpus, so relevance is preserved. cluster_centers_ alone is not enough: centroids are averages, not passages you can put in a prompt.