From ec5b65871d4a2e40851b1a818fbe4165d666bd9f Mon Sep 17 00:00:00 2001 From: n0nuser Date: Mon, 6 Jul 2026 23:11:24 +0200 Subject: [PATCH 1/2] fix: cosine distance space and boundary-aware overlap for oversized paragraphs Chroma collection now created with hnsw:space=cosine to match embedding models' cosine training objective (existing collections need delete + re-ingest to pick this up). _split_long_paragraph in the structural chunker previously raw-sliced oversized paragraphs with zero overlap and no boundary awareness; it now reuses chunk_overlap_chars and snaps splits to sentence/word boundaries, falling back to raw slice only when no boundary is found. --- docs/rag-retrieval.md | 9 ++++- localrag/ingestion/structural_chunker.py | 49 +++++++++++++++++++----- localrag/storage/vector_store.py | 5 ++- tests/test_structural_chunker.py | 29 ++++++++++++++ tests/test_vector_store.py | 7 +++- 5 files changed, 86 insertions(+), 13 deletions(-) diff --git a/docs/rag-retrieval.md b/docs/rag-retrieval.md index 55ddb8d..88b949a 100644 --- a/docs/rag-retrieval.md +++ b/docs/rag-retrieval.md @@ -15,7 +15,14 @@ Configuration lives in `localrag/settings.py` and `.env.example`. `localrag/rag/retriever.py` combines candidates from: -1. Vector search (`VectorStore.query`) ranked by embedding distance. +1. Vector search (`VectorStore.query`) ranked by embedding distance. The Chroma + collection is created with `hnsw:space=cosine` (`localrag/storage/vector_store.py`), + matching the cosine-similarity objective most embedding models (including the + default `nomic-embed-text`) are trained against. Chroma applies collection + metadata only at creation time — an existing collection created before this + setting keeps its original (`l2`) space; delete and re-ingest (or + `POST /collections/rebuild` after a manual `delete_collection`) to pick up + `cosine`. 2. BM25 lexical search (`Bm25Index.query`) ranked by lexical relevance. Candidates are merged by reciprocal rank fusion: diff --git a/localrag/ingestion/structural_chunker.py b/localrag/ingestion/structural_chunker.py index 144bb7c..e20cc36 100644 --- a/localrag/ingestion/structural_chunker.py +++ b/localrag/ingestion/structural_chunker.py @@ -7,6 +7,7 @@ from localrag.settings import Settings _HEADING_PATTERN = re.compile(r"^(#{1,6})\s+(.*)$") +_SENTENCE_BOUNDARY_PATTERN = re.compile(r"(?<=[.!?])\s+") @dataclass @@ -26,23 +27,26 @@ def chunk_document(text: str, file_type: str, settings: Settings) -> list[Chunk] text=cleaned_text, min_chars=settings.chunk_min_chars, max_chars=settings.chunk_max_chars, + overlap_chars=settings.chunk_overlap_chars, ) if file_type in CODE_EXTENSIONS: return _chunk_non_markdown( text=cleaned_text, min_chars=settings.chunk_min_chars, max_chars=settings.chunk_max_chars, + overlap_chars=settings.chunk_overlap_chars, is_code=True, ) return _chunk_non_markdown( text=cleaned_text, min_chars=settings.chunk_min_chars, max_chars=settings.chunk_max_chars, + overlap_chars=settings.chunk_overlap_chars, is_code=False, ) -def _chunk_markdown(text: str, min_chars: int, max_chars: int) -> list[Chunk]: +def _chunk_markdown(text: str, min_chars: int, max_chars: int, overlap_chars: int) -> list[Chunk]: lines = text.splitlines() heading_stack: list[str] = [] sections: list[tuple[str, str]] = [] @@ -74,12 +78,15 @@ def flush_section() -> None: text=text, min_chars=min_chars, max_chars=max_chars, + overlap_chars=overlap_chars, is_code=False, ) chunks: list[Chunk] = [] for section_text, heading_path in sections: - for item in _pack_blocks(section_text, min_chars=min_chars, max_chars=max_chars): + for item in _pack_blocks( + section_text, min_chars=min_chars, max_chars=max_chars, overlap_chars=overlap_chars + ): chunk_type = "markdown_section" if _looks_like_table(item): chunk_type = "markdown_table" @@ -89,16 +96,20 @@ def flush_section() -> None: return chunks -def _chunk_non_markdown(text: str, min_chars: int, max_chars: int, *, is_code: bool) -> list[Chunk]: +def _chunk_non_markdown( + text: str, min_chars: int, max_chars: int, overlap_chars: int, *, is_code: bool +) -> list[Chunk]: chunks: list[Chunk] = [] - for item in _pack_blocks(text, min_chars=min_chars, max_chars=max_chars, is_code=is_code): + for item in _pack_blocks( + text, min_chars=min_chars, max_chars=max_chars, overlap_chars=overlap_chars, is_code=is_code + ): chunk_type = "code_block" if is_code else "text_block" chunks.append(Chunk(text=item, heading_path="", chunk_type=chunk_type)) return chunks def _pack_blocks( # noqa: C901 - text: str, min_chars: int, max_chars: int, *, is_code: bool = False + text: str, min_chars: int, max_chars: int, overlap_chars: int = 0, *, is_code: bool = False ) -> list[str]: blocks = _split_blocks(text=text, is_code=is_code) if not blocks: @@ -118,7 +129,9 @@ def _pack_blocks( # noqa: C901 if current: packed.append(current) current = "" - packed.extend(_split_long_paragraph(normalized, effective_max_chars)) + packed.extend( + _split_long_paragraph(normalized, effective_max_chars, overlap_chars) + ) continue candidate = normalized if not current else f"{current}\n\n{normalized}" @@ -208,17 +221,33 @@ def _split_blocks(text: str, *, is_code: bool) -> list[tuple[str, str]]: # noqa return blocks -def _split_long_paragraph(text: str, max_chars: int) -> list[str]: +def _split_long_paragraph(text: str, max_chars: int, overlap_chars: int) -> list[str]: + text_len = len(text) + safe_overlap = max(0, min(overlap_chars, max_chars - 1)) chunks: list[str] = [] start = 0 - while start < len(text): - chunk = text[start : start + max_chars].strip() + while start < text_len: + end = min(start + max_chars, text_len) + if end < text_len: + end = _snap_to_boundary(text, start, end) + chunk = text[start:end].strip() if chunk: chunks.append(chunk) - start += max_chars + start = end - safe_overlap if end < text_len else text_len return chunks +def _snap_to_boundary(text: str, start: int, end: int) -> int: + window = text[start:end] + sentence_breaks = [start + m.start() for m in _SENTENCE_BOUNDARY_PATTERN.finditer(window)] + if sentence_breaks: + return sentence_breaks[-1] + last_space = window.rfind(" ") + if last_space > 0: + return start + last_space + return end + + def _is_table_row(line: str) -> bool: return line.startswith("|") and "|" in line[1:] diff --git a/localrag/storage/vector_store.py b/localrag/storage/vector_store.py index ae7fcf2..e8e38c0 100644 --- a/localrag/storage/vector_store.py +++ b/localrag/storage/vector_store.py @@ -21,7 +21,10 @@ class VectorStore: def create(cls, persist_path: str, collection_name: str) -> VectorStore: Path(persist_path).mkdir(parents=True, exist_ok=True) client = chromadb.PersistentClient(path=persist_path) - collection = client.get_or_create_collection(name=collection_name) + collection = client.get_or_create_collection( + name=collection_name, + metadata={"hnsw:space": "cosine"}, + ) logger.info( "vector_store_ready persist_path=%s collection=%s", persist_path, diff --git a/tests/test_structural_chunker.py b/tests/test_structural_chunker.py index c1b9f65..70c9692 100644 --- a/tests/test_structural_chunker.py +++ b/tests/test_structural_chunker.py @@ -1,5 +1,7 @@ from __future__ import annotations +from itertools import pairwise + from localrag.ingestion.structural_chunker import chunk_document from localrag.settings import Settings @@ -52,6 +54,33 @@ def test_chunk_document_markdown_splits_oversized_paragraph() -> None: assert all(chunk.heading_path == "Long" for chunk in chunks) +def test_chunk_document_oversized_paragraph_overlaps_between_chunks() -> None: + oversized = "A" * 30 + markdown_text = f"# Long\n\n{oversized}" + settings = Settings(chunk_max_chars=10, chunk_min_chars=1, chunk_overlap_chars=3) + + chunks = chunk_document(markdown_text, ".md", settings) + + body_chunks = [chunk for chunk in chunks if chunk.text != "# Long"] + assert len(body_chunks) > 1 + for prev_chunk, next_chunk in pairwise(body_chunks): + assert prev_chunk.text[-3:] == next_chunk.text[:3] + + +def test_chunk_document_oversized_paragraph_splits_on_sentence_boundary() -> None: + sentence = "This is one sentence." + long_text = " ".join([sentence] * 6) + markdown_text = f"# Notes\n\n{long_text}" + settings = Settings(chunk_max_chars=30, chunk_min_chars=1, chunk_overlap_chars=0) + + chunks = chunk_document(markdown_text, ".md", settings) + + body_chunks = [chunk for chunk in chunks if chunk.text != "# Notes"] + assert len(body_chunks) > 1 + for chunk in body_chunks: + assert chunk.text == sentence + + def test_chunk_document_non_markdown_packs_paragraphs() -> None: text = "First paragraph.\n\nSecond paragraph.\n\nThird paragraph." settings = Settings(chunk_max_chars=40, chunk_min_chars=20) diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py index 3bebe7a..5da02da 100644 --- a/tests/test_vector_store.py +++ b/tests/test_vector_store.py @@ -234,9 +234,13 @@ class FakeClient: def __init__(self, path: str) -> None: self.path = path self.got_collection_name: str | None = None + self.got_metadata: dict[str, object] | None = None - def get_or_create_collection(self, name: str) -> object: + def get_or_create_collection( + self, name: str, metadata: dict[str, object] | None = None + ) -> object: self.got_collection_name = name + self.got_metadata = metadata return sentinel_collection def fake_persistent_client(path: str) -> FakeClient: @@ -248,3 +252,4 @@ def fake_persistent_client(path: str) -> FakeClient: assert created_dir.exists() assert store.collection is sentinel_collection + assert store.client.got_metadata == {"hnsw:space": "cosine"} From b1f19b85fc9225ff8b82bbd3793299a153609ff7 Mon Sep 17 00:00:00 2001 From: n0nuser Date: Mon, 6 Jul 2026 23:13:06 +0200 Subject: [PATCH 2/2] style: apply ruff format --- localrag/ingestion/structural_chunker.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/localrag/ingestion/structural_chunker.py b/localrag/ingestion/structural_chunker.py index e20cc36..ea7fcd8 100644 --- a/localrag/ingestion/structural_chunker.py +++ b/localrag/ingestion/structural_chunker.py @@ -129,9 +129,7 @@ def _pack_blocks( # noqa: C901 if current: packed.append(current) current = "" - packed.extend( - _split_long_paragraph(normalized, effective_max_chars, overlap_chars) - ) + packed.extend(_split_long_paragraph(normalized, effective_max_chars, overlap_chars)) continue candidate = normalized if not current else f"{current}\n\n{normalized}"