refactor(kb-mail): extract chunker to packages/kb-mail (faza mailowa Krok 0)

chunk_text/hard_split/split_paragraphs + 2400/600-char constants move out of
documents_ingest.chunk_embed into kb_mail.chunking so the upcoming
jobs/mail-body-ingest shares the exact same chunker instead of a copy-pasted
drift (the phase-4 lesson for retrieval.py -> packages/kb-retrieval, applied
again). chunk_embed.py re-exports unchanged; zero behavior change, tests moved
1:1 alongside the code (docs/kb/modules/05-faza-mailowa-plan.md, §3).
This commit is contained in:
oskar 2026-07-22 19:00:58 +02:00
parent f0e9533025
commit 348ce10b1d
4 changed files with 202 additions and 177 deletions

View file

@ -1,6 +1,11 @@
"""Chunk + embed job — module 5 phase 2, plan step 6 (docs/kb/modules/05-faza2-plan.md, """Chunk + embed job — module 5 phase 2, plan step 6 (docs/kb/modules/05-faza2-plan.md,
§6 step 6, decision 3). §6 step 6, decision 3).
`chunk_text`/`hard_split`/`split_paragraphs`/`TARGET_CHARS`/`OVERLAP_CHARS` moved to
`kb_mail.chunking` in module 5 faza mailowa, Krok 0 (docs/kb/modules/05-faza-mailowa-plan.md, §3)
so `jobs/mail-body-ingest` shares the exact same chunker instead of a copy-pasted drift; re-exported
here unchanged so nothing importing them from this module breaks.
Pipeline: `envelope(source='paperless').entities[type=content].text` -> chunk (paragraph- Pipeline: `envelope(source='paperless').entities[type=content].text` -> chunk (paragraph-
preferring, ~600 tok/chunk, ~150 tok overlap, hard char-fallback for oversized paragraphs) 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` -> `POST /api/embeddings` (Ollama on SOLARIA, model `bge-m3`) -> `INSERT document_chunk`
@ -60,25 +65,13 @@ from typing import Optional
import aiohttp import aiohttp
import asyncpg import asyncpg
import structlog import structlog
from kb_mail.chunking import OVERLAP_CHARS, TARGET_CHARS, chunk_text, hard_split, split_paragraphs
from kb_retrieval.embed import DEFAULT_MODEL, DEFAULT_OLLAMA_URL, _vector_literal, embed_chunk from kb_retrieval.embed import DEFAULT_MODEL, DEFAULT_OLLAMA_URL, _vector_literal, embed_chunk
_log = structlog.get_logger(__name__) _log = structlog.get_logger(__name__)
EXPECTED_DIM = 1024 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_SQL = """
INSERT INTO document_chunk (envelope_id, chunk_index, text, embedding, model, excluded_reason) INSERT INTO document_chunk (envelope_id, chunk_index, text, embedding, model, excluded_reason)
VALUES ($1, $2, $3, $4::vector, $5, $6) VALUES ($1, $2, $3, $4::vector, $5, $6)
@ -128,68 +121,6 @@ class EmbeddingDimensionError(RuntimeError):
"""Ollama returned a vector of the wrong dimension for the target `document_chunk` schema.""" """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: def extract_content(entities: list) -> str:
"""Pull `entities[type=content].text` out of a document envelope (plan §4.2). Missing """Pull `entities[type=content].text` out of a document envelope (plan §4.2). Missing
or empty content both yield "".""" or empty content both yield ""."""

View file

@ -1,4 +1,7 @@
"""Unit tests for the chunk + embed job — no DB, no real HTTP, no real Ollama.""" """Unit tests for the chunk + embed job — no DB, no real HTTP, no real Ollama.
`chunk_text`/`hard_split`/`split_paragraphs` tests moved to
packages/kb-mail/tests/test_chunking.py (module 5, faza mailowa, Krok 0) alongside the code."""
from __future__ import annotations from __future__ import annotations
import json import json
@ -7,18 +10,13 @@ import pytest
from documents_ingest.chunk_embed import ( from documents_ingest.chunk_embed import (
EmbeddingDimensionError, EmbeddingDimensionError,
OVERLAP_CHARS,
TARGET_CHARS,
chunk_text,
embed_chunk, embed_chunk,
extract_content, extract_content,
fetch_documents, fetch_documents,
fetch_existing_chunk_keys, fetch_existing_chunk_keys,
hard_split,
insert_chunk, insert_chunk,
is_ocr_junk, is_ocr_junk,
run, run,
split_paragraphs,
_vector_literal, _vector_literal,
) )
@ -27,102 +25,6 @@ def _paragraph(char: str, length: int) -> str:
return char * length 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: class TestExtractContent:
def test_finds_content_entity(self): def test_finds_content_entity(self):
entities = [{"type": "filename", "value": "a.pdf"}, {"type": "content", "text": "hello"}] entities = [{"type": "filename", "value": "a.pdf"}, {"type": "content", "text": "hello"}]

View file

@ -0,0 +1,85 @@
"""Shared chunker — module 5, phase mailowa, plan step 0 (docs/kb/modules/05-faza-mailowa-plan.md,
§3). Moved 1:1 out of `jobs/documents_ingest/chunk_embed.py` so both the paperless job and
`jobs/mail-body-ingest` chunk against the exact same tested logic instead of two implementations
drifting apart (the phase-4 lesson for `retrieval.py` -> `packages/kb-retrieval`, applied again).
No new dependencies here `is_ocr_junk` (OCR-corpus calibrated) and the embed client stay in
their own modules; this file is chunking only.
"""
from __future__ import annotations
import re
# ~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.
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")
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.
Empty/whitespace-only text 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

View file

@ -0,0 +1,107 @@
"""Unit tests for the shared chunker — moved 1:1 from
jobs/documents-ingest/tests/test_chunk_embed.py (module 5, faza mailowa, plan §3, Krok 0)."""
from __future__ import annotations
import pytest
from kb_mail.chunking import OVERLAP_CHARS, TARGET_CHARS, chunk_text, hard_split, split_paragraphs
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:])