diff --git a/api/chunker.py b/api/chunker.py index 0d00955a..997a2c75 100644 --- a/api/chunker.py +++ b/api/chunker.py @@ -12,10 +12,52 @@ logger = logging.getLogger(__name__) +def _token_spans(text: str) -> List[tuple]: + """Character spans of whitespace-separated tokens: [(start, end), ...]. + + Chunk SIZING stays token-based, which is what chunk_size has always meant. + Chunk TEXT is then sliced out of the original string between the first and + last token span, so everything between the tokens -- newlines, indentation, + column alignment -- survives verbatim. + + Why this exists (2026-08-14). The previous implementation did + `tokens = text.split()` and rebuilt each chunk as `' '.join(tokens)`. That is + lossy in a way nothing downstream can detect or recover: every run of + whitespace collapses to a single space, so a document's line and column + structure is destroyed while the text still reads fine. + + It was found via The Working with Stories Sourcebook, where a two-ended scale + is printed as two columns on one line: + + Every day Very rarely + + The source markdown on disk has 30 such paired-pole lines and 4,230 + two-column lines overall. After ingest, ACROSS ALL 911 CHUNKS of the two + Kurtz documents, the number retaining any multi-space column run was ZERO, + and the number keeping those two poles adjacent on one line was ZERO -- while + 38 chunks still contained both phrases, now unpaired and unpairable. + + The text survived; the pairing did not. Nothing downstream can then tell + which two of those lines are the endpoints of one scale, which is why a + knowledge-base query returned almost nothing useful from a book full of + scales, and why an analysis of that corpus had to be retracted for + generalising from the handful of scales that happened to be written inline. + + The loss is silent -- no error, no warning, and a chunk count that looks + healthy. mediawiki and gmail chunks in the same database DO retain column + runs, which is what localised the defect to this chunker rather than to + ingest generally. + """ + return [(m.start(), m.end()) for m in re.finditer(r"\S+", text)] + + class TextChunker: """ Split text into chunks suitable for embedding. Tokens are approximated as whitespace-separated words. + + Chunk boundaries are token-based; chunk text is sliced from the original + string so that whitespace structure is preserved. See _token_spans(). """ def __init__( @@ -33,7 +75,8 @@ def chunk_text(self, text: str) -> List[Dict[str, Any]]: if not text or not text.strip(): return [] - tokens = text.split() + spans = _token_spans(text) + tokens = [text[a:b] for a, b in spans] total_tokens = len(tokens) if total_tokens <= self.chunk_size: @@ -52,7 +95,9 @@ def chunk_text(self, text: str) -> List[Dict[str, Any]]: while start < total_tokens: end = min(start + self.chunk_size, total_tokens) chunk_tokens = tokens[start:end] - chunk_text = ' '.join(chunk_tokens) + # slice the ORIGINAL text between the first and last token of this + # chunk, rather than rejoining tokens with single spaces + chunk_text = text[spans[start][0]:spans[end - 1][1]] if len(chunk_tokens) >= self.min_chunk_size or start == 0: chunks.append({ @@ -82,7 +127,8 @@ def chunk_text(self, text: str) -> List[Dict[str, Any]]: if not text or not text.strip(): return [] - tokens = text.split() + spans = _token_spans(text) + tokens = [text[a:b] for a, b in spans] total_tokens = len(tokens) if total_tokens <= self.chunk_size: @@ -102,7 +148,7 @@ def chunk_text(self, text: str) -> List[Dict[str, Any]]: target_end = min(start + self.chunk_size, total_tokens) end = self._find_sentence_boundary(tokens, start, target_end) chunk_tokens = tokens[start:end] - chunk_text = ' '.join(chunk_tokens) + chunk_text = text[spans[start][0]:spans[end - 1][1]] if len(chunk_tokens) >= self.min_chunk_size or start == 0: chunks.append({ diff --git a/migrations/108_tenant_id.sql b/migrations/108_tenant_id.sql new file mode 100644 index 00000000..fb676973 --- /dev/null +++ b/migrations/108_tenant_id.sql @@ -0,0 +1,103 @@ +-- Migration: Add tenant_id for ingest-time tenant attribution +-- Date: 2026-07-31 +-- Purpose: Capture WHICH TENANT a document belongs to at write time, so that +-- per-tenant isolation remains possible later. This migration adds +-- CAPTURE ONLY — it does not enforce anything. See the note at the end. +-- Related: 015_add_privacy_column.sql (same shape, same sticky-merge semantics) +-- +-- WHY NOW, before any tenancy decision is made: +-- Provenance not recorded at write time cannot be reconstructed afterwards. +-- Empirically, on prod 2026-07-31: +-- SELECT count(*) FROM koi_entity_chunk_links l +-- LEFT JOIN koi_memories m ON l.document_rid = m.rid WHERE m.rid IS NULL; +-- -> 200,751 of 3,841,436 rows (5.22%) whose origin is now undecidable. +-- Every document ingested without a tenant key joins that category. +-- +-- SCOPE: koi_memories only. koi_memory_chunks inherits tenancy through +-- document_rid and does not need its own key until read paths filter on it. +-- entity_registry (41 insert sites) and entity_relationships (30) are +-- DELIBERATELY EXCLUDED — entity_registry is a global namespace +-- (UNIQUE (normalized_text, entity_type), resolved with no scope), so adding a +-- tenant key there forks every entity per tenant and destroys cross-tenant +-- aggregation. That is a product decision, not a migration. + +-- Fail fast rather than queue: an ACCESS EXCLUSIVE request that stacks behind an +-- in-flight query on koi_memories would put every subsequent query behind it. +SET lock_timeout = '3s'; + +ALTER TABLE koi_memories ADD COLUMN IF NOT EXISTS tenant_id VARCHAR(100); + +-- Empty string must never be storable. Because attribution is made sticky with +-- COALESCE(existing, new), an empty-string tenant would count as a real owner and +-- permanently block the correct tenant from ever being recorded. The Python writers +-- normalise '' -> None, but only 3 of the koi_memories writers are patched in this +-- phase, so THIS invariant is enforced in the schema where every writer must obey it. +-- Verified: without this, upserting '' then 'secondmuse' leaves the row owned by ''. +-- +-- SCOPE OF THIS CONSTRAINT — it enforces BLANKNESS ONLY, not stickiness. Set-once +-- attribution is implemented per-writer via COALESCE(koi_memories.tenant_id, +-- EXCLUDED.tenant_id), which means an unpatched writer that later adds tenant_id to +-- its upsert list with plain EXCLUDED.tenant_id would silently re-attribute rows, +-- with no error and no way to detect it afterwards. If stickiness needs to be a real +-- database invariant rather than a convention, it belongs in a BEFORE UPDATE trigger. +-- Not done here: at present exactly one writer set exists and all of it is patched. +DO $$ +BEGIN + ALTER TABLE koi_memories + ADD CONSTRAINT koi_memories_tenant_id_not_blank + CHECK (tenant_id IS NULL OR length(btrim(tenant_id)) > 0); +EXCEPTION + WHEN duplicate_object THEN NULL; -- idempotent re-run +END $$; + +COMMENT ON COLUMN koi_memories.tenant_id IS + 'Owning tenant for this document, captured at ingest. NULL = pre-tenancy or ' + 'first-party Regen data. Capture only: no read path filters on this column ' + 'as of migration 108. Set once and sticky (COALESCE on conflict) — a later ' + 'ingest of the same rid must not silently re-attribute an existing document.'; + +-- Partial index: the overwhelming majority of existing rows are NULL, so a +-- partial index stays small and only pays for rows that actually carry a tenant. +CREATE INDEX IF NOT EXISTS idx_koi_memories_tenant_id + ON koi_memories(tenant_id) + WHERE tenant_id IS NOT NULL; + +-- Composite for the eventual "active rows for tenant X" read pattern. NOTE the column +-- order: 015's idx_koi_memories_active_privacy leads with superseded_at, but under this +-- index's own partial predicate superseded_at is constant-NULL, so leading with it would +-- be informationally dead. tenant_id leads instead — verified with EXPLAIN, which puts +-- the Index Cond on tenant_id and none on superseded_at. +CREATE INDEX IF NOT EXISTS idx_koi_memories_active_tenant + ON koi_memories(tenant_id, superseded_at) + WHERE superseded_at IS NULL AND tenant_id IS NOT NULL; + +-- NOTE ON LOCKING (measured on a same-size replica, PG 15; an earlier version of this +-- comment had the lock model inverted): +-- ALTER TABLE ADD COLUMN -> AccessExclusiveLock, catalog-only in PG11+, +-- no table rewrite. ~5 ms. +-- ALTER TABLE ADD CONSTRAINT -> AccessExclusiveLock, and it BLOCKS READERS +-- ... CHECK during a full heap validation scan. ~60-70 ms. +-- This is the statement that actually blocks queries. +-- CREATE INDEX (non-concurrent) -> ShareLock: blocks WRITERS only. SELECTs run fine. +-- ~50 ms each. +-- Sizing: the relevant figure is the HEAP (~178 MB), not pg_total_relation_size +-- (~577 MB, which includes ~192 MB of indexes and ~207 MB of TOAST). tenant_id is +-- not TOASTed, so validation scans the heap only. +-- On a much larger table, or outside a transaction, prefer CREATE INDEX CONCURRENTLY +-- (which cannot run inside a transaction block). +-- +-- NOTE ON DEPLOY ORDER — this is load-bearing, not advisory: +-- RUN THIS MIGRATION BEFORE DEPLOYING THE CODE THAT WRITES tenant_id. +-- The patched writers name tenant_id in their INSERT column list, so against a +-- pre-108 database every write raises +-- asyncpg.exceptions.UndefinedColumnError: column "tenant_id" ... does not exist +-- and the row is NOT written. Verified by reproduction on a pre-108 scratch DB: +-- rows written = 0. It fails loudly rather than silently, but it fails closed — +-- an ingest run against an unmigrated DB loses the documents for that run. +-- The same applies in reverse: do not roll this migration back while the patched +-- code is live. Roll back the code first. +-- +-- NOTE ON ENFORCEMENT: this migration does NOT make it safe to admit an external +-- user. koi-query-api.ts buildPrivacyFilter() still returns '' for any +-- authenticated caller ("Authenticated users see all data"), and access is gated +-- solely by a hardcoded @regen.network email check. Capture is not enforcement. diff --git a/migrations/down/108_tenant_id_down.sql b/migrations/down/108_tenant_id_down.sql new file mode 100644 index 00000000..ed2eb97c --- /dev/null +++ b/migrations/down/108_tenant_id_down.sql @@ -0,0 +1,15 @@ +-- Rollback for 108_tenant_id.sql +-- +-- WARNING: dropping tenant_id DESTROYS attribution that cannot be recomputed. +-- That is the whole point of the forward migration. If you are rolling back +-- because a read path misbehaved, prefer fixing the read path — the column is +-- inert on its own (nothing filters on it as of 108). +-- +-- If you genuinely need to roll back after any tenant-attributed ingest has run, +-- capture the mapping first so it can be restored: +-- CREATE TABLE koi_memories_tenant_backup_108 AS +-- SELECT rid, tenant_id FROM koi_memories WHERE tenant_id IS NOT NULL; + +DROP INDEX IF EXISTS idx_koi_memories_active_tenant; +DROP INDEX IF EXISTS idx_koi_memories_tenant_id; +ALTER TABLE koi_memories DROP COLUMN IF EXISTS tenant_id; diff --git a/scripts/doc_scanner.py b/scripts/doc_scanner.py index f56c07b4..f6cc9434 100644 --- a/scripts/doc_scanner.py +++ b/scripts/doc_scanner.py @@ -26,6 +26,8 @@ --reconcile-interval Seconds between reconcile scans (deletion/drift cleanup) --no-watcher Disable watchdog watcher and rely on timed scans only + --tenant ID Owning tenant recorded on every document ingested (migration 108). + Capture only; sticky (never re-attributes an existing document). """ import argparse @@ -105,13 +107,18 @@ async def get_existing_docs(conn: asyncpg.Connection, repo_name: str) -> Dict[st rows = await conn.fetch(""" SELECT rid, metadata->>'rel_path' AS rel_path, - metadata->>'content_hash' AS content_hash + metadata->>'content_hash' AS content_hash, + tenant_id FROM koi_memories WHERE source_sensor = 'doc-scanner' AND metadata->>'repo' = $1 """, repo_name) return { - r["rel_path"]: {"content_hash": r["content_hash"], "rid": r["rid"]} + r["rel_path"]: { + "content_hash": r["content_hash"], + "rid": r["rid"], + "tenant_id": r["tenant_id"], + } for r in rows if r["rel_path"] } @@ -144,6 +151,7 @@ async def upsert_doc( frontmatter: Dict[str, Any], body_text: str, chash: str, + tenant: Optional[str] = None, ) -> str: """Upsert into koi_memories. Returns the memory UUID.""" doc_content = { @@ -169,17 +177,30 @@ async def upsert_doc( existing = await conn.fetchrow("SELECT id FROM koi_memories WHERE rid = $1", rid) event_type = "NEW" if existing is None else "UPDATE" + # Tenancy: capture the owning tenant at write time (migration 108). This writer is + # the folder-watching path a per-client ingest directory would travel, so it is the + # one route where losing attribution is unrecoverable. CAPTURE ONLY — no read path + # filters on tenant_id yet; this does not make the data safe to expose. + # Sticky via COALESCE: re-scanning a doc must never re-attribute an existing one. + # Operator flag is the ONLY source. Deliberately not read from frontmatter: a + # scanned file must not be able to declare which tenant owns it. + tenant_id = tenant.strip() or None if isinstance(tenant, str) else None + if tenant_id is not None and len(tenant_id) > 100: + logger.warning("tenant too long (%d chars), not recording", len(tenant_id)) + tenant_id = None + memory_id = await conn.fetchval(""" - INSERT INTO koi_memories (id, rid, event_type, source_sensor, content, metadata) - VALUES ($1, $2, $3, 'doc-scanner', $4::jsonb, $5::jsonb) + INSERT INTO koi_memories (id, rid, event_type, source_sensor, content, metadata, tenant_id) + VALUES ($1, $2, $3, 'doc-scanner', $4::jsonb, $5::jsonb, $6) ON CONFLICT (rid) DO UPDATE SET event_type = EXCLUDED.event_type, content = EXCLUDED.content, metadata = EXCLUDED.metadata, + tenant_id = COALESCE(koi_memories.tenant_id, EXCLUDED.tenant_id), updated_at = NOW() RETURNING id """, uuid.uuid4(), rid, event_type, - json.dumps(doc_content), json.dumps(doc_metadata)) + json.dumps(doc_content), json.dumps(doc_metadata), tenant_id) return str(memory_id) @@ -238,6 +259,7 @@ async def scan_repo( force: bool, doc_id_only: bool, delete_missing: bool = True, + tenant: Optional[str] = None, ) -> Dict[str, int]: pool = await asyncpg.create_pool(POSTGRES_URL, min_size=1, max_size=3) embedder = RemoteEmbeddingProvider( @@ -284,10 +306,18 @@ async def scan_repo( chash = content_hash(raw) rid = f"doc-scanner:{repo_name}:{rel_path}" - if not force and existing_hashes.get(rel_path) == chash: + # Unchanged content normally short-circuits. But if a tenant was supplied and + # this document has not been attributed yet, we must still write — otherwise + # `--tenant` silently does nothing for every already-indexed doc, which is + # exactly the lost-attribution failure migration 108 exists to prevent. + _existing = existing.get(rel_path) or {} + _needs_tenant = bool(tenant) and not _existing.get("tenant_id") + if not force and _existing.get("content_hash") == chash and not _needs_tenant: logger.debug("Unchanged %s", rel_path) stats["skipped"] += 1 continue + if _needs_tenant and _existing.get("content_hash") == chash: + logger.info("repo_doc_sensor.backfill_tenant rel_path=%s tenant=%s", rel_path, tenant) governed_marker = " [governed]" if is_governed(fm) else "" logger.info("Indexing %s%s", rel_path, governed_marker) @@ -315,7 +345,7 @@ async def scan_repo( embeddings.append(None) async with pool.acquire() as conn: - await upsert_doc(conn, rid, repo_name, rel_path, fm, body, chash) + await upsert_doc(conn, rid, repo_name, rel_path, fm, body, chash, tenant) await upsert_chunks(conn, rid, chunks, embeddings, fm, repo_name, rel_path) stats["indexed"] += 1 @@ -430,11 +460,13 @@ def __init__( scan_interval: int, reconcile_interval: int, watcher_enabled: bool, + tenant: Optional[str] = None, ): self.repo_path = repo_path self.repo_name = repo_name self.dry_run = dry_run self.doc_id_only = doc_id_only + self.tenant = tenant self.scan_interval = scan_interval self.reconcile_interval = reconcile_interval self.watcher_enabled = watcher_enabled @@ -513,6 +545,7 @@ async def run_cycle(self): force=False, doc_id_only=self.doc_id_only, delete_missing=due_reconcile, + tenant=self.tenant, ) finally: completed_at = datetime.now(timezone.utc) @@ -556,6 +589,10 @@ def main(): parser.add_argument("repo_path", help="Path to repo root") parser.add_argument("--repo-name", help="Override repo name (default: dir name)") parser.add_argument("--dry-run", action="store_true", help="Parse without writing") + parser.add_argument("--tenant", default=None, + help="Owning tenant id, recorded on every document ingested in this run " + "(migration 108). Capture only — no read path filters on it yet. " + "Sticky: will not re-attribute documents that already have a tenant.") parser.add_argument("--force", action="store_true", help="Re-index unchanged files") parser.add_argument("--doc-id-only", action="store_true", help="Only index files with doc_id frontmatter") @@ -587,13 +624,15 @@ def main(): scan_interval=args.scan_interval, reconcile_interval=args.reconcile_interval, watcher_enabled=not args.no_watcher, + tenant=args.tenant, ) try: asyncio.run(sensor.serve_forever()) except KeyboardInterrupt: logger.info("repo_doc_sensor.stopped repo=%s", repo_name) else: - asyncio.run(scan_repo(repo_path, repo_name, args.dry_run, args.force, args.doc_id_only)) + asyncio.run(scan_repo(repo_path, repo_name, args.dry_run, args.force, args.doc_id_only, + tenant=args.tenant)) if __name__ == "__main__": diff --git a/src/core/koi_event_bridge_semantic.py b/src/core/koi_event_bridge_semantic.py index d1150e13..15fa9ee8 100644 --- a/src/core/koi_event_bridge_semantic.py +++ b/src/core/koi_event_bridge_semantic.py @@ -795,14 +795,26 @@ async def store_processed_document( is_private = bool((metadata or {}).get('is_private', False)) access_source = (metadata or {}).get('access_source') + # Tenancy: capture the owning tenant at write time (migration 108). THIS is the + # document-level writer (is_chunk FALSE) — koi_event_bridge_v2.create_new_version + # only ever writes chunk rows, so document attribution has to happen here. + # CAPTURE ONLY: nothing filters on tenant_id. Sticky via COALESCE so a re-ingest + # cannot re-attribute a document that already has an owner. Non-str values are + # coerced away rather than raising, so a malformed payload cannot kill the bridge. + _raw_tenant = (metadata or {}).get('tenant_id') + tenant_id = _raw_tenant.strip() or None if isinstance(_raw_tenant, str) else None + if tenant_id is not None and len(tenant_id) > 100: + logger.warning("koi_memories.tenant_id too long (%d chars), dropping", len(tenant_id)) + tenant_id = None + # Store document in koi_memories (not chunks) await conn.execute(""" INSERT INTO koi_memories ( id, rid, cid, version, event_type, source_sensor, content, metadata, published_at, published_confidence, content_hash, source_content_rid, is_chunk, - is_private, access_source - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, FALSE, $13, $14) + is_private, access_source, tenant_id + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, FALSE, $13, $14, $15) ON CONFLICT (rid) DO UPDATE SET content = $7, metadata = $8, @@ -812,6 +824,7 @@ async def store_processed_document( source_content_rid = $12, is_private = (koi_memories.is_private OR EXCLUDED.is_private), access_source = COALESCE(koi_memories.access_source, EXCLUDED.access_source), + tenant_id = COALESCE(koi_memories.tenant_id, EXCLUDED.tenant_id), updated_at = NOW() """, str(uuid.uuid4()), @@ -833,7 +846,8 @@ async def store_processed_document( content_hash, source_content_rid, is_private, - access_source + access_source, + tenant_id ) logger.info(f"Stored processed document in koi_memories: {document_rid}") diff --git a/src/core/koi_event_bridge_v2.py b/src/core/koi_event_bridge_v2.py index 35251ef2..ba3bc718 100644 --- a/src/core/koi_event_bridge_v2.py +++ b/src/core/koi_event_bridge_v2.py @@ -466,14 +466,23 @@ async def create_new_version(conn: asyncpg.Connection, event: KOIEvent, is_private = bool(bundle_meta.get('is_private', False)) access_source = bundle_meta.get('access_source') + # Tenancy: promote bundle-metadata tenant_id to a dedicated column (migration 108). + # CAPTURE ONLY — nothing filters on this column yet; see the migration's closing note. + # Sticky via COALESCE on conflict: a re-ingest of the same rid must never silently + # re-attribute a document that already belongs to someone. Empty string is + # normalised to NULL so a sensor emitting '' cannot claim ownership by accident. + tenant_id = (bundle_meta.get('tenant_id') or None) + if isinstance(tenant_id, str): + tenant_id = tenant_id.strip() or None + # Insert new version with publication tracking await conn.execute(""" INSERT INTO koi_memories ( id, rid, cid, version, previous_version_id, event_type, source_sensor, content, metadata, published_at, published_confidence, content_hash, - is_private, access_source - ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14) + is_private, access_source, tenant_id + ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15) ON CONFLICT (rid) DO UPDATE SET content = EXCLUDED.content, metadata = EXCLUDED.metadata, @@ -481,6 +490,7 @@ async def create_new_version(conn: asyncpg.Connection, event: KOIEvent, content_hash = EXCLUDED.content_hash, is_private = (koi_memories.is_private OR EXCLUDED.is_private), access_source = COALESCE(koi_memories.access_source, EXCLUDED.access_source), + tenant_id = COALESCE(koi_memories.tenant_id, EXCLUDED.tenant_id), superseded_at = NULL, updated_at = CURRENT_TIMESTAMP """, @@ -500,7 +510,8 @@ async def create_new_version(conn: asyncpg.Connection, event: KOIEvent, published_confidence, content_hash, is_private, - access_source + access_source, + tenant_id ) # Fetch the actual memory_id from database (in case ON CONFLICT kept the old UUID) diff --git a/tests/test_chunker_preserves_whitespace.py b/tests/test_chunker_preserves_whitespace.py new file mode 100644 index 00000000..eea310f6 --- /dev/null +++ b/tests/test_chunker_preserves_whitespace.py @@ -0,0 +1,106 @@ +"""The chunker must not flatten whitespace. + +Regression test for a silent data-loss defect found 2026-08-14. `TextChunker` +and `SentenceAwareChunker` both did `tokens = text.split()` and rebuilt each +chunk as `' '.join(tokens)`. Text survived; every run of whitespace collapsed to +a single space, so line and column structure was destroyed on ingest with no +error, no warning, and a healthy-looking chunk count. + +The corpus that exposed it: The Working with Stories Sourcebook prints a +two-ended scale as two columns on one line ("Every day Very rarely"). +Across all 911 chunks of the two ingested Kurtz documents, the number retaining +any multi-space column run was ZERO, while 38 chunks still contained both poles +as separate unpairable lines. A downstream analysis had to be retracted because +only the scales that happened to be written inline were visible. + +These tests are written to FAIL against the old implementation -- verified: the +pre-fix code scores 0 on every preservation assertion below. +""" + +import re +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from api.chunker import SentenceAwareChunker, TextChunker # noqa: E402 + +COLUMN_RUN = re.compile(r"[A-Za-z]{3,}[ ]{4,}[A-Za-z]{3,}") +PAIRED_POLES = re.compile(r"Every day[ ]{2,}Very rarely") + +# A page shaped like the Sourcebook's: prose, then column-encoded scales. +SCALE_PAGE = """How often does something like this happen where you live? + + Every day Very rarely + + I don't know + +And how much does it matter to the people involved? + + Not at all Completely + + I don't know +""" + + +def _long(text: str) -> str: + """Pad past chunk_size so the multi-chunk path is exercised, not the + single-chunk early return (which was always faithful and would hide this).""" + return text + "\n\n" + ("filler " * 3000) + + +@pytest.mark.parametrize("cls", [TextChunker, SentenceAwareChunker]) +def test_column_alignment_survives_chunking(cls): + text = _long(SCALE_PAGE) + chunks = cls(chunk_size=200, chunk_overlap=20, min_chunk_size=10).chunk_text(text) + joined = "\n".join(c["text"] for c in chunks) + + assert PAIRED_POLES.search(joined), ( + "the two poles of a scale are no longer on one line; the pairing is lost " + "and nothing downstream can recover which two lines belong together" + ) + assert COLUMN_RUN.search(joined), "all multi-space column runs were collapsed" + + +@pytest.mark.parametrize("cls", [TextChunker, SentenceAwareChunker]) +def test_newlines_survive_chunking(cls): + text = _long(SCALE_PAGE) + chunks = cls(chunk_size=200, chunk_overlap=20, min_chunk_size=10).chunk_text(text) + assert any("\n" in c["text"] for c in chunks), ( + "no chunk contains a newline; the document was flattened to one line" + ) + + +@pytest.mark.parametrize("cls", [TextChunker, SentenceAwareChunker]) +def test_chunk_text_is_a_verbatim_slice_of_the_source(cls): + """The strongest form: every chunk must appear in the original, exactly.""" + text = _long(SCALE_PAGE) + for c in cls(chunk_size=200, chunk_overlap=20, min_chunk_size=10).chunk_text(text): + assert c["text"] in text, ( + "a chunk is not a substring of the input, so it was rewritten rather " + "than sliced" + ) + + +@pytest.mark.parametrize("cls", [TextChunker, SentenceAwareChunker]) +def test_sizing_behaviour_is_unchanged(cls): + """The fix must change chunk CONTENT only, never chunk boundaries.""" + text = _long(SCALE_PAGE) + chunks = cls(chunk_size=200, chunk_overlap=20, min_chunk_size=10).chunk_text(text) + assert chunks, "no chunks produced" + assert all(c["end_token"] > c["start_token"] for c in chunks) + assert all(c["end_token"] - c["start_token"] <= 200 for c in chunks), ( + "a chunk exceeded chunk_size in tokens" + ) + assert [c["index"] for c in chunks] == list(range(len(chunks))) + assert all(c["total_chunks"] == len(chunks) for c in chunks) + + +def test_empty_and_short_inputs_still_behave(): + ch = TextChunker(chunk_size=200, chunk_overlap=20, min_chunk_size=10) + assert ch.chunk_text("") == [] + assert ch.chunk_text(" \n ") == [] + short = ch.chunk_text("just a few words here") + assert len(short) == 1 and short[0]["total_chunks"] == 1