Embedding Chunk Planner — Swift source

// Embedding Chunk Planner — pure chunking math for RAG pipelines.
//
// Language: Swift (Swift 5.9, 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 traps (nil instead of an error).
//   - Functionally equivalent to the TS reference: same inputs -> same outputs.
//   - Self-contained: the Swift standard library only (no SwiftPM packages;
//     no need for Foundation). 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.

/// One embedding model's offered dimensions (ascending, Matryoshka shortening
/// included) and pricing: USD per 1M input tokens.
public struct EmbeddingModel: Equatable, Sendable {
    public let id: String
    public let vendor: String
    public let dims: [Int]
    public let inputPerM: Double

    public init(id: String, vendor: String, dims: [Int], inputPerM: Double) {
        self.id = id
        self.vendor = vendor
        self.dims = dims
        self.inputPerM = inputPerM
    }
}

/// How a document splits into overlapping chunks.
public struct ChunkPlan: Equatable, Sendable {
    public let chunks: Int
    public let totalTokensWithOverlap: Int
    public let overheadTokens: Int

    public init(chunks: Int, totalTokensWithOverlap: Int, overheadTokens: Int) {
        self.chunks = chunks
        self.totalTokensWithOverlap = totalTokensWithOverlap
        self.overheadTokens = overheadTokens
    }

    /// The "nothing to embed" plan the TS source returns for zero/negative
    /// input or a non-positive chunk size.
    public static let zero = ChunkPlan(chunks: 0, totalTokensWithOverlap: 0, overheadTokens: 0)
}

/// Chunking knobs, in tokens. Mirrors the TS `Partial<ChunkOptions>`: each
/// field is independently optional — nil falls back to the 512 / 64 default,
/// and setting one leaves the other at its default.
public struct ChunkOptions: Sendable {
    public var chunkSize: Int?
    public var overlap: Int?

    public init(chunkSize: Int? = nil, overlap: Int? = nil) {
        self.chunkSize = chunkSize
        self.overlap = overlap
    }
}

/// Chunk plan plus pricing for one embedding call. The three chunk fields are
/// flattened in (the TS `...plan` spread) so the struct reads like the TS
/// `EmbeddingPlan extends ChunkPlan`.
public struct EmbeddingPlan: Equatable, Sendable {
    public let chunks: Int
    public let totalTokensWithOverlap: Int
    public let overheadTokens: Int
    public let model: EmbeddingModel
    /// One vector per chunk.
    public let vectors: Int
    /// USD: totalTokensWithOverlap / 1e6 * model.inputPerM.
    public let cost: Double

    public init(chunks: Int, totalTokensWithOverlap: Int, overheadTokens: Int,
                model: EmbeddingModel, vectors: Int, cost: Double) {
        self.chunks = chunks
        self.totalTokensWithOverlap = totalTokensWithOverlap
        self.overheadTokens = overheadTokens
        self.model = model
        self.vectors = vectors
        self.cost = cost
    }
}

public enum EmbeddingChunkPlanner {
    /// Embedding model price table — the SSOT for pricing, mirrored from
    /// src/lib/ai/embeddings.ts. Refresh both files together.
    public static let embeddingModels: [EmbeddingModel] = [
        EmbeddingModel(id: "text-embedding-3-small", vendor: "OpenAI", dims: [512, 1536], inputPerM: 0.02),
        EmbeddingModel(id: "text-embedding-3-large", vendor: "OpenAI", dims: [256, 1024, 3072], inputPerM: 0.13),
        EmbeddingModel(id: "embed-english-v3.0", vendor: "Cohere", dims: [512, 1024, 1536], inputPerM: 0.1),
        EmbeddingModel(id: "voyage-3-lite", vendor: "Voyage AI", dims: [512, 1024], inputPerM: 0.02),
    ]

    /// Default knobs: 512-token chunks, 64-token overlap (TS DEFAULT_CHUNK_OPTIONS).
    public static let defaultChunkSize = 512
    public static let defaultOverlap = 64

    /// Look up an embedding model by id. Returns nil for unknown ids.
    public static func getEmbeddingModel(_ id: String) -> EmbeddingModel? {
        embeddingModels.first { $0.id == id }
    }

    /// Plan how `totalTokens` split into overlapping chunks. `opts` may be nil
    /// (both defaults), mirroring the TS optional parameter; each nil field
    /// falls back to the 512 / 64 default independently.
    public static func planChunks(_ totalTokens: Int, opts: ChunkOptions? = nil) -> ChunkPlan {
        let chunkSize = opts?.chunkSize ?? defaultChunkSize
        let overlapRaw = opts?.overlap ?? defaultOverlap

        if chunkSize <= 0 || totalTokens <= 0 {
            return .zero
        }

        // min(max(overlapRaw, 0), chunkSize / 2) — the TS clamp. Overlap that
        // large would never advance, so consecutive chunks always gain at least
        // half a chunk. (chunkSize >= 1 here, so chunkSize - overlap is never zero.)
        let overlap = min(max(overlapRaw, 0), chunkSize / 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).
        let quotient = Double(totalTokens - overlap) / Double(chunkSize - overlap)
        let chunks = max(1, Int(quotient.rounded(.up)))

        let totalTokensWithOverlap = totalTokens + (chunks - 1) * overlap
        return ChunkPlan(chunks: chunks,
                         totalTokensWithOverlap: totalTokensWithOverlap,
                         overheadTokens: totalTokensWithOverlap - totalTokens)
    }

