From 4cf801679365dfcf9291e9c6fd6e40b55e5b1b7d Mon Sep 17 00:00:00 2001 From: oskar Date: Wed, 22 Jul 2026 19:00:58 +0200 Subject: [PATCH] refactor(kb-mail): extract chunker to packages/kb-mail (faza mailowa Krok 0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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). --- .../src/documents_ingest/chunk_embed.py | 81 +------------ .../tests/test_chunk_embed.py | 106 +---------------- packages/kb-mail/src/kb_mail/chunking.py | 85 ++++++++++++++ packages/kb-mail/tests/test_chunking.py | 107 ++++++++++++++++++ 4 files changed, 202 insertions(+), 177 deletions(-) create mode 100644 packages/kb-mail/src/kb_mail/chunking.py create mode 100644 packages/kb-mail/tests/test_chunking.py diff --git a/jobs/documents-ingest/src/documents_ingest/chunk_embed.py b/jobs/documents-ingest/src/documents_ingest/chunk_embed.py index 645ffd7..b6e27d8 100644 --- a/jobs/documents-ingest/src/documents_ingest/chunk_embed.py +++ b/jobs/documents-ingest/src/documents_ingest/chunk_embed.py @@ -1,6 +1,11 @@ """Chunk + embed job — module 5 phase 2, plan step 6 (docs/kb/modules/05-faza2-plan.md, §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- 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` @@ -60,25 +65,13 @@ from typing import Optional import aiohttp import asyncpg 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 _log = structlog.get_logger(__name__) 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, excluded_reason) 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.""" -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 "".""" diff --git a/jobs/documents-ingest/tests/test_chunk_embed.py b/jobs/documents-ingest/tests/test_chunk_embed.py index c9f4124..75a90bc 100644 --- a/jobs/documents-ingest/tests/test_chunk_embed.py +++ b/jobs/documents-ingest/tests/test_chunk_embed.py @@ -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 import json @@ -7,18 +10,13 @@ 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, insert_chunk, is_ocr_junk, run, - split_paragraphs, _vector_literal, ) @@ -27,102 +25,6 @@ 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"}] diff --git a/packages/kb-mail/src/kb_mail/chunking.py b/packages/kb-mail/src/kb_mail/chunking.py new file mode 100644 index 0000000..fc910b4 --- /dev/null +++ b/packages/kb-mail/src/kb_mail/chunking.py @@ -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 diff --git a/packages/kb-mail/tests/test_chunking.py b/packages/kb-mail/tests/test_chunking.py new file mode 100644 index 0000000..dfd05fc --- /dev/null +++ b/packages/kb-mail/tests/test_chunking.py @@ -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:])