oskar/homelab-codex-ws
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

feat(documents-ingest): chunk + embed job (module 5 phase 2 step 6)

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
oskar 2026-07-15 20:41:00 +02:00
parent 5218925408
commit 09bb59a1db
4 changed files with 1056 additions and 0 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

174
jobs/documents-ingest/README.md
View file

@ -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 without operator confirmation (this job reads production Paperless data and
writes production envelope rows on `--apply`). `pytest` passes locally before writes production envelope rows on `--apply`). `pytest` passes locally before
this commit. 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:<pw>@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:<pw>@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.

1
jobs/documents-ingest/pyproject.toml
View file

@ -16,6 +16,7 @@ dependencies = [
[project.scripts] [project.scripts]
documents-ingest = "documents_ingest.extractor:main" documents-ingest = "documents_ingest.extractor:main"
documents-ingest-paperless = "documents_ingest.paperless_adapter:main" documents-ingest-paperless = "documents_ingest.paperless_adapter:main"
documents-ingest-embed = "documents_ingest.chunk_embed:main"
[tool.setuptools.packages.find] [tool.setuptools.packages.find]
where = ["src"] where = ["src"]

426
jobs/documents-ingest/src/documents_ingest/chunk_embed.py Normal file
View file

@ -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:<pw>@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 <oid> <rows>')."""
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()

455
jobs/documents-ingest/tests/test_chunk_embed.py Normal file
View file

@ -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)