Embedding Chunk Planner — C# source

// Embedding Chunk Planner — pure chunking math for RAG pipelines.
//
// Language: C# (C# 12 / .NET 8, standard library only)
// Source:   CosmoDev polyglot showcase port of the Embedding Chunk Planner
//           tool, ported from src/lib/embeddingPlanner.ts (the canonical
//           TypeScript implementation).
// Tool page: https://dev.cosmolabs.org/tools/embedding-chunk-planner
// License:  display source — part of CosmoDev's polyglot tool pages.
//
// Design goals:
//   - Pure + deterministic; never throws.
//   - Functionally equivalent to the TS reference: same inputs -> same outputs.
//   - Self-contained: BCL only (no NuGet packages). The model price table is
//     inlined below, mirrored from src/lib/ai/embeddings.ts — prices NEVER
//     live in the planner itself.
//
// Behavior (mirrors the TS source exactly):
//   - chunkSize <= 0 or totalTokens <= 0 -> ChunkPlan.Zero (nothing to embed).
//   - Negative overlap is treated as 0; overlap then clamps to at most
//     chunkSize / 2 so consecutive chunks always advance.
//   - chunks = Max(1, ceil((totalTokens - overlap) / (chunkSize - overlap)))
//     — a tiny document still yields one chunk.

using System;
using System.Collections.Generic;
using System.Linq;

namespace CosmoDev.Polyglot;

/// <summary>
/// One embedding model's offered dimensions (ascending, Matryoshka shortening
/// included) and pricing: USD per 1M input tokens.
/// </summary>
public sealed record EmbeddingModel(string Id, string Vendor, long[] Dims, double InputPerM);

/// <summary>How a document splits into overlapping chunks.</summary>
public sealed record ChunkPlan(long Chunks, long TotalTokensWithOverlap, long OverheadTokens)
{
    /// <summary>
    /// The "nothing to embed" plan the TS source returns for zero/negative
    /// input or a non-positive chunk size.
    /// </summary>
    public static ChunkPlan Zero { get; } = new(0, 0, 0);
}

/// <summary>
/// Chunk plan plus pricing for one embedding call. The three chunk fields are
/// flattened in (the TS <c>...plan</c> spread) so the record reads like the TS
/// <c>EmbeddingPlan extends ChunkPlan</c>.
/// </summary>
public sealed record EmbeddingPlan(
    long Chunks,
    long TotalTokensWithOverlap,
    long OverheadTokens,
    EmbeddingModel Model,
    long Vectors, // one vector per chunk
    double Cost   // USD: TotalTokensWithOverlap / 1e6 * Model.InputPerM
);

public static class EmbeddingChunkPlanner
{
    /// <summary>
    /// Embedding model price table — the SSOT for pricing, mirrored from
    /// src/lib/ai/embeddings.ts. Refresh both files together.
    /// </summary>
    public static readonly IReadOnlyList<EmbeddingModel> EmbeddingModels =
    [
        new("text-embedding-3-small", "OpenAI", [512, 1536], 0.02),
        new("text-embedding-3-large", "OpenAI", [256, 1024, 3072], 0.13),
        new("embed-english-v3.0", "Cohere", [512, 1024, 1536], 0.1),
        new("voyage-3-lite", "Voyage AI", [512, 1024], 0.02),
    ];

    /// <summary>Default knobs: 512-token chunks, 64-token overlap (TS DEFAULT_CHUNK_OPTIONS).</summary>
    public const long DefaultChunkSize = 512;
    public const long DefaultOverlap = 64;

    /// <summary>Look up an embedding model by id. Returns <c>null</c> for unknown ids.</summary>
    public static EmbeddingModel? GetEmbeddingModel(string id) =>
        EmbeddingModels.FirstOrDefault(m => m.Id == id);

    /// <summary>
    /// Plan how <paramref name="totalTokens"/> split into overlapping chunks.
    /// <paramref name="chunkSize"/> / <paramref name="overlap"/> are each
    /// independently optional — <c>null</c> falls back to the 512 / 64 default.
    /// </summary>
    public static ChunkPlan PlanChunks(long totalTokens, long? chunkSize = null, long? overlap = null)
    {
        long cs = chunkSize ?? DefaultChunkSize;
        long overlapRaw = overlap ?? DefaultOverlap;

        if (cs <= 0 || totalTokens <= 0)
        {
            return ChunkPlan.Zero;
        }

        // Math.Min(Math.Max(overlap, 0), chunkSize / 2) — the TS clamp. Overlap
        // that large would never advance, so consecutive chunks always gain at
        // least half a chunk. (cs >= 1 here, so cs - overlap is never zero.)
        long ov = Math.Min(Math.Max(overlapRaw, 0), cs / 2);

        // Float division + ceil mirrors TS's Math.ceil exactly (a tiny document
        // lands the quotient just below zero; ceil brings it to 0 and Max(1, ..)
        // lifts it back to one chunk).
        long chunks = Math.Max(1L, (long)Math.Ceiling((double)(totalTokens - ov) / (cs - ov)));

        long totalWithOverlap = totalTokens + (chunks - 1) * ov;
        return new ChunkPlan(chunks, totalWithOverlap, totalWithOverlap - totalTokens);
    }

