Embedding Chunk Planner — Python source

"""Embedding Chunk Planner — pure chunking math for RAG pipelines.

Language: Python (3.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 raises.
  - Functionally equivalent to the TS reference: same inputs -> same outputs.
  - Self-contained: stdlib only (no pip 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):
  - chunk_size <= 0 or total_tokens <= 0 -> ChunkPlan(0, 0, 0) (nothing to
    embed).
  - Negative overlap is treated as 0; overlap then clamps to at most
    chunk_size // 2 so consecutive chunks always advance.
  - chunks = max(1, ceil((total_tokens - overlap) / (chunk_size - overlap)))
    — a tiny document still yields one chunk.
"""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Optional, Tuple

__all__ = [
    "ChunkOptions",
    "ChunkPlan",
    "EmbeddingModel",
    "EmbeddingPlan",
    "EMBEDDING_MODELS",
    "DEFAULT_CHUNK_OPTIONS",
    "get_embedding_model",
    "plan_chunks",
    "plan_embedding",
]


@dataclass(frozen=True)
class ChunkOptions:
    """Chunking knobs, in tokens. Defaults: 512-token chunks, 64-token overlap.

    Mirrors the TS ``ChunkOptions``; ``DEFAULT_CHUNK_OPTIONS`` mirrors the TS
    default constant.
    """

    chunk_size: int = 512
    overlap: int = 64


@dataclass(frozen=True)
class ChunkPlan:
    """How a document splits into overlapping chunks."""

    chunks: int
    total_tokens_with_overlap: int
    overhead_tokens: int


@dataclass(frozen=True)
class EmbeddingModel:
    """One embedding model's offered dimensions and pricing.

    ``input_per_m`` is USD per 1M input tokens.
    """

    id: str
    vendor: str
    dims: Tuple[int, ...]
    input_per_m: float


#: Embedding model price table — the SSOT for pricing, mirrored from
#: src/lib/ai/embeddings.ts. Refresh both files together.
EMBEDDING_MODELS: Tuple[EmbeddingModel, ...] = (
    EmbeddingModel("text-embedding-3-small", "OpenAI", (512, 1536), 0.02),
    EmbeddingModel("text-embedding-3-large", "OpenAI", (256, 1024, 3072), 0.13),
    EmbeddingModel("embed-english-v3.0", "Cohere", (512, 1024, 1536), 0.1),
    EmbeddingModel("voyage-3-lite", "Voyage AI", (512, 1024), 0.02),
)

DEFAULT_CHUNK_OPTIONS = ChunkOptions()


def get_embedding_model(model_id: str) -> Optional[EmbeddingModel]:
    """Look up an embedding model by id. Returns ``None`` for unknown ids."""
    return next((m for m in EMBEDDING_MODELS if m.id == model_id), None)


def plan_chunks(
    total_tokens: int,
    chunk_size: Optional[int] = None,
    overlap: Optional[int] = None,
) -> ChunkPlan:
    """Plan how ``total_tokens`` split into overlapping chunks.

    ``chunk_size`` / ``overlap`` are each optional (the TS ``Partial<ChunkOptions>``);
    a ``None`` value falls back to the 512 / 64 default, and passing only one
    leaves the other at its default.
    """
    cs = DEFAULT_CHUNK_OPTIONS.chunk_size if chunk_size is None else chunk_size
    ov_raw = DEFAULT_CHUNK_OPTIONS.overlap if overlap is None else overlap
    if cs <= 0 or total_tokens <= 0:
        return ChunkPlan(0, 0, 0)
    # min(max(overlap, 0), chunk_size // 2) — the TS clamp.
    ov = min(max(ov_raw, 0), cs // 2)
    # Float division + ceil mirrors TS's Math.ceil exactly (a negative
    # numerator below the max(1, ...) floor is handled next).
    chunks = max(1, math.ceil((total_tokens - ov) / (cs - ov)))
    total_with_overlap = total_tokens + (chunks - 1) * ov
    return ChunkPlan(chunks, total_with_overlap, total_with_overlap - total_tokens)


@dataclass(frozen=True)
class EmbeddingPlan(ChunkPlan):
    """Chunk plan plus pricing for one embedding call.

    ``vectors`` is one per chunk; ``cost`` is USD:
    ``total_tokens_with_overlap / 1e6 * model.input_per_m``.
    """

    model: EmbeddingModel
    vectors: int
    cost: float


def plan_embedding(
    total_tokens: int,
    model_id: str,
    dims: int,
    chunk_size: Optional[int] = None,
    overlap: Optional[int] = None,
) -> Optional[EmbeddingPlan]:
    """Chunk a document AND price its embedding for ``model_id`` at ``dims``
    dimensions. Unknown model, or dims the model does not offer -> ``None``."""
    model = get_embedding_model(model_id)
    if model is None or dims not in model.dims:
        return None
    plan = plan_chunks(total_tokens, chunk_size, overlap)
    return EmbeddingPlan(
        chunks=plan.chunks,
        total_tokens_with_overlap=plan.total_tokens_with_overlap,
        overhead_tokens=plan.overhead_tokens,
        model=model,
        vectors=plan.chunks,
        cost=(plan.total_tokens_with_overlap / 1e6) * model.input_per_m,
    )