427 lines
17 KiB
Python
427 lines
17 KiB
Python
|
"""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()
|