    /// <summary>
    /// Chunk a document AND price its embedding for <paramref name="modelId"/>
    /// at <paramref name="dims"/> dimensions. Unknown model, or dims the model
    /// does not offer -> <c>null</c>.
    /// </summary>
    public static EmbeddingPlan? PlanEmbedding(long totalTokens, string modelId, long dims,
        long? chunkSize = null, long? overlap = null)
    {
        EmbeddingModel? model = GetEmbeddingModel(modelId);
        if (model is null || !model.Dims.Contains(dims))
        {
            return null;
        }
        ChunkPlan plan = PlanChunks(totalTokens, chunkSize, overlap);
        return new EmbeddingPlan(
            plan.Chunks,
            plan.TotalTokensWithOverlap,
            plan.OverheadTokens,
            model,
            plan.Chunks,
            plan.TotalTokensWithOverlap / 1e6 * model.InputPerM);
    }
}

// ---------- showcase examples (the canonical suite lives in src/lib) ----------
internal static class Showcase
{
    private static void Main()
    {
        // 1,000 tokens: ceil((1000-64)/(512-64)) = 3 chunks, 2 seams x 64.
        var a = EmbeddingChunkPlanner.PlanChunks(1000);
        Assert(a == new ChunkPlan(3, 1128, 128));

        // A document that fits one chunk has no seam overhead.
        Assert(EmbeddingChunkPlanner.PlanChunks(512) == new ChunkPlan(1, 512, 0));

        // Zero/negative input or non-positive chunk size -> the zero plan.
        Assert(EmbeddingChunkPlanner.PlanChunks(0) == ChunkPlan.Zero);
        Assert(EmbeddingChunkPlanner.PlanChunks(-100) == ChunkPlan.Zero);
        Assert(EmbeddingChunkPlanner.PlanChunks(1000, chunkSize: 0) == ChunkPlan.Zero);
        Assert(EmbeddingChunkPlanner.PlanChunks(1000, chunkSize: -8) == ChunkPlan.Zero);

        // overlap 600 > floor(512/2) = 256 -> clamped to 256.
        Assert(EmbeddingChunkPlanner.PlanChunks(1000, overlap: 600) == new ChunkPlan(3, 1512, 512));

        // Negative overlap clamps to 0: 1000 tokens -> ceil(1000/512) = 2 chunks.
        Assert(EmbeddingChunkPlanner.PlanChunks(1000, overlap: -5) == new ChunkPlan(2, 1000, 0));

        // chunkSize without overlap: ceil((1000-64)/192) = 5 chunks.
        Assert(EmbeddingChunkPlanner.PlanChunks(1000, chunkSize: 256) == new ChunkPlan(5, 1256, 256));

        // Shorter than the overlap still yields one chunk.
        Assert(EmbeddingChunkPlanner.PlanChunks(50, overlap: 64) == new ChunkPlan(1, 50, 0));

        // chunkSize of 1 clamps overlap to 0: ceil(3/1) = 3 chunks.
        Assert(EmbeddingChunkPlanner.PlanChunks(3, chunkSize: 1) == new ChunkPlan(3, 3, 0));

        // Pricing: 1,000 tokens on text-embedding-3-small @ 1536 dims.
        var priced = EmbeddingChunkPlanner.PlanEmbedding(1000, "text-embedding-3-small", 1536);
        Assert(priced is { Chunks: 3, TotalTokensWithOverlap: 1128, Vectors: 3 });
        Assert(Math.Abs(priced.Cost - 0.00002256) < 1e-12); // 1128 / 1e6 * $0.02

        // A single-chunk document on voyage-3-lite @ 512 dims.
        var single = EmbeddingChunkPlanner.PlanEmbedding(512, "voyage-3-lite", 512);
        Assert(single is { Vectors: 1 });
        Assert(Math.Abs(single.Cost - 0.00001024) < 1e-12); // 512 / 1e6 * $0.02

        // Unknown model or unoffered dims -> null.
        Assert(EmbeddingChunkPlanner.PlanEmbedding(1000, "text-embedding-3-small", 999) is null);
        Assert(EmbeddingChunkPlanner.PlanEmbedding(1000, "ghost", 1536) is null);

        // Zero tokens price out to a zero-cost plan.
        var zero = EmbeddingChunkPlanner.PlanEmbedding(0, "text-embedding-3-small", 1536);
        Assert(zero is { Chunks: 0, Vectors: 0 } && zero.Cost == 0.0);
    }

    private static void Assert(bool condition) => System.Diagnostics.Debug.Assert(condition);
}