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:
parent
a8c2071e11
commit
4cf8016793
|
|
@ -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 ""."""
|
||||||
|
|
|
||||||
|
|
@ -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"}]
|
||||||
|
|
|
||||||
85
packages/kb-mail/src/kb_mail/chunking.py
Normal file
85
packages/kb-mail/src/kb_mail/chunking.py
Normal 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
|
||||||
107
packages/kb-mail/tests/test_chunking.py
Normal file
107
packages/kb-mail/tests/test_chunking.py
Normal 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:])
|
||||||
Loading…
Reference in a new issue