    /// Chunk a document AND price its embedding for `modelId` at `dims`
    /// dimensions. Unknown model, or dims the model does not offer -> nil.
    public static func planEmbedding(_ totalTokens: Int, _ modelId: String, dims: Int,
                                     opts: ChunkOptions? = nil) -> EmbeddingPlan? {
        guard let model = getEmbeddingModel(modelId), model.dims.contains(dims) else {
            return nil
        }
        let plan = planChunks(totalTokens, opts: opts)
        return EmbeddingPlan(chunks: plan.chunks,
                             totalTokensWithOverlap: plan.totalTokensWithOverlap,
                             overheadTokens: plan.overheadTokens,
                             model: model,
                             vectors: plan.chunks,
                             cost: Double(plan.totalTokensWithOverlap) / 1e6 * model.inputPerM)
    }
}

// ---------- showcase examples (the canonical suite lives in src/lib) ----------
func runShowcaseExamples() {
    // 1,000 tokens: ceil((1000-64)/(512-64)) = 3 chunks, 2 seams x 64.
    let a = EmbeddingChunkPlanner.planChunks(1000)
    precondition(a == ChunkPlan(chunks: 3, totalTokensWithOverlap: 1128, overheadTokens: 128))

    // A document that fits one chunk has no seam overhead.
    precondition(EmbeddingChunkPlanner.planChunks(512)
        == ChunkPlan(chunks: 1, totalTokensWithOverlap: 512, overheadTokens: 0))

    // Zero/negative input or non-positive chunk size -> the zero plan.
    precondition(EmbeddingChunkPlanner.planChunks(0) == .zero)
    precondition(EmbeddingChunkPlanner.planChunks(-100) == .zero)
    precondition(EmbeddingChunkPlanner.planChunks(1000, opts: ChunkOptions(chunkSize: 0)) == .zero)
    precondition(EmbeddingChunkPlanner.planChunks(1000, opts: ChunkOptions(chunkSize: -8)) == .zero)

    // overlap 600 > floor(512/2) = 256 -> clamped to 256.
    precondition(EmbeddingChunkPlanner.planChunks(1000, opts: ChunkOptions(overlap: 600))
        == ChunkPlan(chunks: 3, totalTokensWithOverlap: 1512, overheadTokens: 512))

    // Negative overlap clamps to 0: 1000 tokens -> ceil(1000/512) = 2 chunks.
    precondition(EmbeddingChunkPlanner.planChunks(1000, opts: ChunkOptions(overlap: -5))
        == ChunkPlan(chunks: 2, totalTokensWithOverlap: 1000, overheadTokens: 0))

    // chunkSize without overlap: ceil((1000-64)/192) = 5 chunks.
    precondition(EmbeddingChunkPlanner.planChunks(1000, opts: ChunkOptions(chunkSize: 256))
        == ChunkPlan(chunks: 5, totalTokensWithOverlap: 1256, overheadTokens: 256))

    // Shorter than the overlap still yields one chunk.
    precondition(EmbeddingChunkPlanner.planChunks(50, opts: ChunkOptions(overlap: 64))
        == ChunkPlan(chunks: 1, totalTokensWithOverlap: 50, overheadTokens: 0))

    // chunkSize of 1 clamps overlap to 0: ceil(3/1) = 3 chunks.
    precondition(EmbeddingChunkPlanner.planChunks(3, opts: ChunkOptions(chunkSize: 1))
        == ChunkPlan(chunks: 3, totalTokensWithOverlap: 3, overheadTokens: 0))

    // Pricing: 1,000 tokens on text-embedding-3-small @ 1536 dims.
    let priced = EmbeddingChunkPlanner.planEmbedding(1000, "text-embedding-3-small", dims: 1536)
    precondition(priced?.chunks == 3 && priced?.totalTokensWithOverlap == 1128 && priced?.vectors == 3)
    precondition(abs((priced?.cost ?? -1) - 0.00002256) < 1e-12) // 1128 / 1e6 * $0.02

    // A single-chunk document on voyage-3-lite @ 512 dims.
    let single = EmbeddingChunkPlanner.planEmbedding(512, "voyage-3-lite", dims: 512)
    precondition(single?.vectors == 1)
    precondition(abs((single?.cost ?? -1) - 0.00001024) < 1e-12) // 512 / 1e6 * $0.02

    // Unknown model or unoffered dims -> nil.
    precondition(EmbeddingChunkPlanner.planEmbedding(1000, "text-embedding-3-small", dims: 999) == nil)
    precondition(EmbeddingChunkPlanner.planEmbedding(1000, "ghost", dims: 1536) == nil)

    // Zero tokens price out to a zero-cost plan.
    let zero = EmbeddingChunkPlanner.planEmbedding(0, "text-embedding-3-small", dims: 1536)
    precondition(zero?.chunks == 0 && zero?.vectors == 0 && zero?.cost == 0.0)
}