diff --git a/benchmarks/README.md b/benchmarks/README.md index 6ec2a42e..db93a515 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -15,6 +15,7 @@ Provenance is tagged inline in [`../docs/benchmarks.md`](../docs/benchmarks.md): | [`batched_retrieval/`](batched_retrieval) | What queries/sec does the public `search_many` SDK path serve as the batch grows — on the CPU kernel, and (on CUDA) the GPU-resident scan? | Any host (local) / CUDA | `uv run python benchmarks/batched_retrieval/run.py` | | [`mps_vs_neon/`](mps_vs_neon) | Does the Apple-GPU (MPS) exact scan beat the default CPU NEON scan on a Mac? | Apple Silicon (local) | see folder README | | [`memory_integrations/`](memory_integrations) | As the memory backend behind LangChain / LlamaIndex / mem0, how does LodeDB compare to each framework's default and common stores (in-memory, FAISS, Chroma, Qdrant) on ingest, recall, footprint, and durable single-add? | Modal A10 / L40S (CUDA) | `modal run benchmarks/memory_integrations/modal_bench.py::main_a10` | +| [`openknowledge_embeddings/`](openknowledge_embeddings) | What does an OpenKnowledge user trade by pointing `baseUrl` at LodeDB instead of OpenAI on retrieval quality, indexing throughput, query latency, cost, and content egress? | Modal CPU (8.0 quota, 16 GB, Ubuntu 22.04) | `modal run benchmarks/openknowledge_embeddings/modal_bench.py::bench --docs 0 --queries 100` | | [`multimodal_image/`](multimodal_image) | Holding CLIP-dimension vectors fixed, how does LodeDB compare to Chroma and Qdrant on image-vector footprint, query latency, recall, and ingest? | Any host (local) | `uv run python benchmarks/multimodal_image/run.py` | | [`embedding_megakernel_spike/`](embedding_megakernel_spike) | Where does single-query MiniLM embedding latency go, and does a fused batch-1 megakernel beat the best low-maintenance compiler path (issue #67 gate)? | Apple M1 (local) / Modal A10 / L40S (CUDA) | `python benchmarks/embedding_megakernel_spike/run.py` | diff --git a/benchmarks/openknowledge_embeddings/README.md b/benchmarks/openknowledge_embeddings/README.md new file mode 100644 index 00000000..64e4b2ba --- /dev/null +++ b/benchmarks/openknowledge_embeddings/README.md @@ -0,0 +1,211 @@ +# openknowledge_embeddings benchmark + +Compares OpenKnowledge's default OpenAI embedding provider with LodeDB's local +OpenAI-compatible `/v1/embeddings` endpoint. OpenKnowledge (Inkeep's local-first +Markdown wiki) accepts any OpenAI-compatible embeddings endpoint for semantic search; +its default provider is OpenAI `text-embedding-3-small`. `lodedb serve` exposes the +same wire contract, so an OpenKnowledge user can point `baseUrl` at a local LodeDB server. + +The question: *what does an OpenKnowledge user trade by pointing `baseUrl` at LodeDB +instead of OpenAI?* + +Metrics only (counts, bytes, latency, provider labels, and scores). No raw documents, +queries, embeddings, or credentials are written. + +**Overall:** LodeDB passes OpenKnowledge's pre-registered quality gate unmodified. On +the full wiki, OpenAI indexes 1.4-2.9x faster for bulk indexing on an 8-vCPU quota +(median 1.8x across four pinned LodeDB runs), while LodeDB has roughly 38x lower p50 +query latency, zero workload cost, and zero external content egress. + +## Axis 1: retrieval quality + +This uses OpenKnowledge's own `packages/server/src/embeddings/eval/semantic-eval.ts` +unmodified, with its committed `eval-set.json`: 22 documents, 28 query pairs, 15 tune +pairs, and 13 held-out pairs. The OpenAI and minilm providers were repeated twice, with +identical numbers across repeats. + +### Reproduce + +From an `open-knowledge` checkout with Bun installed, start LodeDB in another shell: + +```bash +lodedb serve --model minilm --port 8093 +``` + +Run the OpenAI provider: + +```bash +OK_EMBED_SMOKE=1 OK_EMBEDDINGS_API_KEY=sk-... bun run --conditions=development packages/server/src/embeddings/eval/semantic-eval.ts +``` + +Run LodeDB through the same harness: + +```bash +OK_EMBED_SMOKE=1 OK_EMBEDDINGS_API_KEY=placeholder OK_EMBEDDINGS_BASE_URL=http://127.0.0.1:8093/v1 OK_EMBEDDINGS_MODEL=minilm OK_EMBEDDINGS_DIMENSIONS=384 bun run --conditions=development packages/server/src/embeddings/eval/semantic-eval.ts +``` + +The BGE row uses `OK_EMBEDDINGS_MODEL=bge OK_EMBEDDINGS_DIMENSIONS=768` against +`lodedb serve --model bge`. + +### Results + +Held-out split, `n=13`. The shared lexical baseline is MRR `0.4997`, recall@1 +`0.3077`, and recall@5 `0.7692`. + +| provider | MRR@10 | recall@1 | recall@5 | MRR gain vs lexical (min 0.05) | recall@5 gain (min 0.08) | lexical-strong regression (max 0.03) | their FR2 gate | +|---|---:|---:|---:|---:|---:|---:|---| +| openai text-embedding-3-small (1536d) | 0.7974 | 0.6923 | 0.9231 | +0.2978 | +0.1538 | -0.2500 | PASS | +| lodedb minilm (384d, local) | 0.7526 | 0.6154 | 1.0000 | +0.2529 | +0.2308 | -0.1500 | PASS | +| lodedb bge (768d, local) | 0.7333 | 0.5385 | 1.0000 | +0.2337 | +0.2308 | -0.1167 | PASS | + +LodeDB's bge preset (768d) also PASSES the FR2 gate, but scores below minilm through +this API because BGE is an asymmetric model whose query prefix cannot be expressed over +the OpenAI wire. This is why minilm is the documented recommendation. + +LodeDB passes their pre-registered quality gate unmodified. The held-out sample has +only 13 pairs, so the per-metric differences are directional, not significant. In +particular, MRR is `0.7526` vs `0.7974`, while recall@5 is `1.0000` vs `0.9231`. +The tune split selected RRF `k=10` for both providers. + +## Axis 2: throughput, latency, cost, and egress + +### Methodology + +`modal_bench.py` clones `kubernetes/website` at corpus commit +`71d23f81e3479361befc94564e2b955860c03164` and runs against Markdown under `content/en`, +sorted by path. It uses faithful ports of OpenKnowledge's chunker and embedder behavior: + +- Chunking targets 8,000 characters, overlaps 400 characters, and caps a document at + 80 chunks. The port is parity-tested against the OpenKnowledge `chunking.test.ts` + cases in `test_okchunk.py`. +- Documents use greedy sequential batches capped at 96 inputs and 96,000 characters. + Batches are sent sequentially. Queries are sent as one sequential request each. +- Both providers use the same Python HTTP client, request shape, response checks, + Float32 conversion, and normalization. Every request sends a Bearer header. +- The client uses persistent pooled HTTP connections. The `reconnects` counter is 0 + everywhere in the five committed JSONs. An unpooled pilot measured remote query p50 + at 171.5 ms versus 173.3 ms with pooling, so pooling did not materially change remote + latency and the comparison is robust to HTTP client behavior. +- The chunker preserves OpenKnowledge's UTF-16 behavior. Lone surrogates are sanitized + only at the HTTP boundary, immediately before JSON encoding, so UTF-16 parity is + preserved. +- Document requests have 30 second timeouts and query requests have 8 second + timeouts. Retryable statuses are 408, 409, 429, 500, 502, 503, and 504, with up to + 4 retries and a 500 ms exponential backoff. +- The first document batch is a warmup and is excluded from workload timing for both + providers. The LodeDB server start and first batch are reported separately. + +Provider-native dimensions are retained: LodeDB uses 384 dimensions and OpenAI uses +1,536 dimensions. No dimension matching is added for the benchmark. + +### Hardware and setup + +| item | value | +|---|---| +| container | Modal, Ubuntu 22.04, Python 3.11 | +| CPU and memory | CPU=8.0 quota, 16 GB; `os.cpu_count()` reports 24 host cores, but the quota is what matters | +| LodeDB | `lodedb serve --model minilm`, CPU ONNX, loopback, same container | +| OpenAI | measured from a Modal datacenter, a favorable network for the remote provider | + +### Both sides measured identically + +- The same sorted corpus, chunking, batches, timeouts, retry policy, HTTP client, and + sequential query loop are used for both providers. +- Workload timing starts after each provider's warmup batch. Retries and request and + response body bytes are counted for the measured workload. +- Document throughput is measured across all chunks. Query latency is measured across + 100 sequential queries in the full-corpus run and 50 sequential queries in the + 300-document run. + +### Reproduce + +From the repository root: + +```bash +modal run benchmarks/openknowledge_embeddings/modal_bench.py::bench --docs 300 --queries 50 +modal run benchmarks/openknowledge_embeddings/modal_bench.py::bench --docs 0 --queries 100 +``` + +The run needs a Modal secret named `openai-embeddings-bench` containing +`OPENAI_API_KEY`. The metrics-only outputs are [`modal-300docs.json`](results/modal-300docs.json), +[`modal-full-corpus.json`](results/modal-full-corpus.json), and the three pinned LodeDB +repeats: [`modal-full-corpus-lodedb-repeat-1.json`](results/modal-full-corpus-lodedb-repeat-1.json), +[`modal-full-corpus-lodedb-repeat-2.json`](results/modal-full-corpus-lodedb-repeat-2.json), +and [`modal-full-corpus-lodedb-repeat-3.json`](results/modal-full-corpus-lodedb-repeat-3.json). + +### Full corpus results + +The full run contains 2,441 documents, 4,485 chunks, and 22,746,024 characters. +Chunk lengths are p50 6,791, p95 7,999, and max 8,000 characters. 836 documents +produced more than one chunk. Each provider used 257 document batches, with 0 retries +and 0 reconnects on both sides. + +| metric | lodedb (minilm, 384d, loopback CPU) | openai (text-embedding-3-small, 1536d) | +|---|---:|---:| +| document embed wall time | 229.4 s median (main 275.2; range 165.8-355.7) | 121.2 s | +| chunks per second | 20.4 median (main 16.3; range 12.6-27.0) | 37.0 | +| batch latency p50 / p95 / max (ms) | 1,047.5 / 1,205.3 / 1,384.0 (main run) | 391.6 / 741.7 / 2,693.7 | +| query latency p50 / p95 / max (ms), 100 sequential queries | 4.5 / 6.5 / 7.7 (main run; p50 range 4.0-4.8 across four runs) | 173.3 / 346.1 / 2,804.8 | +| provider-reported workload tokens | 5,890,209 (local server estimate) | 5,805,063 (provider-metered) | +| cost per full index | $0 | $0.1161 (at $0.020 per 1M tokens) | +| external egress (request + response body bytes) | 0 | 159,504,761 | +| one-time setup | 3.8 s server start + 3.7 s first batch (ONNX session init + model fetch) | none | + +### Run-to-run variance + +Document embedding throughput on Modal CPU varied about 2x across identical full-corpus +LodeDB runs: 16.3, 27.0, 12.6, and 24.4 chunks/s (median 20.4, range 12.6-27.0). +Wall times were 275.2, 165.8, 355.7, and 183.5 s (median 229.4, range 165.8-355.7). +Query latency did not vary materially: p50 was 4.0-4.8 ms across the four runs. Raw +JSONs for every run are committed. + +### 300-document results + +This run contains 300 documents, 490 chunks, and 30 document batches. + +| metric | lodedb (minilm, 384d, loopback CPU) | openai (text-embedding-3-small, 1536d) | +|---|---:|---:| +| document embed wall time | 27.0 s | 11.1 s | +| chunks per second | 18.1 | 44.1 | +| query latency p50 / p95 / max (ms), 50 sequential queries | 5.6 / 8.7 / 10.3 | 138.9 / 259.9 / 399.1 | +| cost | $0 | $0.0127 | +| external egress (request + response body bytes) | 0 | 18,736,627 | + +### Reading the results + +- **Bulk indexing.** OpenAI is 1.4-2.9x faster on an 8-vCPU quota than the local CPU + container (median 1.8x from 37.0 vs 20.4 chunks per second). LodeDB's median full + corpus wall time was 229.4 s (range 165.8-355.7) versus 121.2 s for OpenAI. Both + index a 2,441-document wiki in minutes. Indexing is a one-time background operation + that scales with local cores. +- **Interactive search latency.** For the per-keystroke user-facing operation, LodeDB + p50 is roughly 38x lower in the main run: 4.5 ms vs 173.3 ms. Its max was 7.7 ms + versus 2,804.8 ms for OpenAI, including one multi-second outlier. Across the four + LodeDB runs, p50 remained 4.0-4.8 ms. +- **Cost and privacy.** LodeDB costs $0 and has 0 external request and response body + bytes. The metered remote API receives the full wiki text and every search query; + total external request and response body traffic is 159,504,761 bytes, or 159.5 MB, + for this corpus. +- **Quality.** LodeDB passes OpenKnowledge's pre-registered quality gate. The held-out + quality sample is small, so the MRR and recall deltas should be treated as + directional rather than significant. + +### Cost and token accounting + +The price constant is pinned to [OpenAI API pricing](https://openai.com/api/pricing), +accessed 2026-07-09: `$0.020` per 1M tokens for `text-embedding-3-small`. The cost +calculation uses only the provider-metered OpenAI workload token count. + +LodeDB's 5,890,209 workload tokens are a local server estimate; +OpenAI's 5,805,063 are provider-metered. These token columns are not directly +comparable. Warmup usage is excluded from both workload totals and costs. + +## Caveats + +- The quality result is a pass against the committed OpenKnowledge harness and its + pre-registered thresholds, not a large evaluation. The held-out split has 13 query + pairs. +- The OpenAI run was measured from a Modal datacenter, which is favorable to the remote + provider's network latency. LodeDB's endpoint is loopback in the same container. +- The full-corpus document timing excludes the one-time LodeDB server start, ONNX + session initialization, and model fetch; those are shown separately in the table. diff --git a/benchmarks/openknowledge_embeddings/bench_core.py b/benchmarks/openknowledge_embeddings/bench_core.py new file mode 100644 index 00000000..d59c8727 --- /dev/null +++ b/benchmarks/openknowledge_embeddings/bench_core.py @@ -0,0 +1,893 @@ +"""Shared OpenKnowledge-compatible corpus, batching, and HTTP benchmark logic. + +This module deliberately has no Modal import so its local smoke path can validate +the same batcher and OpenAI-compatible client against a local LodeDB server. +""" + +from __future__ import annotations + +import argparse +import http.client +import json +import math +import os +import random +import re +import subprocess +import tempfile +import time +import urllib.error +import urllib.parse +import urllib.request +from array import array +from collections.abc import Iterator, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +try: + from benchmarks.openknowledge_embeddings.okchunk import chunk_document, javascript_length +except ModuleNotFoundError: # Supports `python benchmarks/.../bench_core.py --local-smoke`. + from okchunk import chunk_document, javascript_length + +DEFAULT_MAX_BATCH_SIZE = 96 +DEFAULT_MAX_BATCH_CHARS = 96_000 +DOCUMENT_TIMEOUT_SECONDS = 30.0 +QUERY_TIMEOUT_SECONDS = 8.0 +MAX_RETRIES = 4 +BACKOFF_BASE_SECONDS = 0.5 +RETRYABLE_STATUSES = frozenset({408, 409, 429, 500, 502, 503, 504}) +LODEDB_DIMENSIONS = 384 +OPENAI_DIMENSIONS = 1_536 +# Source: OpenAI API pricing, text-embedding-3-small, accessed 2026-07-09. +OPENAI_PRICE_USD_PER_MILLION_TOKENS = 0.020 + +EmbeddingRole = Literal["document", "query"] +ProviderName = Literal["lodedb", "openai"] + + +@dataclass(frozen=True) +class ProviderConfig: + """The request and response settings OpenKnowledge applies to one provider.""" + + name: ProviderName + base_url: str + model: str + requested_dimensions: int | None + expected_dimensions: int + api_key: str + token_usage_label: str + + +@dataclass(frozen=True) +class TrafficSnapshot: + request_body_bytes: int + response_body_bytes: int + + @property + def request_response_body_bytes(self) -> int: + return self.request_body_bytes + self.response_body_bytes + + +@dataclass +class HttpBodyTraffic: + """Body traffic from every HTTP attempt, including retries.""" + + request_body_bytes: int = 0 + response_body_bytes: int = 0 + + def snapshot(self) -> TrafficSnapshot: + return TrafficSnapshot(self.request_body_bytes, self.response_body_bytes) + + +@dataclass(frozen=True) +class BatchResult: + duration_seconds: float + retry_count: int + reconnects: int + reported_total_tokens: int | float + + +@dataclass(frozen=True) +class Corpus: + """The selected wiki documents reduced to the data the workload needs.""" + + chunks: tuple[str, ...] + query_candidates: tuple[str, ...] + stats: dict[str, Any] + + +@dataclass(frozen=True) +class LocalLodeDBServer: + base_url: str + cold_start_seconds: float + + +class EmbeddingRequestError(RuntimeError): + """A request failed after its OpenKnowledge-compatible retry policy.""" + + +class EmbeddingResponseError(RuntimeError): + """A successful HTTP response did not match OpenKnowledge's expectations.""" + + +class _RetryableRequestError(RuntimeError): + """Internal classification for an HTTP status eligible for retry.""" + + +def batch_inputs( + texts: Sequence[str], + *, + max_batch_size: int = DEFAULT_MAX_BATCH_SIZE, + max_batch_chars: int = DEFAULT_MAX_BATCH_CHARS, +) -> list[list[str]]: + """Ports OpenKnowledge's greedy sequential count and character batcher.""" + + batches: list[list[str]] = [] + current: list[str] = [] + chars = 0 + for text in texts: + if current and ( + len(current) >= max_batch_size or chars + javascript_length(text) > max_batch_chars + ): + batches.append(current) + current = [] + chars = 0 + current.append(text) + chars += javascript_length(text) + if current: + batches.append(current) + return batches + + +def _round(value: float | None, digits: int = 6) -> float | None: + return None if value is None else round(value, digits) + + +def _percentile(values: Sequence[float], percentile: float) -> float | None: + if not values: + return None + ordered = sorted(values) + index = (len(ordered) - 1) * percentile + lower = math.floor(index) + upper = math.ceil(index) + if lower == upper: + return ordered[lower] + return ordered[lower] + (ordered[upper] - ordered[lower]) * (index - lower) + + +def _latency_percentiles_ms(values_seconds: Sequence[float]) -> dict[str, float | None]: + values_ms = [value * 1_000.0 for value in values_seconds] + return { + "p50": _round(_percentile(values_ms, 0.50)), + "p95": _round(_percentile(values_ms, 0.95)), + "max": _round(max(values_ms) if values_ms else None), + } + + +def _chunk_char_percentiles(chunks: Sequence[str]) -> dict[str, float | int | None]: + lengths = [javascript_length(chunk) for chunk in chunks] + maximum = max(lengths) if lengths else None + return { + "p50": _round(_percentile(lengths, 0.50)), + "p95": _round(_percentile(lengths, 0.95)), + "max": maximum, + } + + +def _clean_query_candidate(value: str) -> str: + normalized = " ".join(value.replace("`", "").split()) + return normalized[:160].strip() + + +def _document_query_candidates(path: Path, text: str) -> list[str]: + """Produces short deterministic queries from a document's title and headings.""" + + candidates: list[str] = [] + frontmatter_title = re.search(r"(?m)^title:\s*[\"']?(.+?)[\"']?\s*$", text) + if frontmatter_title: + candidates.append(_clean_query_candidate(frontmatter_title.group(1))) + for heading in re.findall(r"(?m)^\s{0,3}#{1,6}\s+(.+?)(?:\s+#+)?\s*$", text): + candidates.append(_clean_query_candidate(heading)) + if len(candidates) >= 3: + break + if not candidates: + candidates.append(_clean_query_candidate(path.stem.replace("-", " "))) + return [candidate for candidate in candidates if candidate] + + +def load_kubernetes_corpus(corpus_root: Path, docs: int = 0) -> Corpus: + """Loads sorted ``content/en/**/*.md`` files and applies OpenKnowledge chunking.""" + + if docs < 0: + raise ValueError("--docs must be zero (all documents) or a positive integer") + markdown_root = corpus_root / "content" / "en" + paths = sorted(markdown_root.rglob("*.md"), key=lambda path: path.as_posix()) + if docs: + paths = paths[:docs] + + chunks: list[str] = [] + query_candidates: list[str] = [] + total_chars = 0 + docs_with_multiple_chunks = 0 + for path in paths: + text = path.read_text(encoding="utf-8") + document_chunks = chunk_document(text) + chunks.extend(document_chunks) + query_candidates.extend(_document_query_candidates(path, text)) + total_chars += javascript_length(text) + if len(document_chunks) > 1: + docs_with_multiple_chunks += 1 + + stats = { + "docs": len(paths), + "total_chars": total_chars, + "chunks": len(chunks), + "chunk_chars": _chunk_char_percentiles(chunks), + "docs_with_more_than_one_chunk": docs_with_multiple_chunks, + } + return Corpus(tuple(chunks), tuple(query_candidates), stats) + + +def generate_queries(corpus: Corpus, query_count: int) -> list[str]: + """Cycles deterministic title and heading candidates to make realistic short queries.""" + + if query_count < 0: + raise ValueError("--queries must be zero or a positive integer") + if query_count == 0: + return [] + if not corpus.query_candidates: + raise ValueError("the selected corpus has no usable document title or heading for queries") + return [ + corpus.query_candidates[index % len(corpus.query_candidates)] + for index in range(query_count) + ] + + +def _normalize_in_place(vector: array[float]) -> None: + """Mirrors OpenKnowledge's Float32Array L2 normalization before vectors are returned.""" + + norm = 0.0 + for value in vector: + norm += value * value + norm = math.sqrt(norm) + if norm > 0.0: + for index, value in enumerate(vector): + vector[index] = value / norm + + +def _reported_total_tokens(payload: dict[str, Any]) -> int | float: + usage = payload.get("usage") + if not isinstance(usage, dict): + return 0 + total_tokens = usage.get("total_tokens", 0) + if isinstance(total_tokens, bool) or not isinstance(total_tokens, (int, float)): + return 0 + return total_tokens + + +def _consume_embedding_response( + payload: Any, + *, + expected_count: int, + expected_dimensions: int, +) -> int | float: + """Validates, orders, Float32-casts, and normalizes as OpenKnowledge does.""" + + if not isinstance(payload, dict): + raise EmbeddingResponseError("embeddings response was not a JSON object") + data = payload.get("data") + if not isinstance(data, list) or len(data) != expected_count: + got = len(data) if isinstance(data, list) else 0 + raise EmbeddingResponseError( + f"embeddings response had {got} vectors, expected {expected_count}" + ) + if not all(isinstance(item, dict) for item in data): + raise EmbeddingResponseError("embeddings response contained a non-object vector item") + + ordered = sorted(data, key=lambda item: item.get("index", 0)) + for item in ordered: + embedding = item.get("embedding") + if not isinstance(embedding, list): + raise EmbeddingResponseError("embeddings response contained a non-array embedding") + if len(embedding) != expected_dimensions: + raise EmbeddingResponseError( + "embeddings response had " + f"{len(embedding)} dimensions, expected {expected_dimensions}" + ) + try: + vector = array("f", embedding) + except (TypeError, ValueError, OverflowError) as exc: + raise EmbeddingResponseError( + "embeddings response contained a non-numeric embedding" + ) from exc + _normalize_in_place(vector) + return _reported_total_tokens(payload) + + +def _sleep_duration_for_retry(retry_number: int) -> float: + """Matches the source embedder's jittered 500 ms exponential backoff.""" + + ceiling_ms = BACKOFF_BASE_SECONDS * 1_000.0 * 2 ** (retry_number - 1) + delay_ms = ceiling_ms / 2.0 + random.random() * (ceiling_ms / 2.0) + return math.floor(delay_ms + 0.5) / 1_000.0 + + +def _sanitize_lone_surrogates(text: str) -> str: + """Replaces unpaired UTF-16 surrogate code units before UTF-8 encoding.""" + + if not any(0xD800 <= ord(character) <= 0xDFFF for character in text): + return text + + sanitized: list[str] = [] + index = 0 + while index < len(text): + code_point = ord(text[index]) + if 0xD800 <= code_point <= 0xDBFF: + if index + 1 < len(text) and 0xDC00 <= ord(text[index + 1]) <= 0xDFFF: + low_surrogate = ord(text[index + 1]) + astral_code_point = ( + 0x10000 + (code_point - 0xD800) * 0x400 + low_surrogate - 0xDC00 + ) + sanitized.append(chr(astral_code_point)) + index += 2 + continue + sanitized.append("\ufffd") + elif 0xDC00 <= code_point <= 0xDFFF: + sanitized.append("\ufffd") + else: + sanitized.append(text[index]) + index += 1 + return "".join(sanitized) + + +class OpenAICompatibleClient: + """A standard-library HTTP port of OpenKnowledge's OpenAI-compatible embedder.""" + + def __init__(self, config: ProviderConfig) -> None: + self.config = config + self.endpoint = f"{config.base_url.rstrip('/')}/embeddings" + endpoint = urllib.parse.urlsplit(self.endpoint) + if endpoint.scheme not in {"http", "https"} or endpoint.hostname is None: + raise ValueError("provider base_url must be an absolute HTTP or HTTPS URL") + self._connection_type = ( + http.client.HTTPSConnection + if endpoint.scheme == "https" + else http.client.HTTPConnection + ) + self._host = endpoint.hostname + self._port = endpoint.port + self._request_target = endpoint.path or "/" + if endpoint.query: + self._request_target = f"{self._request_target}?{endpoint.query}" + self._connection: http.client.HTTPConnection | None = None + self._connection_was_used = False + self.reconnects = 0 + self.traffic = HttpBodyTraffic() + + def close(self) -> None: + """Closes the provider connection once its workload is complete.""" + + if self._connection is not None: + self._connection.close() + self._connection = None + self._connection_was_used = False + + def _new_connection(self, timeout_seconds: float) -> http.client.HTTPConnection: + return self._connection_type(self._host, self._port, timeout=timeout_seconds) + + @staticmethod + def _socket_is_closed(connection: http.client.HTTPConnection) -> bool: + socket = connection.sock + if socket is None: + return True + try: + return socket.fileno() < 0 + except OSError: + return True + + @staticmethod + def _set_timeout(connection: http.client.HTTPConnection, timeout_seconds: float) -> None: + connection.timeout = timeout_seconds + if connection.sock is not None: + # Documents run before queries, so retain the socket and update its timeout in place. + connection.sock.settimeout(timeout_seconds) + + def _connection_for_request( + self, timeout_seconds: float + ) -> tuple[http.client.HTTPConnection, bool]: + if self._connection is None: + self._connection = self._new_connection(timeout_seconds) + return self._connection, False + + if self._connection_was_used and self._socket_is_closed(self._connection): + self.close() + self._connection = self._new_connection(timeout_seconds) + return self._connection, True + + self._set_timeout(self._connection, timeout_seconds) + return self._connection, False + + def _reconnect(self, timeout_seconds: float) -> None: + self.close() + self._connection = self._new_connection(timeout_seconds) + self.reconnects += 1 + + def _post_once(self, body: bytes, timeout_seconds: float) -> bytes: + reconnect_attempted = False + while True: + connection, needs_reconnect = self._connection_for_request(timeout_seconds) + if needs_reconnect: + self.reconnects += 1 + reconnect_attempted = True + + self.traffic.request_body_bytes += len(body) + try: + connection.request( + "POST", + self._request_target, + body=body, + headers={ + "Content-Type": "application/json", + "Authorization": f"Bearer {self.config.api_key}", + }, + ) + self._connection_was_used = True + response = connection.getresponse() + response_body = response.read() + self.traffic.response_body_bytes += len(response_body) + status = response.status + except ( + http.client.RemoteDisconnected, + http.client.BadStatusLine, + http.client.CannotSendRequest, + http.client.ResponseNotReady, + BrokenPipeError, + ConnectionResetError, + ): + if reconnect_attempted: + raise + self._reconnect(timeout_seconds) + reconnect_attempted = True + continue + except OSError: + if not self._socket_is_closed(connection) or reconnect_attempted: + raise + self._reconnect(timeout_seconds) + reconnect_attempted = True + continue + + if not 200 <= status < 300: + if status in RETRYABLE_STATUSES: + raise _RetryableRequestError(f"embeddings request failed: HTTP {status}") + raise EmbeddingRequestError(f"embeddings request failed: HTTP {status}") + return response_body + + def embed_batch(self, texts: Sequence[str], role: EmbeddingRole) -> BatchResult: + """Issues one request with the source embedder's body, timeout, and retry policy.""" + + body_dict: dict[str, Any] = { + "model": self.config.model, + "input": [_sanitize_lone_surrogates(text) for text in texts], + "encoding_format": "float", + } + if self.config.requested_dimensions is not None: + body_dict["dimensions"] = self.config.requested_dimensions + body = json.dumps(body_dict, ensure_ascii=False, separators=(",", ":")).encode("utf-8") + timeout_seconds = DOCUMENT_TIMEOUT_SECONDS if role == "document" else QUERY_TIMEOUT_SECONDS + started = time.monotonic() + attempt = 0 + retries = 0 + reconnects_before = self.reconnects + while True: + try: + response_body = self._post_once(body, timeout_seconds) + payload = json.loads(response_body.decode("utf-8")) + reported_total_tokens = _consume_embedding_response( + payload, + expected_count=len(texts), + expected_dimensions=self.config.expected_dimensions, + ) + return BatchResult( + time.monotonic() - started, + retries, + self.reconnects - reconnects_before, + reported_total_tokens, + ) + except _RetryableRequestError as exc: + failure: BaseException = exc + except ( + http.client.HTTPException, + urllib.error.URLError, + TimeoutError, + OSError, + UnicodeDecodeError, + json.JSONDecodeError, + ) as exc: + failure = exc + except EmbeddingResponseError as exc: + raise EmbeddingRequestError("embeddings response was malformed") from exc + + if attempt >= MAX_RETRIES: + raise EmbeddingRequestError("embeddings request failed after retries") from failure + attempt += 1 + retries += 1 + time.sleep(_sleep_duration_for_retry(attempt)) + + +def _traffic_delta(before: TrafficSnapshot, after: TrafficSnapshot) -> TrafficSnapshot: + return TrafficSnapshot( + request_body_bytes=after.request_body_bytes - before.request_body_bytes, + response_body_bytes=after.response_body_bytes - before.response_body_bytes, + ) + + +def _is_loopback_url(url: str) -> bool: + hostname = urllib.parse.urlparse(url).hostname + return hostname in {"localhost", "127.0.0.1", "::1"} + + +def _traffic_report(traffic: TrafficSnapshot, *, loopback: bool) -> dict[str, int]: + total = traffic.request_response_body_bytes + return { + "request_body_bytes": traffic.request_body_bytes, + "response_body_bytes": traffic.response_body_bytes, + "request_response_body_bytes": total, + "external_request_response_body_bytes": 0 if loopback else total, + "loopback_request_response_body_bytes": total if loopback else 0, + } + + +def _rate(count: int, elapsed_seconds: float) -> float | None: + return None if elapsed_seconds <= 0.0 else _round(count / elapsed_seconds) + + +def measure_document_embeddings( + client: OpenAICompatibleClient, + batches: Sequence[Sequence[str]], + *, + document_count: int, +) -> dict[str, Any]: + """Embeds every document chunk sequentially, exactly as OpenKnowledge's embed() does.""" + + started = time.monotonic() + latencies: list[float] = [] + retries = 0 + reconnects = 0 + reported_total_tokens: int | float = 0 + chunk_count = 0 + for batch in batches: + result = client.embed_batch(batch, "document") + latencies.append(result.duration_seconds) + retries += result.retry_count + reconnects += result.reconnects + reported_total_tokens += result.reported_total_tokens + chunk_count += len(batch) + elapsed = time.monotonic() - started + return { + "wall_seconds": _round(elapsed), + "chunks": chunk_count, + "documents": document_count, + "chunks_per_second": _rate(chunk_count, elapsed), + "documents_per_second": _rate(document_count, elapsed), + "batch_count": len(batches), + "batch_latency_ms": _latency_percentiles_ms(latencies), + "retry_count": retries, + "reconnects": reconnects, + "reported_usage_total_tokens": reported_total_tokens, + } + + +def measure_query_embeddings( + client: OpenAICompatibleClient, + queries: Sequence[str], +) -> dict[str, Any]: + """Embeds one query per sequential request, matching OpenKnowledge search behavior.""" + + started = time.monotonic() + latencies: list[float] = [] + retries = 0 + reconnects = 0 + reported_total_tokens: int | float = 0 + for query in queries: + result = client.embed_batch([query], "query") + latencies.append(result.duration_seconds) + retries += result.retry_count + reconnects += result.reconnects + reported_total_tokens += result.reported_total_tokens + elapsed = time.monotonic() - started + return { + "wall_seconds": _round(elapsed), + "query_count": len(queries), + "latency_ms": _latency_percentiles_ms(latencies), + "retry_count": retries, + "reconnects": reconnects, + "reported_usage_total_tokens": reported_total_tokens, + } + + +def _warmup_document_batch( + client: OpenAICompatibleClient, + batches: Sequence[Sequence[str]], +) -> dict[str, Any]: + if not batches: + return {"performed": False} + result = client.embed_batch(batches[0], "document") + return { + "performed": True, + "input_count": len(batches[0]), + "latency_ms": _round(result.duration_seconds * 1_000.0), + "retry_count": result.retry_count, + "reconnects": result.reconnects, + "reported_usage_total_tokens": result.reported_total_tokens, + } + + +def run_provider_workload( + config: ProviderConfig, + *, + document_batches: Sequence[Sequence[str]], + document_count: int, + queries: Sequence[str], + cold_start_seconds: float | None = None, +) -> dict[str, Any]: + """Runs a provider's warmup, document workload, and sequential query workload.""" + + client = OpenAICompatibleClient(config) + loopback = _is_loopback_url(config.base_url) + + warmup_before = client.traffic.snapshot() + warmup = _warmup_document_batch(client, document_batches) + warmup_after = client.traffic.snapshot() + + workload_before = client.traffic.snapshot() + documents = measure_document_embeddings( + client, + document_batches, + document_count=document_count, + ) + query_metrics = measure_query_embeddings(client, queries) + workload_after = client.traffic.snapshot() + + warmup["http_body_traffic"] = _traffic_report( + _traffic_delta(warmup_before, warmup_after), + loopback=loopback, + ) + workload_traffic = _traffic_report( + _traffic_delta(workload_before, workload_after), + loopback=loopback, + ) + reported_workload_tokens = ( + documents["reported_usage_total_tokens"] + query_metrics["reported_usage_total_tokens"] + ) + cost_usd = ( + reported_workload_tokens / 1_000_000.0 * OPENAI_PRICE_USD_PER_MILLION_TOKENS + if config.name == "openai" + else 0.0 + ) + return { + "model": config.model, + "dimensions": config.expected_dimensions, + "requested_dimensions": config.requested_dimensions, + "cold_start_seconds": _round(cold_start_seconds), + "warmup": warmup, + "documents": documents, + "queries": query_metrics, + "token_usage": { + "document_reported_usage_total_tokens": documents["reported_usage_total_tokens"], + "query_reported_usage_total_tokens": query_metrics["reported_usage_total_tokens"], + "workload_reported_usage_total_tokens": reported_workload_tokens, + "label": config.token_usage_label, + }, + "egress": workload_traffic, + "cost_usd": _round(cost_usd, 12), + "cost_scope": "document and query workload only; warmup is excluded", + } + + +def _server_failure_message(process: subprocess.Popen[str], log: Any) -> str: + log.seek(0) + output = log.read() + tail = output[-4_000:] + return f"lodedb serve exited with code {process.returncode}: {tail}" + + +@contextmanager +def running_lodedb_server( + *, + path: Path = Path("/root/okbench-store"), + port: int = 8_099, + startup_timeout_seconds: float = 900.0, +) -> Iterator[LocalLodeDBServer]: + """Starts the benchmark's CPU ONNX endpoint and measures readiness from Popen.""" + + log = tempfile.TemporaryFile(mode="w+", encoding="utf-8") + command = [ + "lodedb", + "serve", + "--model", + "minilm", + "--port", + str(port), + "--path", + str(path), + "--device", + "cpu", + "--runtime", + "onnx", + ] + started = time.monotonic() + process = subprocess.Popen(command, stdout=log, stderr=subprocess.STDOUT, text=True) + health_url = f"http://127.0.0.1:{port}/healthz" + try: + while True: + if process.poll() is not None: + raise RuntimeError(_server_failure_message(process, log)) + try: + with urllib.request.urlopen(health_url, timeout=1.0) as response: + if response.status == 200: + yield LocalLodeDBServer( + base_url=f"http://127.0.0.1:{port}/v1", + cold_start_seconds=time.monotonic() - started, + ) + break + except (urllib.error.URLError, TimeoutError, OSError): + pass + if time.monotonic() - started >= startup_timeout_seconds: + raise TimeoutError("lodedb serve did not become healthy before the startup timeout") + time.sleep(0.25) + finally: + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=15.0) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=15.0) + log.close() + + +def run_benchmark( + *, + corpus_root: Path, + docs: int, + query_count: int, + provider: str, + openai_api_key: str | None, + lodedb_store_path: Path = Path("/root/okbench-store"), + cpu_count: int | None = None, + corpus_revision: str | None = None, +) -> dict[str, Any]: + """Runs requested providers in one process, always placing LodeDB before OpenAI.""" + + if provider not in {"lodedb", "openai", "both"}: + raise ValueError("--provider must be one of: lodedb, openai, both") + corpus = load_kubernetes_corpus(corpus_root, docs) + queries = generate_queries(corpus, query_count) + document_batches = batch_inputs(corpus.chunks) + results: dict[str, Any] = { + "config": { + "docs_flag": docs, + "docs_selected": corpus.stats["docs"], + "queries_flag": query_count, + "queries_selected": len(queries), + "provider_flag": provider, + "cpu_count": cpu_count if cpu_count is not None else os.cpu_count(), + "document_batch_count": len(document_batches), + "document_batch_size": DEFAULT_MAX_BATCH_SIZE, + "document_batch_chars": DEFAULT_MAX_BATCH_CHARS, + "corpus_revision": corpus_revision, + }, + "corpus": corpus.stats, + "providers": {}, + } + + # Fairness: both providers use this chunker, batcher, retry policy, and HTTP client. + if provider in {"lodedb", "both"}: + with running_lodedb_server(path=lodedb_store_path) as server: + results["providers"]["lodedb"] = run_provider_workload( + ProviderConfig( + name="lodedb", + base_url=server.base_url, + model="minilm", + requested_dimensions=LODEDB_DIMENSIONS, + expected_dimensions=LODEDB_DIMENSIONS, + api_key="local-benchmark-placeholder", + token_usage_label="LodeDB server estimate from character count", + ), + document_batches=document_batches, + document_count=corpus.stats["docs"], + queries=queries, + cold_start_seconds=server.cold_start_seconds, + ) + + if provider in {"openai", "both"}: + if not openai_api_key: + raise RuntimeError("OPENAI_API_KEY is required for the openai provider") + # Remote OpenAI is measured from a datacenter, which is favorable to its network latency. + # OpenKnowledge's config must match model dimensions, so real users receive LodeDB's + # 384 dimensions and OpenAI's native 1536 dimensions rather than an artificial match. + results["providers"]["openai"] = run_provider_workload( + ProviderConfig( + name="openai", + base_url="https://api.openai.com/v1", + model="text-embedding-3-small", + requested_dimensions=None, + expected_dimensions=OPENAI_DIMENSIONS, + api_key=openai_api_key, + token_usage_label="OpenAI provider-reported token total", + ), + document_batches=document_batches, + document_count=corpus.stats["docs"], + queries=queries, + ) + return results + + +def _local_smoke_documents(count: int) -> list[str]: + """Creates non-empty representative Markdown inputs without requiring the wiki clone.""" + + topics = ( + "deployment rollouts", + "service discovery", + "cluster networking", + "pod scheduling", + "persistent storage", + ) + return [ + f"# Kubernetes {topics[index % len(topics)]}\n\n" + f"Local embedding smoke document {index}. " + "This verifies OpenKnowledge-compatible document batching against LodeDB." + for index in range(count) + ] + + +def run_local_smoke(*, base_url: str, docs: int = 20) -> dict[str, Any]: + """Drives roughly 20 short documents through the exact document embedding path.""" + + if docs <= 0: + raise ValueError("--docs must be a positive integer for --local-smoke") + chunks = tuple( + chunk for document in _local_smoke_documents(docs) for chunk in chunk_document(document) + ) + batches = batch_inputs(chunks) + client = OpenAICompatibleClient( + ProviderConfig( + name="lodedb", + base_url=base_url, + model="minilm", + requested_dimensions=LODEDB_DIMENSIONS, + expected_dimensions=LODEDB_DIMENSIONS, + api_key="local-benchmark-placeholder", + token_usage_label="LodeDB server estimate from character count", + ) + ) + before = client.traffic.snapshot() + documents = measure_document_embeddings(client, batches, document_count=docs) + after = client.traffic.snapshot() + return { + "documents": documents, + "http_body_traffic": _traffic_report(_traffic_delta(before, after), loopback=True), + } + + +def _main() -> None: + parser = argparse.ArgumentParser(description="OpenKnowledge embedding benchmark local smoke") + parser.add_argument( + "--local-smoke", action="store_true", help="run against an existing local LodeDB server" + ) + parser.add_argument("--base-url", default="http://127.0.0.1:8099/v1") + parser.add_argument("--docs", type=int, default=20) + args = parser.parse_args() + if not args.local_smoke: + parser.error("pass --local-smoke to run the local client smoke test") + print( + json.dumps( + run_local_smoke(base_url=args.base_url, docs=args.docs), indent=2, sort_keys=True + ) + ) + + +if __name__ == "__main__": + _main() diff --git a/benchmarks/openknowledge_embeddings/modal_bench.py b/benchmarks/openknowledge_embeddings/modal_bench.py new file mode 100644 index 00000000..c9b11ca3 --- /dev/null +++ b/benchmarks/openknowledge_embeddings/modal_bench.py @@ -0,0 +1,133 @@ +"""Benchmark OpenKnowledge's two production embedding-provider choices on Modal. + +Launch this module from the repository root. Its local benchmark import relies on +the repository-root ``benchmarks`` namespace. + + modal run benchmarks/openknowledge_embeddings/modal_bench.py::bench --docs 500 + modal run benchmarks/openknowledge_embeddings/modal_bench.py::bench --provider lodedb +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +import modal + +from benchmarks.openknowledge_embeddings.bench_core import run_benchmark + +_REMOTE_SRC = "/root/lodedb-src" +_REMOTE_BENCHMARK_DIR = "/root/benchmarks/openknowledge_embeddings" +_KUBERNETES_WEBSITE = "/root/k8s-website" +KUBERNETES_WEBSITE_REVISION = "71d23f81e3479361befc94564e2b955860c03164" +_LODEDB_RUNTIME_DEPENDENCIES = ( + "numpy>=2.0.0,<3", + "typer>=0.12.0", + "pyyaml>=6.0.0", + "onnxruntime>=1.20.0,<2", + "transformers>=4.40.0,<5", +) + + +def _build_image() -> modal.Image: + """Builds the local LodeDB source with its CPU ONNX embedding dependencies.""" + + image = ( + modal.Image.from_registry("ubuntu:22.04", add_python="3.11") + .apt_install("build-essential", "curl", "git", "pkg-config", "libopenblas-dev") + # transformers resolves tokenizers and huggingface-hub for MiniLM download and tokenization. + .pip_install(*_LODEDB_RUNTIME_DEPENDENCIES) + .run_commands( + "curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | " + "sh -s -- -y --default-toolchain stable --profile minimal" + ) + .env({"PYTHONPATH": "/root"}) + ) + + # Modal imports this module again in the container. These local paths only exist while + # `modal run` constructs the image, so leave the source-copy steps behind this guard. + if modal.is_local(): + repo_root = Path(__file__).resolve().parents[2] + for relative_path in ( + "pyproject.toml", + "README.md", + "LICENSE", + "NOTICE", + "Cargo.toml", + "Cargo.lock", + ): + image = image.add_local_file( + str(repo_root / relative_path), + remote_path=f"{_REMOTE_SRC}/{relative_path}", + copy=True, + ) + image = image.add_local_dir( + str(repo_root / "third_party" / "turbovec"), + remote_path=f"{_REMOTE_SRC}/third_party/turbovec", + copy=True, + ignore=["**/target/**", "**/__pycache__/**", "**/*.so", "**/*.pyd", "**/*.dylib"], + ) + # The Python extension links lodedb-core by relative path. Its manifest inherits + # package fields from the root Cargo workspace and lodedb-core links lodedb-gpu. + image = image.add_local_dir( + str(repo_root / "crates"), + remote_path=f"{_REMOTE_SRC}/crates", + copy=True, + ignore=["**/target/**", "**/__pycache__/**"], + ) + image = image.add_local_dir( + str(repo_root / "src"), + remote_path=f"{_REMOTE_SRC}/src", + copy=True, + ignore=["**/__pycache__/**", "**/*.pyc", "**/*.so", "**/*.pyd", "**/*.dylib"], + ) + image = image.add_local_dir( + str(Path(__file__).resolve().parent), + remote_path=_REMOTE_BENCHMARK_DIR, + copy=True, + ignore=["**/__pycache__/**", "**/*.pyc", "results/**"], + ) + + return image.run_commands( + 'PATH="$HOME/.cargo/bin:$PATH" python -m pip install --no-deps /root/lodedb-src', + f"git init {_KUBERNETES_WEBSITE}", + f"git -C {_KUBERNETES_WEBSITE} remote add origin https://github.com/kubernetes/website", + f"git -C {_KUBERNETES_WEBSITE} fetch --depth 1 origin {KUBERNETES_WEBSITE_REVISION}", + f"git -C {_KUBERNETES_WEBSITE} checkout --detach FETCH_HEAD", + ) + + +IMAGE = _build_image() +app = modal.App("lodedb-openknowledge-embeddings", image=IMAGE) +OPENAI_SECRET = modal.Secret.from_name("openai-embeddings-bench") + + +@app.function(cpu=8.0, memory=16384, timeout=3600, secrets=[OPENAI_SECRET]) +def run_remote(docs: int, queries: int, provider: str) -> dict[str, Any]: + """Runs selected providers in one CPU container, with LodeDB ordered first.""" + + return run_benchmark( + corpus_root=Path(_KUBERNETES_WEBSITE), + docs=docs, + query_count=queries, + provider=provider, + openai_api_key=os.environ.get("OPENAI_API_KEY"), + lodedb_store_path=Path("/root/okbench-store"), + cpu_count=os.cpu_count(), + corpus_revision=KUBERNETES_WEBSITE_REVISION, + ) + + +def _print_results(results: dict[str, Any]) -> None: + print("BEGIN_RESULTS_JSON") + print(json.dumps(results, indent=2, sort_keys=True)) + print("END_RESULTS_JSON") + + +@app.local_entrypoint() +def bench(docs: int = 0, queries: int = 100, provider: str = "both") -> None: + """Runs the benchmark. ``docs=0`` selects the full sorted Kubernetes Markdown corpus.""" + + _print_results(run_remote.remote(docs=docs, queries=queries, provider=provider)) diff --git a/benchmarks/openknowledge_embeddings/okchunk.py b/benchmarks/openknowledge_embeddings/okchunk.py new file mode 100644 index 00000000..0d9abf5b --- /dev/null +++ b/benchmarks/openknowledge_embeddings/okchunk.py @@ -0,0 +1,121 @@ +"""OpenKnowledge's dependency-free document chunker, ported from TypeScript.""" + +from __future__ import annotations + +from collections.abc import Iterable + +CHUNK_TARGET_CHARS = 8_000 +CHUNK_OVERLAP_CHARS = 400 +MAX_CHUNKS_PER_DOC = 80 +CHUNK_CONFIG_ID = f"c{CHUNK_TARGET_CHARS}-o{CHUNK_OVERLAP_CHARS}-m{MAX_CHUNKS_PER_DOC}" + +# ECMAScript WhiteSpace plus LineTerminator code points used by String.prototype.trim(). +_JS_TRIM_CODE_UNITS = frozenset( + "\u0009\u000a\u000b\u000c\u000d\u0020\u00a0\u1680" + "\u2000\u2001\u2002\u2003\u2004\u2005\u2006\u2007" + "\u2008\u2009\u200a\u2028\u2029\u202f\u205f\u3000\ufeff" +) + + +def _to_js_code_units(text: str) -> str: + """Returns a Python string whose characters are JavaScript UTF-16 code units.""" + + raw = text.encode("utf-16-le", "surrogatepass") + return "".join( + chr(int.from_bytes(raw[index : index + 2], "little")) for index in range(0, len(raw), 2) + ) + + +def _from_js_code_units(code_units: str) -> str: + """Reassembles a Python string from JavaScript UTF-16 code units.""" + + raw = b"".join(unit.encode("utf-16-le", "surrogatepass") for unit in code_units) + return raw.decode("utf-16-le", "surrogatepass") + + +def _js_trim(code_units: str) -> str: + """Implements String.prototype.trim() over the UTF-16 representation.""" + + start = 0 + end = len(code_units) + while start < end and code_units[start] in _JS_TRIM_CODE_UNITS: + start += 1 + while end > start and code_units[end - 1] in _JS_TRIM_CODE_UNITS: + end -= 1 + return code_units[start:end] + + +def javascript_length(text: str) -> int: + """Returns ``text.length`` as JavaScript would report it.""" + + return len(text.encode("utf-16-le", "surrogatepass")) // 2 + + +def _last_index_of_boundary(code_units: str, end: int) -> int: + """Matches Math.max(text.lastIndexOf(' ', end), text.lastIndexOf('\\n', end)).""" + + return max(code_units.rfind(" ", 0, end + 1), code_units.rfind("\n", 0, end + 1)) + + +def chunk_document( + text: str, + *, + target_chars: int | None = None, + overlap_chars: int | None = None, + max_chunks: int | None = None, +) -> list[str]: + """Splits text exactly as OpenKnowledge's ``chunkDocument`` does. + + The TypeScript implementation measures character positions in UTF-16 code units. + Keeping that behavior matters for documents containing astral Unicode characters. + """ + + target = max(1, CHUNK_TARGET_CHARS if target_chars is None else target_chars) + requested_overlap = CHUNK_OVERLAP_CHARS if overlap_chars is None else overlap_chars + overlap = max(0, min(requested_overlap, target - 1)) + chunk_cap = MAX_CHUNKS_PER_DOC if max_chunks is None else max_chunks + source = _to_js_code_units(text) + + if not _js_trim(source): + return [] + if len(source) <= target: + return [_from_js_code_units(_js_trim(source))] + + chunks: list[str] = [] + start = 0 + while start < len(source) and len(chunks) < chunk_cap: + end = min(len(source), start + target) + if end < len(source): + boundary = _last_index_of_boundary(source, end) + if boundary > start + target // 2: + end = boundary + piece = _js_trim(source[start:end]) + if piece: + chunks.append(_from_js_code_units(piece)) + if end >= len(source): + break + next_start = end - overlap + start = next_start if next_start > start else end + return chunks + + +def chunkDocument( # noqa: N802 - matches the OpenKnowledge public helper name. + text: str, + options: dict[str, int] | None = None, +) -> list[str]: + """Compatibility spelling and options shape of OpenKnowledge's TypeScript helper.""" + + values: Iterable[tuple[str, int]] = () if options is None else options.items() + kwargs = { + "target_chars": None, + "overlap_chars": None, + "max_chunks": None, + } + for key, value in values: + if key == "targetChars": + kwargs["target_chars"] = value + elif key == "overlapChars": + kwargs["overlap_chars"] = value + elif key == "maxChunks": + kwargs["max_chunks"] = value + return chunk_document(text, **kwargs) diff --git a/benchmarks/openknowledge_embeddings/results/modal-300docs.json b/benchmarks/openknowledge_embeddings/results/modal-300docs.json new file mode 100644 index 00000000..c6bdcdbf --- /dev/null +++ b/benchmarks/openknowledge_embeddings/results/modal-300docs.json @@ -0,0 +1,155 @@ +{ + "config": { + "corpus_revision": "71d23f81e3479361befc94564e2b955860c03164", + "cpu_count": 24, + "docs_flag": 300, + "docs_selected": 300, + "document_batch_chars": 96000, + "document_batch_count": 30, + "document_batch_size": 96, + "provider_flag": "both", + "queries_flag": 50, + "queries_selected": 50 + }, + "corpus": { + "chunk_chars": { + "max": 8000, + "p50": 6113.5, + "p95": 7999.0 + }, + "chunks": 490, + "docs": 300, + "docs_with_more_than_one_chunk": 147, + "total_chars": 2658713 + }, + "providers": { + "lodedb": { + "cold_start_seconds": 3.277098, + "cost_scope": "document and query workload only; warmup is excluded", + "cost_usd": 0.0, + "dimensions": 384, + "documents": { + "batch_count": 30, + "batch_latency_ms": { + "max": 1066.366488, + "p50": 910.105389, + "p95": 1048.901096 + }, + "chunks": 490, + "chunks_per_second": 18.125954, + "documents": 300, + "documents_per_second": 11.097523, + "reconnects": 0, + "reported_usage_total_tokens": 683529, + "retry_count": 0, + "wall_seconds": 27.03306 + }, + "egress": { + "external_request_response_body_bytes": 0, + "loopback_request_response_body_bytes": 7382070, + "request_body_bytes": 2795771, + "request_response_body_bytes": 7382070, + "response_body_bytes": 4586299 + }, + "model": "minilm", + "queries": { + "latency_ms": { + "max": 10.304891, + "p50": 5.569571, + "p95": 8.662564 + }, + "query_count": 50, + "reconnects": 0, + "reported_usage_total_tokens": 460, + "retry_count": 0, + "wall_seconds": 0.302266 + }, + "requested_dimensions": 384, + "token_usage": { + "document_reported_usage_total_tokens": 683529, + "label": "LodeDB server estimate from character count", + "query_reported_usage_total_tokens": 460, + "workload_reported_usage_total_tokens": 683989 + }, + "warmup": { + "http_body_traffic": { + "external_request_response_body_bytes": 0, + "loopback_request_response_body_bytes": 283675, + "request_body_bytes": 97014, + "request_response_body_bytes": 283675, + "response_body_bytes": 186661 + }, + "input_count": 22, + "latency_ms": 3151.348862, + "performed": true, + "reconnects": 0, + "reported_usage_total_tokens": 23761, + "retry_count": 0 + } + }, + "openai": { + "cold_start_seconds": null, + "cost_scope": "document and query workload only; warmup is excluded", + "cost_usd": 0.01274128, + "dimensions": 1536, + "documents": { + "batch_count": 30, + "batch_latency_ms": { + "max": 904.928219, + "p50": 312.712834, + "p95": 663.50281 + }, + "chunks": 490, + "chunks_per_second": 44.143612, + "documents": 300, + "documents_per_second": 27.026701, + "reconnects": 0, + "reported_usage_total_tokens": 636674, + "retry_count": 0, + "wall_seconds": 11.100134 + }, + "egress": { + "external_request_response_body_bytes": 18736627, + "loopback_request_response_body_bytes": 0, + "request_body_bytes": 2795691, + "request_response_body_bytes": 18736627, + "response_body_bytes": 15940936 + }, + "model": "text-embedding-3-small", + "queries": { + "latency_ms": { + "max": 399.068839, + "p50": 138.903222, + "p95": 259.934613 + }, + "query_count": 50, + "reconnects": 0, + "reported_usage_total_tokens": 390, + "retry_count": 0, + "wall_seconds": 7.720582 + }, + "requested_dimensions": null, + "token_usage": { + "document_reported_usage_total_tokens": 636674, + "label": "OpenAI provider-reported token total", + "query_reported_usage_total_tokens": 390, + "workload_reported_usage_total_tokens": 637064 + }, + "warmup": { + "http_body_traffic": { + "external_request_response_body_bytes": 745930, + "loopback_request_response_body_bytes": 0, + "request_body_bytes": 97013, + "request_response_body_bytes": 745930, + "response_body_bytes": 648917 + }, + "input_count": 22, + "latency_ms": 1248.072161, + "performed": true, + "reconnects": 0, + "reported_usage_total_tokens": 23035, + "retry_count": 0 + } + } + } +} \ No newline at end of file diff --git a/benchmarks/openknowledge_embeddings/results/modal-full-corpus-lodedb-repeat-1.json b/benchmarks/openknowledge_embeddings/results/modal-full-corpus-lodedb-repeat-1.json new file mode 100644 index 00000000..ed90b32f --- /dev/null +++ b/benchmarks/openknowledge_embeddings/results/modal-full-corpus-lodedb-repeat-1.json @@ -0,0 +1,91 @@ +{ + "config": { + "corpus_revision": "71d23f81e3479361befc94564e2b955860c03164", + "cpu_count": 24, + "docs_flag": 0, + "docs_selected": 2441, + "document_batch_chars": 96000, + "document_batch_count": 257, + "document_batch_size": 96, + "provider_flag": "lodedb", + "queries_flag": 100, + "queries_selected": 100 + }, + "corpus": { + "chunk_chars": { + "max": 8000, + "p50": 6791, + "p95": 7999.0 + }, + "chunks": 4485, + "docs": 2441, + "docs_with_more_than_one_chunk": 836, + "total_chars": 22746024 + }, + "providers": { + "lodedb": { + "cold_start_seconds": 2.018696, + "cost_scope": "document and query workload only; warmup is excluded", + "cost_usd": 0.0, + "dimensions": 384, + "documents": { + "batch_count": 257, + "batch_latency_ms": { + "max": 1134.786868, + "p50": 621.547208, + "p95": 768.066616 + }, + "chunks": 4485, + "chunks_per_second": 27.04414, + "documents": 2441, + "documents_per_second": 14.719007, + "reconnects": 0, + "reported_usage_total_tokens": 5889110, + "retry_count": 0, + "wall_seconds": 165.839995 + }, + "egress": { + "external_request_response_body_bytes": 0, + "loopback_request_response_body_bytes": 63139979, + "request_body_bytes": 24218495, + "request_response_body_bytes": 63139979, + "response_body_bytes": 38921484 + }, + "model": "minilm", + "queries": { + "latency_ms": { + "max": 7.021624, + "p50": 4.838521, + "p95": 5.851534 + }, + "query_count": 100, + "reconnects": 0, + "reported_usage_total_tokens": 1099, + "retry_count": 0, + "wall_seconds": 0.487128 + }, + "requested_dimensions": 384, + "token_usage": { + "document_reported_usage_total_tokens": 5889110, + "label": "LodeDB server estimate from character count", + "query_reported_usage_total_tokens": 1099, + "workload_reported_usage_total_tokens": 5890209 + }, + "warmup": { + "http_body_traffic": { + "external_request_response_body_bytes": 0, + "loopback_request_response_body_bytes": 283675, + "request_body_bytes": 97014, + "request_response_body_bytes": 283675, + "response_body_bytes": 186661 + }, + "input_count": 22, + "latency_ms": 2539.229973, + "performed": true, + "reconnects": 0, + "reported_usage_total_tokens": 23761, + "retry_count": 0 + } + } + } +} \ No newline at end of file diff --git a/benchmarks/openknowledge_embeddings/results/modal-full-corpus-lodedb-repeat-2.json b/benchmarks/openknowledge_embeddings/results/modal-full-corpus-lodedb-repeat-2.json new file mode 100644 index 00000000..b0d32dc8 --- /dev/null +++ b/benchmarks/openknowledge_embeddings/results/modal-full-corpus-lodedb-repeat-2.json @@ -0,0 +1,91 @@ +{ + "config": { + "corpus_revision": "71d23f81e3479361befc94564e2b955860c03164", + "cpu_count": 24, + "docs_flag": 0, + "docs_selected": 2441, + "document_batch_chars": 96000, + "document_batch_count": 257, + "document_batch_size": 96, + "provider_flag": "lodedb", + "queries_flag": 100, + "queries_selected": 100 + }, + "corpus": { + "chunk_chars": { + "max": 8000, + "p50": 6791, + "p95": 7999.0 + }, + "chunks": 4485, + "docs": 2441, + "docs_with_more_than_one_chunk": 836, + "total_chars": 22746024 + }, + "providers": { + "lodedb": { + "cold_start_seconds": 3.775641, + "cost_scope": "document and query workload only; warmup is excluded", + "cost_usd": 0.0, + "dimensions": 384, + "documents": { + "batch_count": 257, + "batch_latency_ms": { + "max": 1893.782014, + "p50": 1331.114599, + "p95": 1691.71827 + }, + "chunks": 4485, + "chunks_per_second": 12.607785, + "documents": 2441, + "documents_per_second": 6.861896, + "reconnects": 0, + "reported_usage_total_tokens": 5889110, + "retry_count": 0, + "wall_seconds": 355.732581 + }, + "egress": { + "external_request_response_body_bytes": 0, + "loopback_request_response_body_bytes": 63139359, + "request_body_bytes": 24218495, + "request_response_body_bytes": 63139359, + "response_body_bytes": 38920864 + }, + "model": "minilm", + "queries": { + "latency_ms": { + "max": 8.832993, + "p50": 4.649489, + "p95": 7.00232 + }, + "query_count": 100, + "reconnects": 0, + "reported_usage_total_tokens": 1099, + "retry_count": 0, + "wall_seconds": 0.480675 + }, + "requested_dimensions": 384, + "token_usage": { + "document_reported_usage_total_tokens": 5889110, + "label": "LodeDB server estimate from character count", + "query_reported_usage_total_tokens": 1099, + "workload_reported_usage_total_tokens": 5890209 + }, + "warmup": { + "http_body_traffic": { + "external_request_response_body_bytes": 0, + "loopback_request_response_body_bytes": 283725, + "request_body_bytes": 97014, + "request_response_body_bytes": 283725, + "response_body_bytes": 186711 + }, + "input_count": 22, + "latency_ms": 4143.230979, + "performed": true, + "reconnects": 0, + "reported_usage_total_tokens": 23761, + "retry_count": 0 + } + } + } +} \ No newline at end of file diff --git a/benchmarks/openknowledge_embeddings/results/modal-full-corpus-lodedb-repeat-3.json b/benchmarks/openknowledge_embeddings/results/modal-full-corpus-lodedb-repeat-3.json new file mode 100644 index 00000000..31de4ec2 --- /dev/null +++ b/benchmarks/openknowledge_embeddings/results/modal-full-corpus-lodedb-repeat-3.json @@ -0,0 +1,91 @@ +{ + "config": { + "corpus_revision": "71d23f81e3479361befc94564e2b955860c03164", + "cpu_count": 24, + "docs_flag": 0, + "docs_selected": 2441, + "document_batch_chars": 96000, + "document_batch_count": 257, + "document_batch_size": 96, + "provider_flag": "lodedb", + "queries_flag": 100, + "queries_selected": 100 + }, + "corpus": { + "chunk_chars": { + "max": 8000, + "p50": 6791, + "p95": 7999.0 + }, + "chunks": 4485, + "docs": 2441, + "docs_with_more_than_one_chunk": 836, + "total_chars": 22746024 + }, + "providers": { + "lodedb": { + "cold_start_seconds": 1.76563, + "cost_scope": "document and query workload only; warmup is excluded", + "cost_usd": 0.0, + "dimensions": 384, + "documents": { + "batch_count": 257, + "batch_latency_ms": { + "max": 914.754153, + "p50": 708.495764, + "p95": 767.435396 + }, + "chunks": 4485, + "chunks_per_second": 24.438152, + "documents": 2441, + "documents_per_second": 13.300675, + "reconnects": 0, + "reported_usage_total_tokens": 5889110, + "retry_count": 0, + "wall_seconds": 183.524519 + }, + "egress": { + "external_request_response_body_bytes": 0, + "loopback_request_response_body_bytes": 63139979, + "request_body_bytes": 24218495, + "request_response_body_bytes": 63139979, + "response_body_bytes": 38921484 + }, + "model": "minilm", + "queries": { + "latency_ms": { + "max": 7.435699, + "p50": 4.03816, + "p95": 6.418055 + }, + "query_count": 100, + "reconnects": 0, + "reported_usage_total_tokens": 1099, + "retry_count": 0, + "wall_seconds": 0.425141 + }, + "requested_dimensions": 384, + "token_usage": { + "document_reported_usage_total_tokens": 5889110, + "label": "LodeDB server estimate from character count", + "query_reported_usage_total_tokens": 1099, + "workload_reported_usage_total_tokens": 5890209 + }, + "warmup": { + "http_body_traffic": { + "external_request_response_body_bytes": 0, + "loopback_request_response_body_bytes": 283675, + "request_body_bytes": 97014, + "request_response_body_bytes": 283675, + "response_body_bytes": 186661 + }, + "input_count": 22, + "latency_ms": 2885.97003, + "performed": true, + "reconnects": 0, + "reported_usage_total_tokens": 23761, + "retry_count": 0 + } + } + } +} \ No newline at end of file diff --git a/benchmarks/openknowledge_embeddings/results/modal-full-corpus.json b/benchmarks/openknowledge_embeddings/results/modal-full-corpus.json new file mode 100644 index 00000000..c405d196 --- /dev/null +++ b/benchmarks/openknowledge_embeddings/results/modal-full-corpus.json @@ -0,0 +1,155 @@ +{ + "config": { + "corpus_revision": "71d23f81e3479361befc94564e2b955860c03164", + "cpu_count": 24, + "docs_flag": 0, + "docs_selected": 2441, + "document_batch_chars": 96000, + "document_batch_count": 257, + "document_batch_size": 96, + "provider_flag": "both", + "queries_flag": 100, + "queries_selected": 100 + }, + "corpus": { + "chunk_chars": { + "max": 8000, + "p50": 6791, + "p95": 7999.0 + }, + "chunks": 4485, + "docs": 2441, + "docs_with_more_than_one_chunk": 836, + "total_chars": 22746024 + }, + "providers": { + "lodedb": { + "cold_start_seconds": 3.777442, + "cost_scope": "document and query workload only; warmup is excluded", + "cost_usd": 0.0, + "dimensions": 384, + "documents": { + "batch_count": 257, + "batch_latency_ms": { + "max": 1384.022654, + "p50": 1047.529633, + "p95": 1205.268481 + }, + "chunks": 4485, + "chunks_per_second": 16.299736, + "documents": 2441, + "documents_per_second": 8.871272, + "reconnects": 0, + "reported_usage_total_tokens": 5889110, + "retry_count": 0, + "wall_seconds": 275.157827 + }, + "egress": { + "external_request_response_body_bytes": 0, + "loopback_request_response_body_bytes": 63139359, + "request_body_bytes": 24218495, + "request_response_body_bytes": 63139359, + "response_body_bytes": 38920864 + }, + "model": "minilm", + "queries": { + "latency_ms": { + "max": 7.688164, + "p50": 4.529746, + "p95": 6.523913 + }, + "query_count": 100, + "reconnects": 0, + "reported_usage_total_tokens": 1099, + "retry_count": 0, + "wall_seconds": 0.457105 + }, + "requested_dimensions": 384, + "token_usage": { + "document_reported_usage_total_tokens": 5889110, + "label": "LodeDB server estimate from character count", + "query_reported_usage_total_tokens": 1099, + "workload_reported_usage_total_tokens": 5890209 + }, + "warmup": { + "http_body_traffic": { + "external_request_response_body_bytes": 0, + "loopback_request_response_body_bytes": 283725, + "request_body_bytes": 97014, + "request_response_body_bytes": 283725, + "response_body_bytes": 186711 + }, + "input_count": 22, + "latency_ms": 3696.469377, + "performed": true, + "reconnects": 0, + "reported_usage_total_tokens": 23761, + "retry_count": 0 + } + }, + "openai": { + "cold_start_seconds": null, + "cost_scope": "document and query workload only; warmup is excluded", + "cost_usd": 0.11610126, + "dimensions": 1536, + "documents": { + "batch_count": 257, + "batch_latency_ms": { + "max": 2693.70737, + "p50": 391.558117, + "p95": 741.686651 + }, + "chunks": 4485, + "chunks_per_second": 37.005168, + "documents": 2441, + "documents_per_second": 20.140382, + "reconnects": 0, + "reported_usage_total_tokens": 5804126, + "retry_count": 0, + "wall_seconds": 121.199289 + }, + "egress": { + "external_request_response_body_bytes": 159504761, + "loopback_request_response_body_bytes": 0, + "request_body_bytes": 24218138, + "request_response_body_bytes": 159504761, + "response_body_bytes": 135286623 + }, + "model": "text-embedding-3-small", + "queries": { + "latency_ms": { + "max": 2804.75847, + "p50": 173.323863, + "p95": 346.13049 + }, + "query_count": 100, + "reconnects": 0, + "reported_usage_total_tokens": 937, + "retry_count": 0, + "wall_seconds": 24.728703 + }, + "requested_dimensions": null, + "token_usage": { + "document_reported_usage_total_tokens": 5804126, + "label": "OpenAI provider-reported token total", + "query_reported_usage_total_tokens": 937, + "workload_reported_usage_total_tokens": 5805063 + }, + "warmup": { + "http_body_traffic": { + "external_request_response_body_bytes": 745905, + "loopback_request_response_body_bytes": 0, + "request_body_bytes": 97013, + "request_response_body_bytes": 745905, + "response_body_bytes": 648892 + }, + "input_count": 22, + "latency_ms": 1063.635641, + "performed": true, + "reconnects": 0, + "reported_usage_total_tokens": 23035, + "retry_count": 0 + } + } + } +} \ No newline at end of file diff --git a/benchmarks/openknowledge_embeddings/test_bench_core.py b/benchmarks/openknowledge_embeddings/test_bench_core.py new file mode 100644 index 00000000..bb084fef --- /dev/null +++ b/benchmarks/openknowledge_embeddings/test_bench_core.py @@ -0,0 +1,173 @@ +"""HTTP client regression checks for the OpenKnowledge embedding benchmark.""" + +import http.client +import json +import sys +from pathlib import Path + +_BENCHMARK_DIR = Path(__file__).resolve().parent +if str(_BENCHMARK_DIR) not in sys.path: + sys.path.insert(0, str(_BENCHMARK_DIR)) + +from bench_core import ( # noqa: E402 + DOCUMENT_TIMEOUT_SECONDS, + QUERY_TIMEOUT_SECONDS, + OpenAICompatibleClient, + ProviderConfig, +) +from okchunk import chunk_document # noqa: E402 + + +class _FakeSocket: + def __init__(self) -> None: + self.timeouts: list[float] = [] + + def fileno(self) -> int: + return 1 + + def settimeout(self, timeout_seconds: float) -> None: + self.timeouts.append(timeout_seconds) + + +class _FakeResponse: + def __init__(self, body: bytes) -> None: + self.status = 200 + self._body = body + self.was_read = False + + def read(self) -> bytes: + self.was_read = True + return self._body + + +class _FakeConnection: + def __init__( + self, + responses: list[_FakeResponse], + errors: list[BaseException] | None = None, + ) -> None: + self.sock: _FakeSocket | None = _FakeSocket() + self.timeout: float | None = None + self.requests: list[tuple[str, str, bytes, dict[str, str]]] = [] + self.responses = responses + self.errors = [] if errors is None else errors + self.closed = False + + def request( + self, + method: str, + target: str, + *, + body: bytes, + headers: dict[str, str], + ) -> None: + self.requests.append((method, target, body, headers)) + if self.errors: + raise self.errors.pop(0) + + def getresponse(self) -> _FakeResponse: + return self.responses.pop(0) + + def close(self) -> None: + self.closed = True + self.sock = None + + +def _response() -> _FakeResponse: + return _FakeResponse( + json.dumps( + { + "data": [{"index": 0, "embedding": [1.0]}], + "usage": {"total_tokens": 1}, + } + ).encode("utf-8") + ) + + +def _client_with_connections( + connections: list[_FakeConnection], +) -> tuple[OpenAICompatibleClient, list[float]]: + client = OpenAICompatibleClient( + ProviderConfig( + name="lodedb", + base_url="http://127.0.0.1:8099/v1", + model="minilm", + requested_dimensions=None, + expected_dimensions=1, + api_key="test-key", + token_usage_label="test", + ) + ) + connection_timeouts: list[float] = [] + + def new_connection(timeout_seconds: float) -> _FakeConnection: + connection_timeouts.append(timeout_seconds) + connection = connections.pop(0) + connection.timeout = timeout_seconds + return connection + + client._new_connection = new_connection # type: ignore[method-assign] + return client, connection_timeouts + + +def test_client_reuses_connection_updates_timeout_and_sanitizes_document_input() -> None: + document_response = _response() + query_response = _response() + connection = _FakeConnection([document_response, query_response]) + client, connection_timeouts = _client_with_connections([connection]) + lone_surrogate_chunk = chunk_document("a" * 7_999 + "\U0001f600")[0] + + document_result = client.embed_batch([lone_surrogate_chunk], "document") + query_result = client.embed_batch(["cluster networking"], "query") + + assert document_result.retry_count == 0 + assert document_result.reconnects == 0 + assert query_result.retry_count == 0 + assert query_result.reconnects == 0 + assert connection_timeouts == [DOCUMENT_TIMEOUT_SECONDS] + assert connection.sock is not None + assert connection.sock.timeouts == [QUERY_TIMEOUT_SECONDS] + assert len(connection.requests) == 2 + assert document_response.was_read + assert query_response.was_read + document_body = json.loads(connection.requests[0][2].decode("utf-8")) + assert document_body["input"] == ["a" * 7_999 + "\ufffd"] + + +def test_remote_disconnect_reconnects_without_consuming_retry_budget() -> None: + failed_connection = _FakeConnection([], [http.client.RemoteDisconnected("closed")]) + reconnected_response = _response() + replacement_connection = _FakeConnection([reconnected_response]) + client, connection_timeouts = _client_with_connections( + [failed_connection, replacement_connection] + ) + + result = client.embed_batch(["pod scheduling"], "document") + + assert result.retry_count == 0 + assert result.reconnects == 1 + assert client.reconnects == 1 + assert connection_timeouts == [DOCUMENT_TIMEOUT_SECONDS, DOCUMENT_TIMEOUT_SECONDS] + assert failed_connection.closed + assert len(failed_connection.requests) == 1 + assert len(replacement_connection.requests) == 1 + assert reconnected_response.was_read + + +def test_closed_keep_alive_socket_reconnects_without_consuming_retry_budget() -> None: + initial_connection = _FakeConnection([_response()]) + replacement_connection = _FakeConnection([_response()]) + client, connection_timeouts = _client_with_connections( + [initial_connection, replacement_connection] + ) + + client.embed_batch(["service discovery"], "document") + initial_connection.sock = None + result = client.embed_batch(["persistent storage"], "document") + + assert result.retry_count == 0 + assert result.reconnects == 1 + assert client.reconnects == 1 + assert connection_timeouts == [DOCUMENT_TIMEOUT_SECONDS, DOCUMENT_TIMEOUT_SECONDS] + assert initial_connection.closed + assert len(replacement_connection.requests) == 1 diff --git a/benchmarks/openknowledge_embeddings/test_okchunk.py b/benchmarks/openknowledge_embeddings/test_okchunk.py new file mode 100644 index 00000000..82927a27 --- /dev/null +++ b/benchmarks/openknowledge_embeddings/test_okchunk.py @@ -0,0 +1,78 @@ +"""Concrete expectations ported from OpenKnowledge's chunking.test.ts.""" + +import json +import sys +from pathlib import Path + +import pytest + +_BENCHMARK_DIR = Path(__file__).resolve().parent +if str(_BENCHMARK_DIR) not in sys.path: + sys.path.insert(0, str(_BENCHMARK_DIR)) + +from bench_core import _sanitize_lone_surrogates # noqa: E402 +from okchunk import ( # noqa: E402 + CHUNK_OVERLAP_CHARS, + CHUNK_TARGET_CHARS, + MAX_CHUNKS_PER_DOC, + chunk_document, +) + + +def test_blank_input_yields_no_chunks() -> None: + assert chunk_document("") == [] + assert chunk_document(" \n\t ") == [] + + +def test_short_document_is_one_trimmed_chunk() -> None: + assert chunk_document(" hello world ") == ["hello world"] + + +def test_document_at_target_boundary_stays_one_chunk() -> None: + text = "a" * CHUNK_TARGET_CHARS + assert chunk_document(text) == [text] + + +def test_long_document_splits_with_overlap_and_preserves_content() -> None: + words = [f"word{index}" for index in range(1_200)] + chunks = chunk_document(" ".join(words)) + + assert len(chunks) > 1 + joined = " ".join(chunks) + for word in (words[0], words[600], words[-1]): + assert word in joined + + first_tail = chunks[0][-CHUNK_OVERLAP_CHARS // 2 :] + assert first_tail.strip().split(" ")[0] in chunks[1] + + +def test_chunks_never_exceed_target_length() -> None: + text = "lorem ipsum dolor sit amet " * 2_000 + assert all(len(chunk) <= CHUNK_TARGET_CHARS for chunk in chunk_document(text)) + + +def test_max_chunk_cap_applies_to_pathologically_large_input() -> None: + text = "x" * CHUNK_TARGET_CHARS * (MAX_CHUNKS_PER_DOC + 50) + assert len(chunk_document(text)) <= MAX_CHUNKS_PER_DOC + + +def test_unbroken_run_makes_forward_progress() -> None: + text = "a" * CHUNK_TARGET_CHARS * 3 + chunks = chunk_document(text) + assert len(chunks) > 1 + assert len(chunks) <= MAX_CHUNKS_PER_DOC + + +def test_astral_hard_cut_preserves_parity_and_sanitizes_request_text() -> None: + chunks = chunk_document("a" * 7_999 + "\U0001f600") + + assert 0xD800 <= ord(chunks[0][-1]) <= 0xDBFF + with pytest.raises(UnicodeEncodeError): + json.dumps({"input": [chunks[0]]}, ensure_ascii=False).encode("utf-8") + + sanitized = _sanitize_lone_surrogates(chunks[0]) + serialized = json.dumps({"input": [sanitized]}, ensure_ascii=False).encode("utf-8") + assert b"\xef\xbf\xbd" in serialized + + well_formed = "Kubernetes \U0001f600" + assert _sanitize_lone_surrogates(well_formed) == well_formed