diff --git a/src/openjarvis/connectors/chunker.py b/src/openjarvis/connectors/chunker.py new file mode 100644 index 00000000..f0515497 --- /dev/null +++ b/src/openjarvis/connectors/chunker.py @@ -0,0 +1,284 @@ +"""Type-aware semantic chunker for Deep Research ingestion. + +Splits text based on document type, never splitting mid-sentence. +Returns ``ChunkResult`` dataclass objects with section metadata and +inherited parent metadata. + +Splitting strategy by doc_type +------------------------------- +- ``event``, ``contact`` : Always a single chunk; never split. +- ``email`` : Split on reply boundaries (``On … wrote:``), + then sentence-split within each part. +- ``message`` : Split on double-newline boundaries, accumulate + into chunks up to *max_tokens*. +- ``document``, ``note``, + anything else : Split on ``## Heading`` section boundaries → + paragraph boundaries (``\\n\\n``) within sections → + sentence boundaries as a last resort. + +Token counting uses whitespace splitting: ``len(text.split())``. +""" + +from __future__ import annotations + +import re +from dataclasses import dataclass, field +from typing import Any, Dict, List, Optional + +# --------------------------------------------------------------------------- +# Public types +# --------------------------------------------------------------------------- + +_SENTENCE_SPLIT_RE = re.compile(r'(?<=[.!?])\s+(?=[A-Z"])') +_SECTION_RE = re.compile(r"(?m)^##\s+(.+)$") +_REPLY_BOUNDARY_RE = re.compile(r"(?m)^On .+wrote:\s*$") + + +@dataclass(slots=True) +class ChunkResult: + """A single chunk produced by ``SemanticChunker.chunk()``.""" + + content: str + index: int = 0 + metadata: Dict[str, Any] = field(default_factory=dict) + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _count_tokens(text: str) -> int: + """Approximate token count via whitespace splitting.""" + return len(text.split()) + + +def _split_sentences(text: str) -> List[str]: + """Split *text* into sentences using the canonical regex. + + The regex splits after sentence-ending punctuation (``.``, ``!``, ``?``) + followed by whitespace and a capital letter or a double-quote. + """ + parts = _SENTENCE_SPLIT_RE.split(text) + return [p.strip() for p in parts if p.strip()] + + +def _accumulate( + segments: List[str], + *, + max_tokens: int, + sep: str = " ", +) -> List[str]: + """Greedily merge *segments* into chunks up to *max_tokens* tokens. + + A segment that is already larger than *max_tokens* is placed in its own + chunk; it is never split further by this function. + """ + chunks: List[str] = [] + current_parts: List[str] = [] + current_tokens = 0 + + for seg in segments: + seg_tokens = _count_tokens(seg) + if current_parts and current_tokens + seg_tokens > max_tokens: + chunks.append(sep.join(current_parts)) + current_parts = [seg] + current_tokens = seg_tokens + else: + current_parts.append(seg) + current_tokens += seg_tokens + + if current_parts: + chunks.append(sep.join(current_parts)) + + return chunks + + +def _sentence_chunks(text: str, *, max_tokens: int) -> List[str]: + """Split *text* by sentences and accumulate into max_tokens chunks.""" + sentences = _split_sentences(text) + if not sentences: + stripped = text.strip() + return [stripped] if stripped else [] + return _accumulate(sentences, max_tokens=max_tokens, sep=" ") + + +def _paragraph_chunks(text: str, *, max_tokens: int) -> List[str]: + """Split *text* on paragraph breaks (``\\n\\n``), then by sentences if needed.""" + paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()] + result: List[str] = [] + for para in paragraphs: + if _count_tokens(para) <= max_tokens: + result.append(para) + else: + result.extend(_sentence_chunks(para, max_tokens=max_tokens)) + return result + + +# --------------------------------------------------------------------------- +# SemanticChunker +# --------------------------------------------------------------------------- + + +class SemanticChunker: + """Split text based on document type without breaking mid-sentence. + + Parameters + ---------- + max_tokens: + Soft upper limit on chunk size measured in whitespace-delimited tokens + (i.e. ``len(text.split())``). Single unsplittable segments may exceed + this limit. + """ + + def __init__(self, max_tokens: int = 512) -> None: + self.max_tokens = max_tokens + + # ------------------------------------------------------------------ + # Public API + # ------------------------------------------------------------------ + + def chunk( + self, + text: str, + *, + doc_type: str = "document", + metadata: Optional[Dict[str, Any]] = None, + ) -> List[ChunkResult]: + """Split *text* into ``ChunkResult`` objects. + + Parameters + ---------- + text: The raw text to split. + doc_type: Controls the splitting strategy (see module docstring). + metadata: Parent metadata dict; copied into every chunk's ``metadata``. + + Returns + ------- + A list of ``ChunkResult`` objects with sequential 0-based ``index`` + values. Returns an empty list if *text* is empty or whitespace-only. + """ + if not text or not text.strip(): + return [] + + parent_meta: Dict[str, Any] = dict(metadata or {}) + + if doc_type in ("event", "contact"): + raw_chunks = self._chunk_atomic(text) + elif doc_type == "email": + raw_chunks = self._chunk_email(text) + elif doc_type == "message": + raw_chunks = self._chunk_message(text) + else: + # "document", "note", or any unknown type + raw_chunks = self._chunk_document(text) + + results: List[ChunkResult] = [] + for idx, (content, extra_meta) in enumerate(raw_chunks): + merged: Dict[str, Any] = dict(parent_meta) + merged.update(extra_meta) + results.append(ChunkResult(content=content, index=idx, metadata=merged)) + + return results + + # ------------------------------------------------------------------ + # Strategy implementations + # ------------------------------------------------------------------ + + def _chunk_atomic(self, text: str) -> List[tuple[str, Dict[str, Any]]]: + """Return the entire text as a single chunk (event / contact).""" + return [(text, {})] + + def _chunk_email(self, text: str) -> List[tuple[str, Dict[str, Any]]]: + """Split on reply boundaries; sentence-split each part.""" + # Split the email into parts on "On ... wrote:" lines. + # re.split with a capturing group keeps the boundary in results, + # so we re-attach the header to the following segment. + boundaries = _REPLY_BOUNDARY_RE.split(text) + + # Each boundary match is a separator; reassemble so the "On … wrote:" + # line stays with the content that follows it (the quoted block). + raw_parts: List[str] = [] + if boundaries: + # The first element is the text before the first boundary (the + # main reply body). + raw_parts.append(boundaries[0]) + # Subsequent elements alternate: matched boundary, then text after. + # Because we used split() (not findall), the boundaries themselves + # are not in the list — only the text segments between them. + # So boundaries[1:] are the segments after each matched header. + # We need to re-find the headers to reassemble. + headers = _REPLY_BOUNDARY_RE.findall(text) + for header, body in zip(headers, boundaries[1:]): + # We found the header text via findall; reconstruct the part. + part = (header.strip() + "\n" + body).strip() + raw_parts.append(part) + + chunks: List[tuple[str, Dict[str, Any]]] = [] + for part in raw_parts: + part = part.strip() + if not part: + continue + if _count_tokens(part) <= self.max_tokens: + chunks.append((part, {})) + else: + for sub in _sentence_chunks(part, max_tokens=self.max_tokens): + if sub: + chunks.append((sub, {})) + + return chunks if chunks else [(text.strip(), {})] + + def _chunk_message(self, text: str) -> List[tuple[str, Dict[str, Any]]]: + """Split on double-newline boundaries and accumulate up to max_tokens.""" + paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()] + raw_chunks = _accumulate(paragraphs, max_tokens=self.max_tokens, sep="\n\n") + return [(c, {}) for c in raw_chunks if c] + + def _chunk_document(self, text: str) -> List[tuple[str, Dict[str, Any]]]: + """Split on ## headings → paragraphs → sentences.""" + # Find all ## heading positions + section_matches = list(_SECTION_RE.finditer(text)) + + if not section_matches: + # No headings — fall back to paragraph/sentence splitting + raw_chunks = _paragraph_chunks(text, max_tokens=self.max_tokens) + return [(c, {}) for c in raw_chunks if c] + + # Build (title, body_text) pairs for each section + sections: List[tuple[str, str]] = [] + for i, m in enumerate(section_matches): + title = m.group(1).strip() + body_start = m.end() + body_end = ( + section_matches[i + 1].start() + if i + 1 < len(section_matches) + else len(text) + ) + body = text[body_start:body_end].strip() + sections.append((title, body)) + + # Check for preamble text before the first heading + preamble = text[: section_matches[0].start()].strip() + result: List[tuple[str, Dict[str, Any]]] = [] + + if preamble: + for c in _paragraph_chunks(preamble, max_tokens=self.max_tokens): + if c: + result.append((c, {})) + + for title, body in sections: + section_meta: Dict[str, Any] = {"section": title} + if not body: + # Empty section — emit a placeholder chunk with just the title + result.append((title, section_meta)) + continue + + para_chunks = _paragraph_chunks(body, max_tokens=self.max_tokens) + for c in para_chunks: + if c: + result.append((c, dict(section_meta))) + + return result if result else [(text.strip(), {})] + + +__all__ = ["ChunkResult", "SemanticChunker"] diff --git a/tests/connectors/test_chunker.py b/tests/connectors/test_chunker.py new file mode 100644 index 00000000..086d804a --- /dev/null +++ b/tests/connectors/test_chunker.py @@ -0,0 +1,264 @@ +"""Tests for SemanticChunker — type-aware text splitting.""" + +from __future__ import annotations + +import pytest + +from openjarvis.connectors.chunker import ChunkResult, SemanticChunker + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture() +def chunker() -> SemanticChunker: + return SemanticChunker(max_tokens=50) + + +# --------------------------------------------------------------------------- +# 1. Short message stays as single chunk +# --------------------------------------------------------------------------- + + +def test_short_message_single_chunk(chunker: SemanticChunker) -> None: + """A message shorter than max_tokens is returned as a single chunk.""" + text = "Hello, world! How are you today?" + results = chunker.chunk(text, doc_type="message") + assert len(results) == 1 + assert results[0].content == text + assert results[0].index == 0 + + +# --------------------------------------------------------------------------- +# 2. Long document splits on ## Heading sections, metadata has section key +# --------------------------------------------------------------------------- + + +def test_document_splits_on_headings() -> None: + """Documents with ## headings produce chunks with section metadata.""" + chunker = SemanticChunker(max_tokens=512) + text = ( + "## Introduction\n" + "This is the introduction paragraph. It explains the context.\n\n" + "## Methods\n" + "This section describes the methods used in the study.\n\n" + "## Results\n" + "Here are the results of the experiment." + ) + results = chunker.chunk(text, doc_type="document") + # Each ## section should produce at least one chunk + sections = {r.metadata.get("section") for r in results} + assert "Introduction" in sections + assert "Methods" in sections + assert "Results" in sections + + +# --------------------------------------------------------------------------- +# 3. Within a section, splits on paragraph boundaries +# --------------------------------------------------------------------------- + + +def test_document_splits_on_paragraphs() -> None: + """Within a section, long content splits on double-newline paragraph breaks.""" + # Use a small max_tokens to force splitting within a section + chunker = SemanticChunker(max_tokens=15) + # Build a section with two paragraphs, each > 15 tokens + para1 = " ".join(["word"] * 20) # 20 tokens + para2 = " ".join(["text"] * 20) # 20 tokens + text = f"## Section One\n{para1}\n\n{para2}" + results = chunker.chunk(text, doc_type="document") + # Both paragraphs should be separate chunks + assert len(results) >= 2 + # All chunks from this section carry section metadata + for r in results: + assert r.metadata.get("section") == "Section One" + + +# --------------------------------------------------------------------------- +# 4. Never splits mid-sentence (chunks end with . ? or ! except possibly last) +# --------------------------------------------------------------------------- + + +def test_no_mid_sentence_splits() -> None: + """Chunks (except possibly the last) must end with sentence-ending punctuation.""" + chunker = SemanticChunker(max_tokens=10) + # Build text with clearly delimited sentences in a document section + text = ( + "## Analysis\n" + "The first result was positive. The second outcome was negative. " + "The third finding was inconclusive. The final conclusion is pending." + ) + results = chunker.chunk(text, doc_type="document") + # All chunks except possibly the last should end with sentence punctuation + for r in results[:-1]: + stripped = r.content.rstrip() + assert stripped[-1] in {".", "?", "!"}, ( + f"Non-final chunk does not end with sentence punctuation: {stripped!r}" + ) + + +# --------------------------------------------------------------------------- +# 5. Email thread splits on reply boundaries +# --------------------------------------------------------------------------- + + +def test_email_splits_on_reply_boundaries() -> None: + """Emails split on 'On ... wrote:' reply headers.""" + chunker = SemanticChunker(max_tokens=512) + text = ( + "Hi Alice, please see my comments below.\n\n" + "On Mon, Jan 1, 2024, Alice Smith wrote:\n" + "> Original message here.\n" + "> More original text.\n\n" + "On Sun, Dec 31, 2023, Bob Jones wrote:\n" + "> Even earlier message content." + ) + results = chunker.chunk(text, doc_type="email") + # Should produce more than one chunk due to reply boundaries + assert len(results) >= 2 + + +# --------------------------------------------------------------------------- +# 6. Event stays as single chunk +# --------------------------------------------------------------------------- + + +def test_event_always_single_chunk() -> None: + """Events are never split regardless of length.""" + chunker = SemanticChunker(max_tokens=5) + text = " ".join(["word"] * 100) # 100 tokens, well above max_tokens=5 + results = chunker.chunk(text, doc_type="event") + assert len(results) == 1 + assert results[0].content == text + + +# --------------------------------------------------------------------------- +# 7. Contact stays as single chunk +# --------------------------------------------------------------------------- + + +def test_contact_always_single_chunk() -> None: + """Contacts are never split regardless of length.""" + chunker = SemanticChunker(max_tokens=5) + text = " ".join(["info"] * 100) # 100 tokens, well above max_tokens=5 + results = chunker.chunk(text, doc_type="contact") + assert len(results) == 1 + assert results[0].content == text + + +# --------------------------------------------------------------------------- +# 8. Parent metadata inherited to all chunks +# --------------------------------------------------------------------------- + + +def test_parent_metadata_inherited() -> None: + """All chunks carry the parent metadata passed to chunk().""" + chunker = SemanticChunker(max_tokens=10) + parent_meta = {"source": "gmail", "doc_id": "abc-123", "priority": "high"} + text = ( + "## Section A\n" + "First sentence of section A. Second sentence of section A. " + "Third sentence here. Fourth sentence concludes.\n\n" + "## Section B\n" + "First sentence of section B. Second sentence of section B." + ) + results = chunker.chunk(text, doc_type="document", metadata=parent_meta) + for r in results: + assert r.metadata.get("source") == "gmail" + assert r.metadata.get("doc_id") == "abc-123" + assert r.metadata.get("priority") == "high" + + +# --------------------------------------------------------------------------- +# 9. Chunks have sequential 0-based indexes +# --------------------------------------------------------------------------- + + +def test_sequential_chunk_indexes() -> None: + """Chunks are indexed sequentially from 0 across all splits.""" + chunker = SemanticChunker(max_tokens=10) + text = ( + "## Alpha\n" + "Sentence one ends here. Sentence two ends here. Sentence three ends here.\n\n" + "## Beta\n" + "Sentence four ends here. Sentence five ends here." + ) + results = chunker.chunk(text, doc_type="document") + for i, r in enumerate(results): + assert r.index == i, f"Expected index {i}, got {r.index}" + + +# --------------------------------------------------------------------------- +# Additional edge cases +# --------------------------------------------------------------------------- + + +def test_note_doc_type_treated_as_document() -> None: + """doc_type='note' uses the same document splitting strategy.""" + chunker = SemanticChunker(max_tokens=512) + text = ( + "## My Note\n" + "This is a note with section headings.\n\n" + "## Another Section\n" + "More content here." + ) + results = chunker.chunk(text, doc_type="note") + sections = {r.metadata.get("section") for r in results} + assert "My Note" in sections + assert "Another Section" in sections + + +def test_unknown_doc_type_treated_as_document() -> None: + """Unknown doc_type uses the document splitting strategy.""" + chunker = SemanticChunker(max_tokens=512) + text = "## Header\nContent under the header.\n\n" + results = chunker.chunk(text, doc_type="unknown_type") + sections = {r.metadata.get("section") for r in results} + assert "Header" in sections + + +def test_empty_text_returns_empty_list() -> None: + """Empty string input returns an empty list.""" + chunker = SemanticChunker(max_tokens=512) + results = chunker.chunk("", doc_type="document") + assert results == [] + + +def test_whitespace_only_text_returns_empty_list() -> None: + """Whitespace-only input returns an empty list.""" + chunker = SemanticChunker(max_tokens=512) + results = chunker.chunk(" \n\n \t ", doc_type="document") + assert results == [] + + +def test_chunk_result_is_dataclass() -> None: + """ChunkResult has the expected fields with correct defaults.""" + cr = ChunkResult(content="hello") + assert cr.content == "hello" + assert cr.index == 0 + assert cr.metadata == {} + + +def test_message_accumulates_into_max_tokens() -> None: + """Message chunks accumulate paragraphs up to max_tokens.""" + chunker = SemanticChunker(max_tokens=20) + # Each paragraph is 8 tokens; two fit in 20 but three would exceed 20 (8+8+8=24) + para = "one two three four five six seven eight" # 8 tokens + text = f"{para}\n\n{para}\n\n{para}\n\n{para}" + results = chunker.chunk(text, doc_type="message") + # Should not fit all 4 paragraphs in one chunk (would be 32 tokens) + assert len(results) >= 2 + # Each chunk should be within or just at max_tokens (greedy accumulation) + for r in results[:-1]: + assert len(r.content.split()) <= 20 * 2 # some flexibility for joining + + +def test_document_no_headings_uses_paragraphs() -> None: + """Documents without ## headings fall back to paragraph splitting.""" + chunker = SemanticChunker(max_tokens=15) + para1 = " ".join(["alpha"] * 20) # 20 tokens, forces split + para2 = " ".join(["beta"] * 20) + text = f"{para1}\n\n{para2}" + results = chunker.chunk(text, doc_type="document") + assert len(results) >= 2