RAG Chunk Comparator — Java source

// RAG Chunk Comparator — chunk one document three ways and report the stats
// that matter for retrieval.
//
// Language: Java (17+, zero dependencies)
// Port of src/lib/ragChunkComparator.ts (the canonical TypeScript
// implementation). Field names stay camelCase to match the TS surface.
// Tool page: https://dev.cosmolabs.org/tools/rag-chunk-comparator

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public final class RagChunkComparator {

    public enum ChunkStrategy { FIXED, SENTENCE, MARKDOWN }

    /** Target chunk size in tokens + fixed-strategy overlap. */
    public record ChunkOptions(int sizeTokens, int overlapTokens) {
        public ChunkOptions(int sizeTokens) { this(sizeTokens, 0); }
    }

    public record Chunk(int index, String text, int tokens, String heading) {
        public Chunk(int index, String text, int tokens) { this(index, text, tokens, null); }
    }

    /** Share of boundaries on sentence ends is 0-1; 1 when no boundaries. */
    public record StrategyStats(int count, int minTokens, int maxTokens, int avgTokens,
                                double sentenceBoundaryShare) {}

    public record StrategyResult(ChunkStrategy strategy, List<Chunk> chunks, StrategyStats stats) {}

    private static final Pattern SENTENCE_SPLIT =
            Pattern.compile("(?<=[.!?]) +");
    private static final Pattern ENDS_SENTENCE =
            Pattern.compile("[.!?][\"'\\)\\]]?$");
    private static final Pattern HEADING =
            Pattern.compile("^(#{1,6})\\s+(.*)$");

    /**
     * The {@code type: 'prose'} path of the tokenEstimator, inlined: every
     * non-empty line costs max(1, round(length / 4)) tokens; empty text is 0.
     */
    private static int tok(String s) {
        if (s == null || s.isEmpty()) return 0;
        int tokens = 0;
        for (String line : s.split("\n", -1)) {
            if (!line.isEmpty()) {
                long per = Math.round(line.length() / 4.0);
                tokens += Math.max(1, (int) per);
            }
        }
        return tokens;
    }

    /** Split on sentence enders followed by whitespace or end of text. */
    public static List<String> splitSentences(String text) {
        String collapsed = text.replaceAll("\\s+", " ").trim();
        List<String> out = new ArrayList<>();
        for (String s : SENTENCE_SPLIT.split(collapsed)) {
            if (!s.isEmpty()) out.add(s);
        }
        return out;
    }

    private static boolean endsSentence(String s) {
        return ENDS_SENTENCE.matcher(s.trim()).find();
    }

    /**
     * Greedy character accumulation to a token target (overlapping allowed).
     * @throws IllegalArgumentException on impossible options (the TS RangeError contract).
     */
    public static List<Chunk> chunkFixed(String text, ChunkOptions opts) {
        if (opts.sizeTokens() <= 0)
            throw new IllegalArgumentException("sizeTokens must be > 0");
        if (opts.overlapTokens() < 0 || opts.overlapTokens() >= opts.sizeTokens())
            throw new IllegalArgumentException("overlapTokens must be in [0, sizeTokens)");
        String clean = text.trim();
        if (clean.isEmpty()) return new ArrayList<>();
        // ~4 chars per prose token: step by tokens, verify with the estimator.
        int charStep = Math.max(1, (int) Math.round(opts.sizeTokens() * 4.0));
        int overlapChars = (int) Math.round(opts.overlapTokens() * 4.0);
        List<Chunk> chunks = new ArrayList<>();
        int start = 0;
        while (start < clean.length()) {
            int end = Math.min(start + charStep, clean.length());
            // Prefer cutting at whitespace near the target.
            if (end < clean.length()) {
                int cut = clean.lastIndexOf(' ', end);
                if (cut > start) end = cut;
            }
            String piece = clean.substring(start, end).trim();
            if (!piece.isEmpty()) chunks.add(new Chunk(chunks.size(), piece, tok(piece)));
            if (end >= clean.length()) break;
            start = Math.max(end - overlapChars, start + 1);
        }
        return chunks;
    }

    /**
     * Group whole sentences up to the token target; boundaries never split a
     * sentence. A single sentence larger than the target becomes its own chunk.
     */
    public static List<Chunk> chunkBySentences(String text, ChunkOptions opts) {
        if (opts.sizeTokens() <= 0)
            throw new IllegalArgumentException("sizeTokens must be > 0");
        List<String> sentences = splitSentences(text);
        if (sentences.isEmpty()) return new ArrayList<>();
        List<Chunk> chunks = new ArrayList<>();
        List<String> current = new ArrayList<>();
        int[] currentTokens = {0};
        // A small closure over mutable state, mirroring the TS flush().
        Runnable flush = () -> {
            if (current.isEmpty()) return;
            String piece = String.join(" ", current);
            chunks.add(new Chunk(chunks.size(), piece, tok(piece)));
            current.clear();
            currentTokens[0] = 0;
        };
        for (String sentence : sentences) {
            int t = tok(sentence);
            if (currentTokens[0] > 0 && currentTokens[0] + t > opts.sizeTokens()) flush.run();
            current.add(sentence);
            currentTokens[0] += t;
        }
        flush.run();
        return chunks;
    }

    /**
     * Split on markdown headings; oversized sections fall back to sentence
     * grouping, and every chunk carries the section heading.
     */
    public static List<Chunk> chunkMarkdown(String text, ChunkOptions opts) {
        if (opts.sizeTokens() <= 0)
            throw new IllegalArgumentException("sizeTokens must be > 0");
        record Section(String heading, List<String> body) {}
        List<Section> sections = new ArrayList<>();
        String[] pendingHeading = {null};
        List<String> pendingBody = new ArrayList<>();
        for (String line : text.split("\n", -1)) {
            Matcher m = HEADING.matcher(line);
            if (m.matches()) {
                if (!pendingBody.isEmpty())
                    sections.add(new Section(pendingHeading[0], new ArrayList<>(pendingBody)));
                pendingHeading[0] = m.group(2).trim();
                pendingBody.clear();
            } else {
                pendingBody.add(line);
            }
        }
        if (!pendingBody.isEmpty())
            sections.add(new Section(pendingHeading[0], new ArrayList<>(pendingBody)));

        List<Chunk> chunks = new ArrayList<>();
        for (Section section : sections) {
            String clean = String.join("\n", section.body()).trim();
            if (clean.isEmpty()) continue;
            String whole = section.heading() != null
                    ? "# " + section.heading() + "\n" + clean : clean;
            if (tok(whole) <= opts.sizeTokens()) {
                chunks.add(new Chunk(chunks.size(), whole, tok(whole), section.heading()));
                continue;
            }
            for (Chunk c : chunkBySentences(clean, opts)) {
                chunks.add(new Chunk(chunks.size(), c.text(), c.tokens(), section.heading()));
            }
        }
        return chunks;
    }

    private static StrategyResult statsFor(ChunkStrategy strategy, List<Chunk> chunks) {
        int count = chunks.size();
        int min = 0, max = 0;
        long sum = 0;
        if (count > 0) {
            min = Integer.MAX_VALUE;
            for (Chunk c : chunks) {
                min = Math.min(min, c.tokens());
                max = Math.max(max, c.tokens());
                sum += c.tokens();
            }
        }
        int avg = count > 0 ? (int) Math.round((double) sum / count) : 0;
        int boundaries = 0, ending = 0;
        for (int i = 0; i + 1 < count; i++) {
            boundaries++;
            if (endsSentence(chunks.get(i).text())) ending++;
        }
        double share = boundaries > 0 ? (double) ending / boundaries : 1.0;
        return new StrategyResult(strategy, chunks,
                new StrategyStats(count, min, max, avg, share));
    }

    /** Run all three strategies over one document and report comparable stats. */
    public static Map<ChunkStrategy, StrategyResult> compareStrategies(String text, ChunkOptions opts) {
        Map<ChunkStrategy, StrategyResult> out = new LinkedHashMap<>();
        out.put(ChunkStrategy.FIXED, statsFor(ChunkStrategy.FIXED, chunkFixed(text, opts)));
        out.put(ChunkStrategy.SENTENCE, statsFor(ChunkStrategy.SENTENCE, chunkBySentences(text, opts)));
        out.put(ChunkStrategy.MARKDOWN, statsFor(ChunkStrategy.MARKDOWN, chunkMarkdown(text, opts)));
        return out;
    }

    private RagChunkComparator() {}
}