From 09bb59a1dbd9a2928e24ad26d8f09c485b9d54ab Mon Sep 17 00:00:00 2001 From: oskar Date: Wed, 15 Jul 2026 20:41:00 +0200 Subject: [PATCH] feat(documents-ingest): chunk + embed job (module 5 phase 2 step 6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds documents-ingest-embed: chunks source='paperless' envelope content (paragraph-preferring, ~600 tok/chunk, ~150 tok overlap, hard char-fallback for oversized paragraphs per plan §2 decision 3), embeds each chunk via Ollama (bge-m3, dim validated against document_chunk's VECTOR(1024) on every response) and inserts into document_chunk. Lives in documents-ingest per the plan's own recommendation (§6 step 6) rather than a new package — reuses the job family's existing idempotency/stats-balance/dry-run conventions (paperless_adapter.py, gmail-header-backfill). A 5-angle multi-agent code review of the initial implementation surfaced three real bugs, fixed here: hard_split() could infinite-loop if --chunk-overlap >= --chunk-size (now guarded in both hard_split() and main()); insert_chunk() wasn't error-isolated like embed_chunk(), so a DB write failure would crash the whole run instead of being counted and skipped; and ON CONFLICT DO NOTHING's outcome was discarded, so a silently skipped row (the known gap where document_chunk's UNIQUE constraint doesn't include `model`) would have been miscounted as a successful insert - now tracked separately as chunks_conflict_skipped and treated as a run failure. Smoke-tested and run to completion live on SOLARIA against the real Ollama instance and kb-postgres@PIHA: dry-run matched the known phase-2-step-5 figures exactly (186 fetched, 26 empty_content, 2684 chunks planned), a --limit 10 apply + idempotent re-run + DB/distance sanity checks all passed, and the full 186-document run inserted 2683/2684 chunks (1 isolated error - Ollama's runtime context window rejected one pathological dot-leader table-of-contents chunk that tokenized far more densely than estimated; documented as a known limitation, not fixed here given it's a single-chunk edge case). Timing: ~0.79s/chunk average on CPU (SOLARIA's Ollama runs GPU-less per the recent GPU-reservation-disabled fix), ~35 min wall-clock for the full pilot - the real input for scoping the later mail-corpus embedding phase (plan §7's GPU-based estimate doesn't hold here). pytest: 101 passed. Co-Authored-By: Claude Sonnet 5 --- jobs/documents-ingest/README.md | 174 +++++++ jobs/documents-ingest/pyproject.toml | 1 + .../src/documents_ingest/chunk_embed.py | 426 ++++++++++++++++ .../tests/test_chunk_embed.py | 455 ++++++++++++++++++ 4 files changed, 1056 insertions(+) create mode 100644 jobs/documents-ingest/src/documents_ingest/chunk_embed.py create mode 100644 jobs/documents-ingest/tests/test_chunk_embed.py diff --git a/jobs/documents-ingest/README.md b/jobs/documents-ingest/README.md index a264f33..c3aad43 100644 --- a/jobs/documents-ingest/README.md +++ b/jobs/documents-ingest/README.md @@ -301,3 +301,177 @@ the live Paperless API, over SSH — **not executed as part of this change** without operator confirmation (this job reads production Paperless data and writes production envelope rows on `--apply`). `pytest` passes locally before this commit. + +--- + +## Phase 2 step 6 — `documents-ingest-embed` (chunk + embed) + +Module 5, phase 2, plan step 6 (`docs/kb/modules/05-faza2-plan.md`, §6 step 6, +§2 decision 3). Reads `entities[type=content].text` off every `source='paperless'` +envelope, chunks it, calls Ollama (`POST /api/embeddings`, model `bge-m3`) for +each chunk, and inserts the result into `document_chunk` +(`services/kb-postgres/init/002_chunks.sql`). This job only ever `INSERT`s into +`document_chunk` — `envelope` is read-only here, and `services/ollama/` is +untouched. + +### Where it runs + +**On SOLARIA** (that's where Ollama lives), against kb-postgres@PIHA over +Tailscale — the reverse of the other jobs in this package, which run on PIHA. +`--ollama-url` defaults to `http://localhost:11434` (Ollama on the same node); +`--dsn` needs PIHA's Tailscale address, e.g. +`postgresql://kb:@piha:5433/kb`. + +### Install + +```bash +pip install -e packages/kb-mail/ +pip install -e jobs/documents-ingest/ +``` + +### Usage + +```bash +# Dry run (default) — chunk and count only, no Ollama calls, no DB writes: +documents-ingest-embed --dsn postgresql://kb:@piha:5433/kb + +# Smoke-test slice: +documents-ingest-embed --dsn ... --apply --limit 10 + +# Full run: +documents-ingest-embed --dsn ... --apply +``` + +### Chunking (plan §2 decision 3) + +Paragraph-preferring: splits on blank-line boundaries, greedily packs +paragraphs up to `--chunk-size` characters (default 2400, ≈600 tokens at a +~4 chars/token heuristic — no local bge-m3 tokenizer available offline), +`--chunk-overlap` characters of trailing context carried into the next chunk +(default 600, ≈150 tokens). A paragraph that alone exceeds `--chunk-size` +falls back to a hard character-based sliding window — Paperless OCR text has +no page-break markers (plan §1.2), so there's nothing else to split large, +unbroken text on. A document with empty OCR content (the 26 `empty_content` +documents from phase 2 step 5) yields zero chunks and is counted separately, +not as an error. + +### Idempotency + +A pre-fetched set of `(envelope_id, chunk_index)` pairs already embedded with +`--model` skips re-embedding on rerun — no wasted Ollama calls. +`document_chunk`'s own `UNIQUE (envelope_id, chunk_index)` + +`ON CONFLICT DO NOTHING` is the second line of defense; `insert_chunk`'s +command tag is checked so a silently-skipped row is counted as +`chunks_conflict_skipped`, never miscounted as `chunks_inserted`. Note that +uniqueness is on `(envelope_id, chunk_index)` only, not `model` — +re-embedding with a *different* model hits this path and that embedding is +discarded (wasted work, correctly reported via `chunks_conflict_skipped`, +but not persisted). Out of scope for this single-model pilot; the real fix +for whoever indexes a second model later is `UNIQUE (envelope_id, +chunk_index, model)` at the schema layer. + +A DB write failure for one chunk (dropped connection, unexpected bytes) is +isolated the same way an embed failure is — counted as `chunks_errors`, +never aborting the rest of the run. + +### Dimension guard + +Every embedding response's length is checked against +`document_chunk.embedding`'s `VECTOR(1024)` column. A mismatch raises +`EmbeddingDimensionError` and aborts the whole run immediately — never +silently indexes vectors of the wrong dimension. + +### Chunk size/overlap validation + +`--chunk-overlap` must be smaller than `--chunk-size` — the sliding-window +hard-split fallback advances by `chunk_size - chunk_overlap` per step, so an +overlap `>=` size would never advance and hang. `main()` rejects this +combination before opening a DB connection; `hard_split()` itself also +raises `ValueError` as a second line of defense for direct callers. + +### Stats must balance + +``` +documents_fetched = empty_content + documents_chunked +chunks_total = chunks_already_embedded + chunks_inserted + + chunks_conflict_skipped + chunks_errors +``` + +`main()` exits 1 on `chunks_errors > 0`, `chunks_conflict_skipped > 0`, or if +either balance breaks. The summary line also reports +`avg_embed_seconds_per_chunk` — CPU-only Ollama timing, the input for +deciding whether/how to scale this to the mail corpus later (plan §7). + +### Tests + +```bash +pip install -e packages/kb-mail/ +pip install -e jobs/documents-ingest/ +cd jobs/documents-ingest && pytest +``` + +Pure unit tests, no DB or real HTTP — `run()` is tested by monkeypatching +`asyncpg.connect` (fake connection) and `aiohttp.ClientSession` (fake session +serving a canned embedding vector, or a 500 for a chosen prompt to exercise +error isolation). Covers: chunking (paragraph boundaries, overlap, empty +document, document shorter than one chunk, oversized paragraph hard-fallback, +the overlap-must-be-smaller-than-size guard), `extract_content`, idempotency +(pre-existing keys skipped, no Ollama calls made for them, a second `--apply` +run embeds nothing new, existing keys are correctly scoped to `--model`), +dimension-mismatch abort, isolated per-chunk embed and insert errors, +`ON CONFLICT` no-ops counted separately from real inserts, and the +stats-balance invariant. + +### Known limitation — Ollama context-length rejections on pathological chunks + +Ollama's *runtime* context window for a model can be smaller than the +model's advertised max (bge-m3 supports 8192 tokens, but Ollama's default +`num_ctx` is lower) — and some OCR text tokenizes far more densely than the +~4-chars/token heuristic this job uses to size chunks. Concretely: a table- +of-contents page made almost entirely of dot-leader formatting +(`". . . . . . ."`, repeated hundreds of times) hit this on the pilot run — +Ollama returned `500 {"error":"the input length exceeds the context +length"}` for one 2400-char chunk that should have been well within budget +by character count alone. The job isolates this exactly like any other embed +failure (`chunks_errors`, logged, run continues), so it never crashes a +run — but it also never automatically shrinks and retries the offending +chunk. Given how rare this was (1 chunk out of 2684 in the full pilot, all +from one document's dot-leader ToC), it's left as a known gap rather than +fixed here; a real fix would be either a smaller/adaptive chunk size for +low-character-entropy text, or a shrink-and-retry loop on this specific +Ollama error. + +### Definition of Done + +Per `CLAUDE.md`: `pytest` passes locally (101 tests). Smoke-tested and then +run to completion live on SOLARIA against the real Ollama instance and +kb-postgres@PIHA: + +- Dry-run: 186 fetched, 26 `empty_content`, 2684 chunks planned — matches + the known phase-2-step-5 figures exactly. +- `--apply --limit 10`: 64 chunks embedded, 0 errors, avg ≈0.83s/chunk on CPU. +- Re-run of the same slice: fully idempotent — 0 Ollama calls, 0 inserts. +- Full `--apply` (all 186 documents): **2683/2684 chunks inserted, 1 isolated + error** (see "Known limitation" above) — `chunks_errors=1` correctly + produced a non-zero exit rather than silently reporting success. + `document_chunk` ends at 2683 rows across 160 distinct envelopes, matching + `documents_chunked`. A `document_chunk_envelope_idx`-backed count and an + `ORDER BY embedding <=> ...` nearest-neighbor sanity query both look + correct (top match is the reference chunk itself at distance 0; next + nearest are chunks of the same source document). +- **Timing (CPU-only, no GPU driver on SOLARIA)**: ≈0.79s/chunk average + across 2683 real embeddings (2115.8s total embed time), ≈13.2s/document + average across the 160 chunked documents, ≈35 minutes wall-clock for the + full 186-document pilot. This is the real-world input for scaling this + pipeline to the much larger mail corpus later (plan §7 assumed GPU-based + "minutes for the whole pilot"; SOLARIA's Ollama currently runs CPU-only + per the recent GPU-reservation-disabled fix). The 186-document pilot's + ≈13.2s/document average is dominated by Paperless' long OCR text (≈22k + chars/doc average, per plan §1.2) — 225 030 mail envelopes will have a + very different, likely much shorter, per-envelope chunk count (email + bodies vs. scanned multi-page PDFs), so this number doesn't extrapolate + directly to a mail-corpus estimate. What it does establish: at + ≈0.79s/chunk sequential CPU embedding, any corpus with a non-trivial + average chunk count per item will need either a GPU driver fix, + concurrent/batched Ollama calls, or both, before a full mail-corpus run + is practical — flagged for whoever picks up the mail-indexer phase. diff --git a/jobs/documents-ingest/pyproject.toml b/jobs/documents-ingest/pyproject.toml index 8265a09..9cead27 100644 --- a/jobs/documents-ingest/pyproject.toml +++ b/jobs/documents-ingest/pyproject.toml @@ -16,6 +16,7 @@ dependencies = [ [project.scripts] documents-ingest = "documents_ingest.extractor:main" documents-ingest-paperless = "documents_ingest.paperless_adapter:main" +documents-ingest-embed = "documents_ingest.chunk_embed:main" [tool.setuptools.packages.find] where = ["src"] diff --git a/jobs/documents-ingest/src/documents_ingest/chunk_embed.py b/jobs/documents-ingest/src/documents_ingest/chunk_embed.py new file mode 100644 index 0000000..3d3ff91 --- /dev/null +++ b/jobs/documents-ingest/src/documents_ingest/chunk_embed.py @@ -0,0 +1,426 @@ +"""Chunk + embed job — module 5 phase 2, plan step 6 (docs/kb/modules/05-faza2-plan.md, +§6 step 6, decision 3). + +Pipeline: `envelope(source='paperless').entities[type=content].text` -> chunk (paragraph- +preferring, ~600 tok/chunk, ~150 tok overlap, hard char-fallback for oversized paragraphs) +-> `POST /api/embeddings` (Ollama on SOLARIA, model `bge-m3`) -> `INSERT document_chunk` +(`services/kb-postgres/init/002_chunks.sql`, untouched by this change). + +Scope is deliberately narrow to the pilot (plan §6 step 7): only `source='paperless'` +envelopes. Mail chunking/embedding is a later phase (plan §1.1) and will reuse this same +`chunk_text()` + `document_chunk` table, not a new pipeline. + +Runs on SOLARIA (needs Ollama on localhost) against kb-postgres@PIHA over Tailscale: + +Install (from repo root): + pip install -e packages/kb-mail/ + pip install -e jobs/documents-ingest/ + +Usage: + # Dry run (default) — chunk and count, no Ollama calls, no DB writes: + documents-ingest-embed --dsn postgresql://kb:@piha:5433/kb + + # Real run: + documents-ingest-embed --dsn ... --apply + + # Smoke-test slice: + documents-ingest-embed --dsn ... --apply --limit 10 + +DSN can come from KB_DSN, Ollama URL from OLLAMA_URL (default http://localhost:11434). + +Idempotency: a pre-fetched set of existing (envelope_id, chunk_index) pairs for this +`model` skips chunks already embedded — no re-embedding, no wasted Ollama calls on rerun. +`insert_chunk`'s own `ON CONFLICT (envelope_id, chunk_index) DO NOTHING` is the second line +of defense; its command tag is checked so a silently-skipped row is counted as +`chunks_conflict_skipped`, never miscounted as `chunks_inserted`. Note: that UNIQUE +constraint is on (envelope_id, chunk_index) only, not including `model` — re-embedding the +same pilot with a *different* model would hit this path and its embedding would be +discarded (a real chunk of work wasted, correctly reported, but not persisted). Out of +scope for this pilot (single model, bge-m3); flagged as a follow-up if/when a second model +is ever indexed (the real fix is `UNIQUE (envelope_id, chunk_index, model)` at the schema +layer). +""" +from __future__ import annotations + +import argparse +import asyncio +import json +import os +import re +import sys +import time +from typing import Optional + +import aiohttp +import asyncpg +import structlog + +_log = structlog.get_logger(__name__) + +DEFAULT_OLLAMA_URL = "http://localhost:11434" +DEFAULT_MODEL = "bge-m3" +EXPECTED_DIM = 1024 + +# Plan decision 3: ~600 tok/chunk, ~150 tok overlap. No local bge-m3 tokenizer available +# offline, so tokens are approximated at ~4 chars/token (common heuristic) -> ~2400 chars +# target, ~600 chars overlap. Paragraph boundaries are preferred; a hard character-based +# sliding window is only used when a single paragraph alone exceeds the target (plan §1.2: +# Paperless OCR text has no page-break markers to split on instead). +CHARS_PER_TOKEN = 4 +TARGET_TOKENS = 600 +OVERLAP_TOKENS = 150 +TARGET_CHARS = TARGET_TOKENS * CHARS_PER_TOKEN +OVERLAP_CHARS = OVERLAP_TOKENS * CHARS_PER_TOKEN + +_PARA_SPLIT = re.compile(r"\n\s*\n") + +_INSERT_SQL = """ +INSERT INTO document_chunk (envelope_id, chunk_index, text, embedding, model) +VALUES ($1, $2, $3, $4::vector, $5) +ON CONFLICT (envelope_id, chunk_index) DO NOTHING +""" + + +class EmbeddingDimensionError(RuntimeError): + """Ollama returned a vector of the wrong dimension for the target `document_chunk` schema.""" + + +def split_paragraphs(text: str) -> list[str]: + """Split on blank-line boundaries; drop empty fragments.""" + return [p.strip() for p in _PARA_SPLIT.split(text) if p.strip()] + + +def hard_split(text: str, size: int = TARGET_CHARS, overlap: int = OVERLAP_CHARS) -> list[str]: + """Fixed-size sliding window over raw characters — fallback for a paragraph too big to fit.""" + if overlap >= size: + # step = size - overlap would be <= 0 below, so `start` would never advance past `n` + # and the while loop would spin forever, growing `chunks` without bound. + raise ValueError(f"chunk_overlap ({overlap}) must be smaller than chunk_size ({size})") + if len(text) <= size: + return [text] + step = size - overlap + chunks = [] + start = 0 + n = len(text) + while start < n: + end = min(start + size, n) + chunks.append(text[start:end]) + if end == n: + break + start += step + return chunks + + +def chunk_text(text: str, size: int = TARGET_CHARS, overlap: int = OVERLAP_CHARS) -> list[str]: + """Paragraph-preferring chunker with a hard character-fallback (plan §2 decision 3). + + Empty/whitespace-only text (e.g. the 26 empty_content Paperless documents, plan §1.2) + returns []. A document shorter than one chunk returns exactly one chunk (the full text). + """ + text = (text or "").strip() + if not text: + return [] + + paragraphs = split_paragraphs(text) + + chunks: list[str] = [] + buffer = "" + for para in paragraphs: + if len(para) > size: + if buffer: + chunks.append(buffer) + buffer = "" + chunks.extend(hard_split(para, size, overlap)) + continue + + candidate = f"{buffer}\n\n{para}" if buffer else para + if len(candidate) <= size: + buffer = candidate + continue + + chunks.append(buffer) + tail = buffer[-overlap:] if buffer else "" + buffer = f"{tail}\n\n{para}" if tail else para + + if buffer: + chunks.append(buffer) + return chunks + + +def extract_content(entities: list) -> str: + """Pull `entities[type=content].text` out of a document envelope (plan §4.2). Missing + or empty content both yield "".""" + for entity in entities or []: + if isinstance(entity, dict) and entity.get("type") == "content": + return entity.get("text") or "" + return "" + + +def _decode_jsonb(value: object) -> object: + """asyncpg may return jsonb as a str or an already-decoded object.""" + if value is None: + return None + if isinstance(value, str): + return json.loads(value) + return value + + +def _vector_literal(embedding: list[float]) -> str: + """Render a Python float list as a pgvector input literal, e.g. '[0.1,0.2,...]'.""" + return "[" + ",".join(repr(v) for v in embedding) + "]" + + +async def fetch_documents(conn: asyncpg.Connection, limit: Optional[int], offset: Optional[int]) -> list: + """`source='paperless'` envelopes, ordered by id for stable --limit/--offset slicing.""" + query = "SELECT id, entities FROM envelope WHERE source = 'paperless' ORDER BY id" + params: list = [] + if limit is not None: + params.append(limit) + query += f" LIMIT ${len(params)}" + if offset: + params.append(offset) + query += f" OFFSET ${len(params)}" + return await conn.fetch(query, *params) + + +async def fetch_existing_chunk_keys(conn: asyncpg.Connection, model: str) -> set[tuple[str, int]]: + """(envelope_id, chunk_index) pairs already embedded with this model — idempotency + dry-run preview.""" + rows = await conn.fetch( + "SELECT envelope_id, chunk_index FROM document_chunk WHERE model = $1", model + ) + return {(r["envelope_id"], r["chunk_index"]) for r in rows} + + +async def embed_chunk( + session: aiohttp.ClientSession, base_url: str, model: str, text: str +) -> tuple[list[float], float]: + """POST /api/embeddings on Ollama for one chunk. Returns (embedding, elapsed_seconds).""" + t0 = time.monotonic() + async with session.post(f"{base_url}/api/embeddings", json={"model": model, "prompt": text}) as resp: + resp.raise_for_status() + data = await resp.json() + elapsed = time.monotonic() - t0 + + embedding = data.get("embedding") + if not embedding: + raise ValueError(f"ollama response missing 'embedding': {data!r}") + return embedding, elapsed + + +async def insert_chunk( + conn: asyncpg.Connection, envelope_id: str, chunk_index: int, text: str, + embedding: list[float], model: str, +) -> str: + """Returns asyncpg's command tag (e.g. 'INSERT 0 1' or 'INSERT 0 0' if ON CONFLICT + DO NOTHING skipped the row) so the caller can tell a real insert from a no-op.""" + return await conn.execute( + _INSERT_SQL, envelope_id, chunk_index, text, _vector_literal(embedding), model + ) + + +def _rows_affected(command_tag: str) -> int: + """Parse the row count out of an asyncpg INSERT command tag ('INSERT ').""" + return int(command_tag.rsplit(" ", 1)[-1]) + + +async def run( + dsn: str, + ollama_url: str = DEFAULT_OLLAMA_URL, + model: str = DEFAULT_MODEL, + limit: Optional[int] = None, + offset: Optional[int] = None, + apply: bool = False, + chunk_size: int = TARGET_CHARS, + chunk_overlap: int = OVERLAP_CHARS, +) -> dict: + """Chunk + embed one --limit/--offset slice of `source='paperless'` envelopes. + + dry-run (apply=False) chunks and counts everything — no Ollama calls, no DB writes; + `chunks_inserted` reports what *would* be written (mirrors the rest of this job family). + Returns stats that must always balance: + + documents_fetched = empty_content + documents_chunked + chunks_total = chunks_already_embedded + chunks_inserted + + chunks_conflict_skipped + chunks_errors + + `chunks_conflict_skipped` counts inserts where `ON CONFLICT (envelope_id, chunk_index) + DO NOTHING` silently discarded the row (the pre-fetched `existing` set is the first line + of defense against this and should make it rare; see the module docstring's note on the + UNIQUE constraint not including `model`) — tracked separately so a silent no-op is never + miscounted as a successful write. + + Every embedding response's dimension is checked against `EXPECTED_DIM` (raises + EmbeddingDimensionError and aborts the whole run on mismatch) — never silently indexes + a vector that doesn't match the VECTOR(1024) column. A DB write failure for one chunk + (bad bytes, a dropped connection) is isolated and counted as `chunks_errors`, same as an + embed failure — it never aborts the run for the rest of the slice. + """ + stats = { + "documents_fetched": 0, + "empty_content": 0, + "documents_chunked": 0, + "chunks_total": 0, + "chunks_already_embedded": 0, + "chunks_inserted": 0, + "chunks_conflict_skipped": 0, + "chunks_errors": 0, + "embed_calls": 0, + "embed_seconds_total": 0.0, + } + + conn = await asyncpg.connect(dsn) + try: + docs = await fetch_documents(conn, limit, offset) + existing = await fetch_existing_chunk_keys(conn, model) + + session: Optional[aiohttp.ClientSession] = None + if apply: + session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=120)) + + try: + for row in docs: + stats["documents_fetched"] += 1 + envelope_id = row["id"] + entities = _decode_jsonb(row["entities"]) or [] + content = extract_content(entities) + chunks = chunk_text(content, chunk_size, chunk_overlap) + + if not chunks: + stats["empty_content"] += 1 + continue + + stats["documents_chunked"] += 1 + for idx, chunk in enumerate(chunks): + stats["chunks_total"] += 1 + key = (envelope_id, idx) + if key in existing: + stats["chunks_already_embedded"] += 1 + continue + + if not apply: + stats["chunks_inserted"] += 1 + continue + + assert session is not None + try: + embedding, elapsed = await embed_chunk(session, ollama_url, model, chunk) + except Exception: + _log.warning( + "skip.embed_error", envelope_id=envelope_id, chunk_index=idx, exc_info=True + ) + stats["chunks_errors"] += 1 + continue + + if len(embedding) != EXPECTED_DIM: + raise EmbeddingDimensionError( + f"ollama model={model!r} returned dim={len(embedding)}, " + f"expected {EXPECTED_DIM} (document_chunk.embedding is VECTOR({EXPECTED_DIM}))" + ) + + stats["embed_calls"] += 1 + stats["embed_seconds_total"] += elapsed + + try: + command_tag = await insert_chunk(conn, envelope_id, idx, chunk, embedding, model) + except Exception: + _log.warning( + "skip.insert_error", envelope_id=envelope_id, chunk_index=idx, exc_info=True + ) + stats["chunks_errors"] += 1 + continue + + existing.add(key) + if _rows_affected(command_tag) == 0: + _log.warning( + "chunk.conflict_skipped", envelope_id=envelope_id, chunk_index=idx, model=model + ) + stats["chunks_conflict_skipped"] += 1 + else: + stats["chunks_inserted"] += 1 + finally: + if session is not None: + await session.close() + finally: + await conn.close() + + balance_docs = stats["empty_content"] + stats["documents_chunked"] + balance_chunks = ( + stats["chunks_already_embedded"] + stats["chunks_inserted"] + + stats["chunks_conflict_skipped"] + stats["chunks_errors"] + ) + if balance_docs != stats["documents_fetched"] or balance_chunks != stats["chunks_total"]: + _log.error("stats_mismatch", **stats) + + _log.info("run_complete", apply=apply, model=model, **stats) + return stats + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Chunk source='paperless' envelope content and embed via Ollama/bge-m3 " + "into document_chunk (module 5, phase 2 — plan §6 step 6)." + ) + parser.add_argument("--dsn", default=os.environ.get("KB_DSN"), + help="asyncpg DSN for kb-postgres (or set KB_DSN env var)") + parser.add_argument("--ollama-url", default=os.environ.get("OLLAMA_URL", DEFAULT_OLLAMA_URL), + help=f"Ollama base URL (default: {DEFAULT_OLLAMA_URL}, or set OLLAMA_URL)") + parser.add_argument("--model", default=os.environ.get("OLLAMA_EMBED_MODEL", DEFAULT_MODEL), + help=f"Ollama embedding model (default: {DEFAULT_MODEL})") + parser.add_argument("--limit", type=int, default=None, + help="Max documents to process (default: all)") + parser.add_argument("--offset", type=int, default=0, + help="Slice offset, ordered by envelope id (default: 0)") + parser.add_argument("--chunk-size", type=int, default=TARGET_CHARS, + help=f"Target chunk size in characters (default: {TARGET_CHARS} ~= {TARGET_TOKENS} tok)") + parser.add_argument("--chunk-overlap", type=int, default=OVERLAP_CHARS, + help=f"Chunk overlap in characters (default: {OVERLAP_CHARS} ~= {OVERLAP_TOKENS} tok)") + parser.add_argument("--apply", action="store_true", + help="Actually call Ollama and insert chunks. Default is dry-run (chunk + count only).") + args = parser.parse_args() + + if not args.dsn: + _log.error("missing_dsn", hint="pass --dsn or set KB_DSN") + sys.exit(1) + if args.chunk_overlap >= args.chunk_size: + _log.error( + "invalid_chunk_params", chunk_size=args.chunk_size, chunk_overlap=args.chunk_overlap, + hint="--chunk-overlap must be smaller than --chunk-size", + ) + sys.exit(1) + + try: + stats = asyncio.run( + run( + dsn=args.dsn, + ollama_url=args.ollama_url, + model=args.model, + limit=args.limit, + offset=args.offset, + apply=args.apply, + chunk_size=args.chunk_size, + chunk_overlap=args.chunk_overlap, + ) + ) + except EmbeddingDimensionError as exc: + _log.error("dim_mismatch_abort", error=str(exc)) + sys.exit(1) + + mode = "APPLY" if args.apply else "DRY-RUN" + avg_embed = ( + stats["embed_seconds_total"] / stats["embed_calls"] if stats["embed_calls"] else 0.0 + ) + _log.info("summary", mode=mode, avg_embed_seconds_per_chunk=round(avg_embed, 4), **stats) + + balanced = ( + stats["documents_fetched"] == stats["empty_content"] + stats["documents_chunked"] + and stats["chunks_total"] == ( + stats["chunks_already_embedded"] + stats["chunks_inserted"] + + stats["chunks_conflict_skipped"] + stats["chunks_errors"] + ) + ) + failed = stats["chunks_errors"] > 0 or stats["chunks_conflict_skipped"] > 0 or not balanced + sys.exit(1 if failed else 0) + + +if __name__ == "__main__": + main() diff --git a/jobs/documents-ingest/tests/test_chunk_embed.py b/jobs/documents-ingest/tests/test_chunk_embed.py new file mode 100644 index 0000000..e50af53 --- /dev/null +++ b/jobs/documents-ingest/tests/test_chunk_embed.py @@ -0,0 +1,455 @@ +"""Unit tests for the chunk + embed job — no DB, no real HTTP, no real Ollama.""" +from __future__ import annotations + +import json + +import pytest + +from documents_ingest.chunk_embed import ( + EmbeddingDimensionError, + OVERLAP_CHARS, + TARGET_CHARS, + chunk_text, + embed_chunk, + extract_content, + fetch_documents, + fetch_existing_chunk_keys, + hard_split, + run, + split_paragraphs, + _vector_literal, +) + + +def _paragraph(char: str, length: int) -> str: + return char * length + + +class TestSplitParagraphs: + def test_splits_on_blank_lines(self): + text = "first para\n\nsecond para\n\nthird para" + assert split_paragraphs(text) == ["first para", "second para", "third para"] + + def test_no_blank_lines_returns_single_paragraph(self): + text = "line one\nline two\nline three" + assert split_paragraphs(text) == [text] + + def test_drops_empty_fragments(self): + text = "a\n\n\n\n\n\nb" + assert split_paragraphs(text) == ["a", "b"] + + +class TestHardSplit: + def test_overlap_equal_to_size_raises(self): + # step = size - overlap would be 0 -> `start` never advances -> infinite loop. + with pytest.raises(ValueError): + hard_split("x" * 1000, size=100, overlap=100) + + def test_overlap_greater_than_size_raises(self): + with pytest.raises(ValueError): + hard_split("x" * 1000, size=100, overlap=150) + + def test_short_text_returns_single_chunk(self): + assert hard_split("short", size=100, overlap=10) == ["short"] + + def test_splits_with_overlap(self): + text = "A" * 5000 + chunks = hard_split(text, size=2400, overlap=600) + assert len(chunks) == 3 + # consecutive windows overlap by exactly `overlap` characters + assert chunks[0][-600:] == chunks[1][:600] + assert chunks[1][-600:] == chunks[2][:600] + # covers the full text, in order, no gaps + assert chunks[0] + chunks[1][600:] + chunks[2][600:] == text + + def test_last_chunk_reaches_end_of_text(self): + text = "B" * 5000 + chunks = hard_split(text, size=2400, overlap=600) + assert chunks[-1] == text[-len(chunks[-1]):] + assert text.endswith(chunks[-1]) + + +class TestChunkText: + def test_empty_content_returns_no_chunks(self): + assert chunk_text("") == [] + assert chunk_text(None) == [] + assert chunk_text(" \n\n ") == [] + + def test_document_shorter_than_one_chunk_returns_single_chunk(self): + text = "A short document." + assert chunk_text(text, size=TARGET_CHARS, overlap=OVERLAP_CHARS) == [text] + + def test_multi_paragraph_document_splits_on_boundaries(self): + para1 = _paragraph("a", 1000) + para2 = _paragraph("b", 1000) + para3 = _paragraph("c", 1000) + text = f"{para1}\n\n{para2}\n\n{para3}" + + chunks = chunk_text(text, size=2400, overlap=600) + + assert len(chunks) == 2 + assert para1 in chunks[0] + assert para2 in chunks[0] + assert para3 in chunks[-1] + + def test_overlap_between_consecutive_chunks(self): + para1 = _paragraph("a", 1000) + para2 = _paragraph("b", 1000) + para3 = _paragraph("c", 1000) + text = f"{para1}\n\n{para2}\n\n{para3}" + + chunks = chunk_text(text, size=2400, overlap=600) + + # the tail of chunk[0] (the overlap window) reappears at the start of chunk[1] + assert chunks[0][-600:] == chunks[1][: len(chunks[0][-600:])] + + def test_single_oversized_paragraph_falls_back_to_hard_split(self): + text = "X" * 6000 # no blank lines at all + chunks = chunk_text(text, size=2400, overlap=600) + assert len(chunks) > 1 + assert all(len(c) <= 2400 for c in chunks) + + def test_paragraph_larger_than_target_is_hard_split_within_mixed_document(self): + small = _paragraph("s", 100) + huge = _paragraph("h", 6000) + text = f"{small}\n\n{huge}" + + chunks = chunk_text(text, size=2400, overlap=600) + + assert chunks[0] == small + assert len(chunks) > 2 + assert all(len(c) <= 2400 for c in chunks[1:]) + + +class TestExtractContent: + def test_finds_content_entity(self): + entities = [{"type": "filename", "value": "a.pdf"}, {"type": "content", "text": "hello"}] + assert extract_content(entities) == "hello" + + def test_missing_content_entity_returns_empty_string(self): + entities = [{"type": "filename", "value": "a.pdf"}] + assert extract_content(entities) == "" + + def test_none_text_returns_empty_string(self): + entities = [{"type": "content", "text": None}] + assert extract_content(entities) == "" + + def test_empty_entities_list(self): + assert extract_content([]) == "" + assert extract_content(None) == "" + + +class TestVectorLiteral: + def test_formats_as_bracketed_csv(self): + assert _vector_literal([0.1, 0.2, -0.3]) == "[0.1,0.2,-0.3]" + + +class _FakeEmbedResponse: + def __init__(self, payload, status=200): + self._payload = payload + self._status = status + + async def __aenter__(self): + return self + + async def __aexit__(self, *exc): + return False + + async def json(self): + return self._payload + + def raise_for_status(self): + if self._status >= 400: + raise RuntimeError(f"HTTP {self._status}") + + +class _FakeOllamaSession: + """Serves a fixed embedding vector for every /api/embeddings POST, or errors by filename.""" + + def __init__(self, dim=1024, fail_for=None): + self._dim = dim + self._fail_for = fail_for or set() + self.requests: list[dict] = [] + + def post(self, url, json): + self.requests.append({"url": url, "json": json}) + if json["prompt"] in self._fail_for: + return _FakeEmbedResponse({}, status=500) + return _FakeEmbedResponse({"embedding": [0.01] * self._dim}) + + async def close(self): + pass + + +class TestEmbedChunk: + async def test_returns_embedding_and_elapsed(self): + session = _FakeOllamaSession(dim=1024) + embedding, elapsed = await embed_chunk(session, "http://fake-ollama", "bge-m3", "hello world") + assert len(embedding) == 1024 + assert elapsed >= 0 + assert session.requests == [ + {"url": "http://fake-ollama/api/embeddings", "json": {"model": "bge-m3", "prompt": "hello world"}} + ] + + async def test_missing_embedding_key_raises(self): + class _EmptyResponse(_FakeEmbedResponse): + pass + + class _Session: + def post(self, url, json): + return _EmptyResponse({}) + + with pytest.raises(ValueError): + await embed_chunk(_Session(), "http://fake-ollama", "bge-m3", "text") + + +class _FakeConn: + def __init__(self, docs=None, existing_keys=None, existing_model="bge-m3", execute_results=None): + self._docs = docs or [] + self._existing_keys = list(existing_keys or []) + self._existing_model = existing_model + # Optional queue of command tags returned by successive execute() calls, in order + # (e.g. ["INSERT 0 0"] to simulate an ON CONFLICT DO NOTHING no-op). Defaults to a + # real insert every time. + self._execute_results = list(execute_results) if execute_results is not None else None + self.execute_calls: list[tuple] = [] + + async def fetch(self, query, *params): + if "FROM document_chunk" in query: + # Mirrors the real `WHERE model = $1` filter: existing_keys were "written" under + # existing_model, so a query for a different model must not see them. + if params and params[0] != self._existing_model: + return [] + return [{"envelope_id": eid, "chunk_index": idx} for eid, idx in self._existing_keys] + return self._docs + + async def execute(self, query, *params): + self.execute_calls.append(params) + if self._execute_results is not None: + return self._execute_results.pop(0) + return "INSERT 0 1" + + async def close(self): + pass + + +def _row(envelope_id, content): + return {"id": envelope_id, "entities": json.dumps([{"type": "content", "text": content}])} + + +class TestFetchHelpers: + async def test_fetch_documents_no_limit_offset(self): + conn = _FakeConn(docs=[_row("paperless:1", "hi")]) + docs = await fetch_documents(conn, limit=None, offset=None) + assert docs == [_row("paperless:1", "hi")] + + async def test_fetch_existing_chunk_keys(self): + conn = _FakeConn(existing_keys=[("paperless:1", 0), ("paperless:1", 1)]) + keys = await fetch_existing_chunk_keys(conn, model="bge-m3") + assert keys == {("paperless:1", 0), ("paperless:1", 1)} + + +class TestRun: + def _patch(self, monkeypatch, conn, ollama_session): + async def _fake_connect(dsn): + return conn + monkeypatch.setattr("documents_ingest.chunk_embed.asyncpg.connect", _fake_connect) + + def _fake_session_factory(*args, **kwargs): + return ollama_session + monkeypatch.setattr("documents_ingest.chunk_embed.aiohttp.ClientSession", _fake_session_factory) + + async def test_dry_run_counts_without_calling_ollama_or_db(self, monkeypatch): + conn = _FakeConn(docs=[_row("paperless:1", "short doc")]) + ollama = _FakeOllamaSession() + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", apply=False) + + assert stats["documents_fetched"] == 1 + assert stats["documents_chunked"] == 1 + assert stats["chunks_total"] == 1 + assert stats["chunks_inserted"] == 1 + assert ollama.requests == [] + assert conn.execute_calls == [] + + async def test_empty_content_counted_separately_not_as_error(self, monkeypatch): + conn = _FakeConn(docs=[_row("paperless:1", ""), _row("paperless:2", "some text")]) + ollama = _FakeOllamaSession() + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", apply=False) + + assert stats["empty_content"] == 1 + assert stats["documents_chunked"] == 1 + assert stats["documents_fetched"] == 2 + assert stats["chunks_errors"] == 0 + + async def test_apply_embeds_and_inserts(self, monkeypatch): + conn = _FakeConn(docs=[_row("paperless:1", "some real content")]) + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", apply=True) + + assert stats["chunks_inserted"] == 1 + assert stats["embed_calls"] == 1 + assert len(conn.execute_calls) == 1 + assert len(ollama.requests) == 1 + + async def test_idempotent_skips_already_embedded_chunks(self, monkeypatch): + conn = _FakeConn( + docs=[_row("paperless:1", "some real content")], + existing_keys=[("paperless:1", 0)], + ) + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", apply=True) + + assert stats["chunks_already_embedded"] == 1 + assert stats["chunks_inserted"] == 0 + assert ollama.requests == [] + assert conn.execute_calls == [] + + async def test_rerun_after_apply_inserts_nothing_new(self, monkeypatch): + doc = _row("paperless:1", "some real content") + + conn1 = _FakeConn(docs=[doc]) + ollama1 = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn1, ollama1) + first = await run(dsn="postgresql://fake", apply=True) + assert first["chunks_inserted"] == 1 + + conn2 = _FakeConn(docs=[doc], existing_keys=[("paperless:1", 0)]) + ollama2 = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn2, ollama2) + second = await run(dsn="postgresql://fake", apply=True) + + assert second["chunks_inserted"] == 0 + assert second["chunks_already_embedded"] == 1 + assert ollama2.requests == [] + + async def test_dimension_mismatch_aborts(self, monkeypatch): + conn = _FakeConn(docs=[_row("paperless:1", "some real content")]) + ollama = _FakeOllamaSession(dim=768) # wrong dim + self._patch(monkeypatch, conn, ollama) + + with pytest.raises(EmbeddingDimensionError): + await run(dsn="postgresql://fake", apply=True) + + async def test_embed_error_is_isolated_and_counted(self, monkeypatch): + # Two chunks: force one to fail via a document long enough to produce 2 chunks + # (chunk_size/overlap kept small to make the test fast and explicit). + big = "a" * 100 + "\n\n" + "b" * 100 + conn = _FakeConn(docs=[_row("paperless:1", big)]) + ollama = _FakeOllamaSession(dim=1024, fail_for={"a" * 100}) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", apply=True, chunk_size=100, chunk_overlap=20) + + assert stats["chunks_errors"] == 1 + assert stats["chunks_inserted"] == 1 + assert stats["chunks_total"] == ( + stats["chunks_already_embedded"] + stats["chunks_inserted"] + + stats["chunks_conflict_skipped"] + stats["chunks_errors"] + ) + + async def test_insert_error_is_isolated_and_counted(self, monkeypatch): + # Two chunks; the first chunk's INSERT raises (simulated DB blip) but the run + # continues and the second chunk still gets embedded and inserted normally. + big = "a" * 100 + "\n\n" + "b" * 100 + conn = _FakeConn(docs=[_row("paperless:1", big)]) + + real_execute = conn.execute + + async def _flaky_execute(query, *params): + if params[1] == 0: # chunk_index 0 + conn.execute_calls.append(params) + raise RuntimeError("simulated db error") + return await real_execute(query, *params) + + conn.execute = _flaky_execute + + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", apply=True, chunk_size=100, chunk_overlap=20) + + assert stats["chunks_errors"] == 1 + assert stats["chunks_inserted"] == 1 + assert stats["chunks_total"] == ( + stats["chunks_already_embedded"] + stats["chunks_inserted"] + + stats["chunks_conflict_skipped"] + stats["chunks_errors"] + ) + + async def test_conflict_skipped_counted_separately_from_inserted(self, monkeypatch): + # ON CONFLICT DO NOTHING no-op: command tag reports 0 rows affected. + conn = _FakeConn( + docs=[_row("paperless:1", "some real content")], execute_results=["INSERT 0 0"] + ) + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", apply=True) + + assert stats["chunks_conflict_skipped"] == 1 + assert stats["chunks_inserted"] == 0 + assert stats["chunks_total"] == ( + stats["chunks_already_embedded"] + stats["chunks_inserted"] + + stats["chunks_conflict_skipped"] + stats["chunks_errors"] + ) + + async def test_existing_keys_scoped_to_model(self, monkeypatch): + # existing_keys were "written" under a different model — the WHERE model = $1 + # filter must not treat them as already-embedded for the current --model. + conn = _FakeConn( + docs=[_row("paperless:1", "some real content")], + existing_keys=[("paperless:1", 0)], + existing_model="some-other-model", + ) + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", apply=True, model="bge-m3") + + assert stats["chunks_already_embedded"] == 0 + assert stats["chunks_inserted"] == 1 + + async def test_stats_balance_documents_and_chunks(self, monkeypatch): + conn = _FakeConn(docs=[ + _row("paperless:1", "content one"), + _row("paperless:2", ""), + _row("paperless:3", "content three"), + ]) + ollama = _FakeOllamaSession(dim=1024) + self._patch(monkeypatch, conn, ollama) + + stats = await run(dsn="postgresql://fake", apply=True) + + assert stats["documents_fetched"] == stats["empty_content"] + stats["documents_chunked"] + assert stats["chunks_total"] == ( + stats["chunks_already_embedded"] + stats["chunks_inserted"] + + stats["chunks_conflict_skipped"] + stats["chunks_errors"] + ) + + async def test_limit_and_offset_passed_through_to_query(self, monkeypatch): + captured = {} + + class _Conn(_FakeConn): + async def fetch(self, query, *params): + if "FROM document_chunk" in query: + return [] + captured["query"] = query + captured["params"] = params + return [] + + conn = _Conn() + ollama = _FakeOllamaSession() + self._patch(monkeypatch, conn, ollama) + + await run(dsn="postgresql://fake", apply=False, limit=10, offset=5) + + assert "LIMIT $1" in captured["query"] + assert "OFFSET $2" in captured["query"] + assert captured["params"] == (10, 5)