diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 000000000..76a03635a --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,9 @@ +# Changelog + +All notable changes to this project are documented here. + +## Unreleased + +### Added + +- Meaning-unit embeddings chunking for `/v1/batch/embeddings`: header, paragraph, sentence, and `data:image` units keep source offsets so naruon can search SKU lines and senders without mixing them into a due-date vector. The naruon one-vector-per-input reduce is unchanged; read `meaning_units` for unit-level search. diff --git a/README.md b/README.md index 65f57dd4c..3319cae86 100644 --- a/README.md +++ b/README.md @@ -255,6 +255,7 @@ python tests/test_paper_contracts.py python tests/test_admin_contract.py python tests/test_conventions.py python tests/test_api_contract.py +python tests/test_meaning_unit_chunking.py python tests/test_security_hardening.py python tests/test_repository_security_metadata.py python tests/test_product_planning_contract.py diff --git a/contextual_orchestrator/api_contract.py b/contextual_orchestrator/api_contract.py index fae9fba0b..84a6746e4 100644 --- a/contextual_orchestrator/api_contract.py +++ b/contextual_orchestrator/api_contract.py @@ -416,7 +416,7 @@ "/v1/batch/embeddings": { "post": { "operationId": "create_batch_embeddings_job", - "summary": "Submit a bulk, latency-tolerant embeddings batch (token-split, routed via pg-llm-batch, cost-recorded)", + "summary": "Submit a bulk, latency-tolerant embeddings batch (meaning-unit split, routed via pg-llm-batch, cost-recorded)", "security": [{"inference_bearer_auth": []}], "requestBody": { "required": True, diff --git a/contextual_orchestrator/batch_routing.py b/contextual_orchestrator/batch_routing.py index d07a48d25..52ab06280 100644 --- a/contextual_orchestrator/batch_routing.py +++ b/contextual_orchestrator/batch_routing.py @@ -407,6 +407,9 @@ class EmbeddingBatchRequest: part_index: int = 0 part_count: int = 1 token_count: int = 0 + source_start: int = 0 + source_end: int = 0 + unit_kind: str = "paragraph_unit" def to_jsonl_line(self, endpoint: str = "/v1/embeddings") -> Dict[str, Any]: """Render this request as an OpenAI Batch API embeddings JSONL line.""" diff --git a/contextual_orchestrator/cost_router.py b/contextual_orchestrator/cost_router.py index bfbe159db..15934b9fc 100644 --- a/contextual_orchestrator/cost_router.py +++ b/contextual_orchestrator/cost_router.py @@ -18,7 +18,6 @@ from __future__ import annotations -import re from typing import Any, Dict, List, Optional from .batch_routing import ( @@ -36,12 +35,12 @@ ) from .cost_ledger import CostLedger, PriceBook from .kv_config import InMemoryConfigStore +from .meaning_unit_chunking import MeaningUnitChunk, split_meaning_units from .token_counting import HeuristicTokenCounter, build_token_counter _EMBEDDING_CONFIG_CATEGORY = "routing" _DEFAULT_EMBEDDING_MAX_TOKENS_PER_REQUEST = 280_000 _DEFAULT_EMBEDDING_MAX_CHARS_PER_PART = 240_000 -_EMBEDDING_UNIT_RE = re.compile(r"\S+\s*|\s+", re.UNICODE) class CostRoutingCoordinator: @@ -301,7 +300,7 @@ def _build_embedding_requests( model: str, attribution: Dict[str, Any], ) -> tuple[List[EmbeddingBatchRequest], List[int], Dict[str, int]]: - """Map original embedding inputs into token-budgeted provider parts.""" + """Map original embedding inputs into meaning-unit provider parts.""" max_tokens, max_chars = self._embedding_request_limits() requests: List[EmbeddingBatchRequest] = [] part_counts: List[int] = [] @@ -312,16 +311,19 @@ def _build_embedding_requests( ) part_count = len(parts) part_counts.append(part_count) - for part_index, (part_text, token_count) in enumerate(parts): + for part_index, chunk in enumerate(parts): requests.append( EmbeddingBatchRequest( - input_text=part_text, + input_text=chunk.chunk_text, model=model, attribution=dict(attribution), source_index=source_index, part_index=part_index, part_count=part_count, - token_count=token_count, + token_count=chunk.token_count, + source_start=chunk.source_start, + source_end=chunk.source_end, + unit_kind=chunk.unit_kind, ) ) return requests, part_counts, { @@ -362,87 +364,14 @@ def _split_embedding_input( model: str, max_tokens: int, max_chars: int, - ) -> List[tuple[str, int]]: - """Split one original embedding input into provider-safe map parts.""" - if text == "": - return [("", 0)] - parts = self._force_token_safe_chunks( - text, model=model, max_tokens=max_tokens, max_chars=max_chars - ) - return parts or [("", 0)] - - def _force_token_safe_chunks( - self, - text: str, - *, - model: str, - max_tokens: int, - max_chars: int, - ) -> List[tuple[str, int]]: - """Recursively split text until each chunk fits token and char budgets.""" - if text == "": - return [("", 0)] - if len(text) > max_chars: - chunks: List[tuple[str, int]] = [] - for start in range(0, len(text), max_chars): - chunks.extend( - self._force_token_safe_chunks( - text[start : start + max_chars], - model=model, - max_tokens=max_tokens, - max_chars=max_chars, - ) - ) - return chunks - - token_count = self._count_embedding_tokens(text, model) - if token_count <= max_tokens or len(text) <= 1: - return [(text, token_count)] - - units = _EMBEDDING_UNIT_RE.findall(text) - if len(units) > 1: - chunks = [] - current = "" - for unit in units: - candidate = f"{current}{unit}" - if current and ( - len(candidate) > max_chars - or self._count_embedding_tokens(candidate, model) > max_tokens - ): - chunks.extend( - self._force_token_safe_chunks( - current, - model=model, - max_tokens=max_tokens, - max_chars=max_chars, - ) - ) - current = unit - else: - current = candidate - if current: - chunks.extend( - self._force_token_safe_chunks( - current, - model=model, - max_tokens=max_tokens, - max_chars=max_chars, - ) - ) - if len(chunks) > 1 or (chunks and chunks[0][0] != text): - return chunks - - midpoint = max(1, len(text) // 2) - return self._force_token_safe_chunks( - text[:midpoint], - model=model, - max_tokens=max_tokens, - max_chars=max_chars, - ) + self._force_token_safe_chunks( - text[midpoint:], + ) -> List[MeaningUnitChunk]: + """Split one original embedding input into meaning-unit map parts.""" + return split_meaning_units( + text, model=model, max_tokens=max_tokens, max_chars=max_chars, + count_tokens=self._count_embedding_tokens, ) def _count_embedding_tokens(self, text: str, model: str) -> int: @@ -501,10 +430,15 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: "prompt_tokens": max(0, prompt_tokens), "model": item.model, "attribution": dict(request.attribution) if request else {}, + "chunk_text": request.input_text if request else "", + "source_start": request.source_start if request else 0, + "source_end": request.source_end if request else 0, + "unit_kind": request.unit_kind if request else "paragraph_unit", } ) embeddings: List[Dict[str, Any]] = [] + meaning_units: List[Dict[str, Any]] = [] token_counts: List[int] = [] total_cost_amount = 0.0 currency_code = "USD" @@ -542,6 +476,18 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: ), } ) + for part in parts: + meaning_units.append( + { + "source_index": source_index, + "part_index": part["part_index"], + "unit_kind": part["unit_kind"], + "source_start": part["source_start"], + "source_end": part["source_end"], + "chunk_text": part["chunk_text"], + "embedding": part["embedding"], + } + ) document = { "batch_id": batch_id, @@ -553,8 +499,10 @@ def embeddings_batch_document(self, batch_id: str) -> Dict[str, Any]: "total_tokens": sum(token_counts), "part_count": len(requests), "input_part_counts": part_counts, + "meaning_units": meaning_units, "map_reduce": { "strategy": "token_budgeted_embedding_parts_weighted_average", + "meaning_unit_strategy": "header_paragraph_sentence_image", **part_limits, }, "cost_amount": round(total_cost_amount, 6), diff --git a/contextual_orchestrator/meaning_unit_chunking.py b/contextual_orchestrator/meaning_unit_chunking.py new file mode 100644 index 000000000..d2d7504e9 --- /dev/null +++ b/contextual_orchestrator/meaning_unit_chunking.py @@ -0,0 +1,390 @@ +"""Meaning-unit chunking for embeddings search. + +Token-midpoint splits mix unrelated facts (invoice due date vs SKU line vs +sender) into one vector. This module walks untrusted text in source order and +emits header blocks, paragraphs, sentences, and embedded ``data:image`` URIs +as separate units, falling back to word then character splits only when a +single unit exceeds the provider budget. + +Grounding: passage-level retrieval (Karpukhin et al., 2020) and late chunking +(Günther et al., 2024). No new runtime dependency — stdlib scanning only. +""" + +from __future__ import annotations + +from dataclasses import dataclass +import re +from typing import Callable, List + +TokenCountFn = Callable[[str, str], int] + +_HEADER_LINE_RE = re.compile(r"^[A-Za-z][A-Za-z0-9-]*:[ \t].+") +_IMAGE_RE = re.compile(r"data:image/[A-Za-z0-9.+-]+;base64,[A-Za-z0-9+/=]+") +_BLANK_LINE_RE = re.compile(r"\n[ \t]*\n") +_SENTENCE_RE = re.compile(r".+?(?:[.!?](?=\s|$)|$)", re.DOTALL) +_WORD_UNIT_RE = re.compile(r"\S+\s*|\s+", re.UNICODE) + +_KIND_HEADER = "header_block" +_KIND_PARAGRAPH = "paragraph_unit" +_KIND_SENTENCE = "sentence_unit" +_KIND_IMAGE = "embedded_image" +_KIND_TOKEN = "token_fallback" + + +@dataclass(frozen=True) +class MeaningUnitChunk: + """One searchable embedding part with its original source span. + + Attributes: + chunk_text: Text forwarded to the embedding provider. + source_start: Inclusive offset into the original input. + source_end: Exclusive offset into the original input. + unit_kind: Meaning-unit class (header, paragraph, sentence, image, or fallback). + token_count: Token estimate used for the provider budget. + """ + + chunk_text: str + source_start: int + source_end: int + unit_kind: str + token_count: int + + +def split_meaning_units( + text: str, + *, + model: str, + max_tokens: int, + max_chars: int, + count_tokens: TokenCountFn, +) -> List[MeaningUnitChunk]: + """Split ``text`` into meaning units that each fit the provider budget. + + Args: + text: Original embedding input. May contain email headers, paragraphs, + and ``data:image`` URIs. + model: Embedding model name forwarded to ``count_tokens``. + max_tokens: Inclusive token ceiling per emitted chunk. + max_chars: Inclusive character ceiling per emitted chunk. + count_tokens: ``(text, model) -> int`` token estimator. + + Returns: + Source-ordered chunks. An empty input yields one empty paragraph unit + so callers can keep a 1:1 source row. + """ + if text == "": + return [ + MeaningUnitChunk( + chunk_text="", + source_start=0, + source_end=0, + unit_kind=_KIND_PARAGRAPH, + token_count=0, + ) + ] + safe_tokens = max(1, int(max_tokens)) + safe_chars = max(1, int(max_chars)) + chunks: List[MeaningUnitChunk] = [] + for start, end, kind in _scan_top_level_units(text): + chunks.extend( + _fit_unit( + text, + start, + end, + kind, + model=model, + max_tokens=safe_tokens, + max_chars=safe_chars, + count_tokens=count_tokens, + ) + ) + return chunks or [ + MeaningUnitChunk( + chunk_text=text, + source_start=0, + source_end=len(text), + unit_kind=_KIND_PARAGRAPH, + token_count=_safe_count(count_tokens, text, model), + ) + ] + + +def _scan_top_level_units(text: str) -> List[tuple[int, int, str]]: + """Return ``(start, end, kind)`` spans that cover ``text`` in order.""" + spans: List[tuple[int, int, str]] = [] + position = 0 + header_end = _leading_header_end(text) + if header_end > 0: + spans.append((0, header_end, _KIND_HEADER)) + position = header_end + length = len(text) + while position < length: + while position < length and text[position] in "\r\n": + position += 1 + if position >= length: + break + image = _IMAGE_RE.match(text, position) + if image is not None: + spans.append((image.start(), image.end(), _KIND_IMAGE)) + position = image.end() + continue + next_image = _IMAGE_RE.search(text, position) + blank = _BLANK_LINE_RE.search(text, position) + end = length + if blank is not None: + end = min(end, blank.start()) + if next_image is not None: + end = min(end, next_image.start()) + if end <= position: + end = min(length, position + 1) + spans.append((position, end, _KIND_PARAGRAPH)) + position = end + return spans + + +def _leading_header_end(text: str) -> int: + """Return the exclusive end of a leading RFC822-style header block.""" + if not _HEADER_LINE_RE.match(text): + return 0 + position = 0 + length = len(text) + while position < length: + line_end = text.find("\n", position) + line = text[position:] if line_end < 0 else text[position:line_end] + if line.endswith("\r"): + line = line[:-1] + if not _HEADER_LINE_RE.match(line): + break + position = length if line_end < 0 else line_end + 1 + return position + + +def _fit_unit( + text: str, + start: int, + end: int, + kind: str, + *, + model: str, + max_tokens: int, + max_chars: int, + count_tokens: TokenCountFn, +) -> List[MeaningUnitChunk]: + """Emit ``text[start:end]`` as one chunk or budget-safe children.""" + unit_text = text[start:end] + if unit_text == "": + return [] + token_count = _safe_count(count_tokens, unit_text, model) + if token_count <= max_tokens and len(unit_text) <= max_chars: + return [ + MeaningUnitChunk( + chunk_text=unit_text, + source_start=start, + source_end=end, + unit_kind=kind, + token_count=token_count, + ) + ] + if kind in {_KIND_PARAGRAPH, _KIND_HEADER}: + sentences = list(_sentence_spans(unit_text, start)) + if len(sentences) > 1: + return _pack_spans( + text, + sentences, + default_kind=_KIND_SENTENCE, + model=model, + max_tokens=max_tokens, + max_chars=max_chars, + count_tokens=count_tokens, + ) + return _fallback_split( + unit_text, + start, + model=model, + max_tokens=max_tokens, + max_chars=max_chars, + count_tokens=count_tokens, + ) + + +def _sentence_spans(unit_text: str, origin: int) -> List[tuple[int, int]]: + """Return sentence spans relative to the original document.""" + spans: List[tuple[int, int]] = [] + for match in _SENTENCE_RE.finditer(unit_text): + piece = match.group(0) + if not piece.strip(): + continue + spans.append((origin + match.start(), origin + match.end())) + return spans + + +def _pack_spans( + text: str, + spans: List[tuple[int, int]], + *, + default_kind: str, + model: str, + max_tokens: int, + max_chars: int, + count_tokens: TokenCountFn, +) -> List[MeaningUnitChunk]: + """Pack adjacent spans until the next one would exceed the budget.""" + chunks: List[MeaningUnitChunk] = [] + pack_start: int | None = None + pack_end: int | None = None + for start, end in spans: + candidate_start = start if pack_start is None else pack_start + candidate = text[candidate_start:end] + token_count = _safe_count(count_tokens, candidate, model) + if pack_start is not None and ( + token_count > max_tokens or len(candidate) > max_chars + ): + packed = text[pack_start:pack_end] + chunks.append( + MeaningUnitChunk( + chunk_text=packed, + source_start=pack_start, + source_end=pack_end or pack_start, + unit_kind=default_kind, + token_count=_safe_count(count_tokens, packed, model), + ) + ) + pack_start = None + pack_end = None + piece = text[start:end] + if ( + _safe_count(count_tokens, piece, model) > max_tokens + or len(piece) > max_chars + ): + chunks.extend( + _fallback_split( + piece, + start, + model=model, + max_tokens=max_tokens, + max_chars=max_chars, + count_tokens=count_tokens, + ) + ) + continue + if pack_start is None: + pack_start = start + pack_end = end + if pack_start is not None and pack_end is not None: + packed = text[pack_start:pack_end] + chunks.append( + MeaningUnitChunk( + chunk_text=packed, + source_start=pack_start, + source_end=pack_end, + unit_kind=default_kind, + token_count=_safe_count(count_tokens, packed, model), + ) + ) + return chunks + + +def _fallback_split( + unit_text: str, + origin: int, + *, + model: str, + max_tokens: int, + max_chars: int, + count_tokens: TokenCountFn, +) -> List[MeaningUnitChunk]: + """Word-pack, then character-split, a single oversized meaning unit.""" + if len(unit_text) > max_chars: + chunks: List[MeaningUnitChunk] = [] + for offset in range(0, len(unit_text), max_chars): + piece = unit_text[offset : offset + max_chars] + chunks.extend( + _fallback_split( + piece, + origin + offset, + model=model, + max_tokens=max_tokens, + max_chars=max_chars, + count_tokens=count_tokens, + ) + ) + return chunks + token_count = _safe_count(count_tokens, unit_text, model) + if token_count <= max_tokens or len(unit_text) <= 1: + kind = _KIND_TOKEN if token_count > max_tokens else _KIND_SENTENCE + return [ + MeaningUnitChunk( + chunk_text=unit_text, + source_start=origin, + source_end=origin + len(unit_text), + unit_kind=kind, + token_count=token_count, + ) + ] + units = _WORD_UNIT_RE.findall(unit_text) + if len(units) > 1: + chunks = [] + current = "" + current_origin = origin + cursor = origin + for unit in units: + candidate = f"{current}{unit}" + if current and ( + len(candidate) > max_chars + or _safe_count(count_tokens, candidate, model) > max_tokens + ): + chunks.append( + MeaningUnitChunk( + chunk_text=current, + source_start=current_origin, + source_end=current_origin + len(current), + unit_kind=_KIND_TOKEN, + token_count=_safe_count(count_tokens, current, model), + ) + ) + current_origin = cursor + current = unit + else: + current = candidate + cursor += len(unit) + if current: + chunks.append( + MeaningUnitChunk( + chunk_text=current, + source_start=current_origin, + source_end=current_origin + len(current), + unit_kind=_KIND_TOKEN, + token_count=_safe_count(count_tokens, current, model), + ) + ) + if len(chunks) > 1 or (chunks and chunks[0].chunk_text != unit_text): + return chunks + midpoint = max(1, len(unit_text) // 2) + left = unit_text[:midpoint] + right = unit_text[midpoint:] + return _fallback_split( + left, + origin, + model=model, + max_tokens=max_tokens, + max_chars=max_chars, + count_tokens=count_tokens, + ) + _fallback_split( + right, + origin + midpoint, + model=model, + max_tokens=max_tokens, + max_chars=max_chars, + count_tokens=count_tokens, + ) + + +def _safe_count(count_tokens: TokenCountFn, text: str, model: str) -> int: + """Count tokens, treating adapter failures as a whitespace word count.""" + try: + value = int(count_tokens(text, model)) + except Exception: + value = len(text.split()) + if text and value <= 0: + return 1 + return max(0, value) diff --git a/docs/architecture.md b/docs/architecture.md index c0f63a81e..98e25aec9 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -37,6 +37,7 @@ This repository implements the interface and control plane, not the trained coor - `WorkflowStep.access`: Conductor-style visibility control. - `ModelClient`: OpenAI-compatible HTTP client, with `mock://` for local checks. - `contextual_orchestrator.server`: small `/v1/chat/completions` HTTP server. +- `contextual_orchestrator.meaning_unit_chunking`: passage-level embeddings split (Karpukhin et al., 2020; Günther et al., 2024) so naruon mail/DOM search keeps sender, paragraph, and embedded-image spans. The deliberate simplification is the policy. The paper systems learn routing and topology from rewards; this lab uses deterministic keyword scoring so the repo runs without training data, GPUs, or vendor credentials. diff --git a/docs/library_research.md b/docs/library_research.md index 42c7fa95c..9a7b090c7 100644 --- a/docs/library_research.md +++ b/docs/library_research.md @@ -53,6 +53,19 @@ Extraction triggers: Until those triggers exist, Ponytail recommends strengthening the current single-repo product instead of splitting it. +## Meaning-unit embeddings chunking + +Researched before adding `contextual_orchestrator/meaning_unit_chunking.py`: + +| Library | Decision | Evidence | +|---|---|---| +| [LangChain RecursiveCharacterTextSplitter](https://python.langchain.com/docs/how_to/recursive_text_splitter/) | Skip | Adds a runtime dependency and splits on character separators, not source-offset meaning units. | +| [LlamaIndex SentenceSplitter](https://docs.llamaindex.ai/en/stable/module_guides/supporting_modules/node_parser_modules/) | Skip | Pulls a document-index stack this gateway does not run. | +| [Chonkie](https://github.com/chonkie-inc/chonkie) | Skip | Extra package for a scan this stdlib regex already covers. | +| stdlib `re` + `HeuristicTokenCounter` | Select | Header/paragraph/sentence/image scan with source offsets, no new lockfile entry. | + +Skipped: a learned late-chunking encoder, OCR/object tags on `data:image` URIs, and a second embeddings store. Those wait for a naruon consumer that persists `meaning_units` and a licensed vision adapter. + ## Required For New Designs Every new subsystem design must update this file before implementation starts. The entry must name the existing libraries researched, the selected library or stdlib alternative, and the custom code that was deliberately skipped. diff --git a/docs/papers/README.md b/docs/papers/README.md index 65a89d2af..c86795062 100644 --- a/docs/papers/README.md +++ b/docs/papers/README.md @@ -35,6 +35,27 @@ redistribution; each is cited below with its arXiv identifier. the responsive path. Distributed under the arXiv non-exclusive license / CC BY as marked on arXiv. +## Meaning-unit embeddings (searchable chunks) + +- **Dense Passage Retrieval for Open-Domain Question Answering** — Vladimir + Karpukhin, Barlas Oguz, Sewon Min, Patrick Lewis, Ledell Wu, Sergey Edunov, + Danqi Chen, Wen-tau Yih. (2020). *Proceedings of the 2020 Conference on + Empirical Methods in Natural Language Processing (EMNLP)*, 6769–6781. + https://doi.org/10.18653/v1/2020.emnlp-main.550 + `dpr-passage-retrieval-2004.04906.pdf` + Grounds **passage-sized retrieval units**: a due-date sentence and a SKU + packing list must not share one vector if a buyer is going to search either + fact. arXiv:2004.04906; distributed under the arXiv non-exclusive license. + +- **Late Chunking: Contextual Chunk Embeddings Using Long-Context Embedding + Models** — Michael Günther, Isabelle Mohr, Daniel James Williams, Bo Wang, + Han Xiao. (2024). arXiv:2409.04701. https://doi.org/10.48550/arXiv.2409.04701 + `late-chunking-2409.04701.pdf` + Grounds **chunk-then-embed with source spans**: keep header, paragraph, + sentence, and embedded-image units, then reduce only for the naruon + one-vector-per-input contract. Distributed under the arXiv non-exclusive + license. + ## Batch execution / load balancing The external `pg-llm-batch` service carries its own grounding papers, including diff --git a/docs/papers/dpr-passage-retrieval-2004.04906.pdf b/docs/papers/dpr-passage-retrieval-2004.04906.pdf new file mode 100644 index 000000000..d751cc541 Binary files /dev/null and b/docs/papers/dpr-passage-retrieval-2004.04906.pdf differ diff --git a/docs/papers/late-chunking-2409.04701.pdf b/docs/papers/late-chunking-2409.04701.pdf new file mode 100644 index 000000000..84126371d Binary files /dev/null and b/docs/papers/late-chunking-2409.04701.pdf differ diff --git a/docs/rest_api_design.md b/docs/rest_api_design.md index 9378e5a37..79a61c43a 100644 --- a/docs/rest_api_design.md +++ b/docs/rest_api_design.md @@ -15,7 +15,7 @@ |---|---|---| | `GET` | `/openapi.json` | API contract | | `POST` | `/v1/chat/completions` | Compatibility chat endpoint | -| `POST` | `/v1/batch/embeddings` | Submit a bulk, latency-tolerant embeddings batch; oversized inputs are token-split before routing via pg-llm-batch | +| `POST` | `/v1/batch/embeddings` | Submit a bulk, latency-tolerant embeddings batch; oversized inputs are split on header/paragraph/sentence/image meaning units before routing via pg-llm-batch. Next action: read `meaning_units` when you need SKU-level or sender-level search, and keep `embeddings` for the one-vector-per-input naruon reduce. | | `GET` | `/v1/batch/embeddings/{batch_id}` | Poll an embeddings batch; returns reduced vectors + recorded cost once completed | | `GET` | `/api/v1/agent_pools` | List model agents | | `GET` | `/api/v1/orchestration_policies/default_policy` | Read active policy | diff --git a/fuzz/targets.py b/fuzz/targets.py index d0c344462..3208e2936 100644 --- a/fuzz/targets.py +++ b/fuzz/targets.py @@ -28,6 +28,7 @@ from typing import Any from contextual_orchestrator import server +from contextual_orchestrator.meaning_unit_chunking import split_meaning_units from contextual_orchestrator.orchestrator import ( ModelAgent, TaskOrchestrator, @@ -188,3 +189,33 @@ def exercise_orchestration(prompt: str, mode: str) -> None: continue assert frame.startswith("data: ") json.loads(frame[len("data: "):]) + + +def exercise_meaning_unit_chunking(text: str) -> None: + """Drive meaning-unit splitting over arbitrary untrusted text. + + Invariants: every successful result is a list of chunks whose source + spans are ordered, non-empty unless the input is empty, and reconstruct + to a substring of the original input. + """ + chunks = split_meaning_units( + text, + model="fuzz-embed", + max_tokens=8, + max_chars=64, + count_tokens=lambda piece, _model: max(1, len(piece.split()) or (1 if piece else 0)), + ) + assert isinstance(chunks, list) + previous_end = 0 + for chunk in chunks: + assert isinstance(chunk.chunk_text, str) + assert isinstance(chunk.unit_kind, str) and chunk.unit_kind + assert chunk.source_start >= 0 + assert chunk.source_end >= chunk.source_start + assert chunk.source_start >= previous_end or chunk.chunk_text == "" + assert text[chunk.source_start : chunk.source_end] == chunk.chunk_text + assert chunk.token_count >= 0 + previous_end = chunk.source_end + if text == "": + assert len(chunks) == 1 + assert chunks[0].chunk_text == "" diff --git a/tests/fuzz/test_fuzz_properties.py b/tests/fuzz/test_fuzz_properties.py index 7e7b3f347..2c4eec5d3 100644 --- a/tests/fuzz/test_fuzz_properties.py +++ b/tests/fuzz/test_fuzz_properties.py @@ -18,6 +18,7 @@ from fuzz.targets import ( exercise_agent_config, + exercise_meaning_unit_chunking, exercise_orchestration, exercise_redaction, exercise_request_body, @@ -108,3 +109,9 @@ def test_redaction_never_crashes_and_is_idempotent(text: str) -> None: ) def test_orchestration_on_arbitrary_prompt(prompt: str, mode: str) -> None: exercise_orchestration(prompt, mode) + + +@_SETTINGS +@given(st.text(max_size=2048)) +def test_meaning_unit_chunking_never_crashes(text: str) -> None: + exercise_meaning_unit_chunking(text) diff --git a/tests/test_meaning_unit_chunking.py b/tests/test_meaning_unit_chunking.py new file mode 100644 index 000000000..c11b65a0f --- /dev/null +++ b/tests/test_meaning_unit_chunking.py @@ -0,0 +1,209 @@ +"""Meaning-unit embeddings chunking: real invoice/email search accuracy. + +Buyers embed AP mail and naruon DOM excerpts to retrieve *one* fact +(due date vs SKU vs sender). Token-midpoint splits mix those facts into one +vector. These tests use a real accounts-payable email and assert the splitter +keeps each meaning unit searchable on its own. +""" + +from __future__ import annotations + +from pathlib import Path +import sys + +sys.path.insert(0, str(Path(__file__).resolve().parents[1])) + +from contextual_orchestrator import ( # noqa: E402 + CostRoutingCoordinator, + InMemoryConfigStore, + ModelAgent, + TaskOrchestrator, +) +from contextual_orchestrator.batch_routing import ( # noqa: E402 + BatchJob, + EmbeddingBatchRequest, + EmbeddingBatchResultItem, +) +from contextual_orchestrator.meaning_unit_chunking import ( # noqa: E402 + MeaningUnitChunk, + split_meaning_units, +) +from contextual_orchestrator.token_counting import HeuristicTokenCounter # noqa: E402 + +# One-pixel PNG so the image unit is a real data URI, not a placeholder. +_PACKING_SLIP_PNG = ( + "data:image/png;base64," + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==" +) + +ACCOUNTS_PAYABLE_EMAIL = ( + "From: billing@acme.example\n" + "To: ap@buyer.example\n" + "\n" + "Invoice 1042 is due on 15 March 2026. Remit to Acme Treasury only.\n" + "\n" + "The packing list names SKU-77 and SKU-88. Do not pay SKU-99.\n" + "\n" + f"{_PACKING_SLIP_PNG}\n" +) + + +class _RecordingEmbeddingBackend: + """Records mapped embedding parts so tests can inspect meaning units.""" + + name = "recording" + + def __init__(self) -> None: + self.requests: list[EmbeddingBatchRequest] = [] + + def submit(self, requests, metadata=None): + self.requests = list(requests) + results = [ + EmbeddingBatchResultItem( + custom_id=request.custom_id, + index=position, + embedding=[float(request.source_index), float(request.part_index)], + prompt_tokens=request.token_count, + model=request.model, + ) + for position, request in enumerate(self.requests) + ] + self._results = results + return BatchJob( + job_id="meaning-unit-embeddings", + backend=self.name, + status="completed", + request_count=len(self.requests), + ) + + def poll(self, job): + return {"job_id": job.job_id, "status": "completed", "is_complete": True} + + def retrieve(self, job): + return list(self._results) + + +def _count_tokens(text: str, model: str = "") -> int: + return HeuristicTokenCounter(tokens_per_word=1.0).count_text(text, model) + + +def test_invoice_paragraphs_stay_separable_for_sku_search() -> None: + """A due-date paragraph must not share a vector with the SKU packing list.""" + chunks = split_meaning_units( + ACCOUNTS_PAYABLE_EMAIL, + model="text-embedding-test", + max_tokens=16, + max_chars=240_000, + count_tokens=_count_tokens, + ) + texts = [chunk.chunk_text for chunk in chunks] + due_date_chunks = [text for text in texts if "Invoice 1042" in text] + sku_chunks = [text for text in texts if "SKU-77" in text] + assert due_date_chunks, texts + assert sku_chunks, texts + assert all("SKU-77" not in text for text in due_date_chunks) + assert all("Invoice 1042" not in text for text in sku_chunks) + + +def test_sender_block_is_its_own_meaning_unit() -> None: + """From/To headers must be searchable without the invoice body.""" + chunks = split_meaning_units( + ACCOUNTS_PAYABLE_EMAIL, + model="text-embedding-test", + max_tokens=16, + max_chars=240_000, + count_tokens=_count_tokens, + ) + header_chunks = [chunk for chunk in chunks if chunk.unit_kind == "header_block"] + assert header_chunks + assert "billing@acme.example" in header_chunks[0].chunk_text + assert "Invoice 1042" not in header_chunks[0].chunk_text + + +def test_embedded_image_keeps_source_offsets() -> None: + """A data-URI image stays one unit and points at its original location.""" + chunks = split_meaning_units( + ACCOUNTS_PAYABLE_EMAIL, + model="text-embedding-test", + max_tokens=16, + max_chars=240_000, + count_tokens=_count_tokens, + ) + image_chunks = [chunk for chunk in chunks if chunk.unit_kind == "embedded_image"] + assert len(image_chunks) == 1 + image = image_chunks[0] + assert image.chunk_text.startswith("data:image/") + restored = ACCOUNTS_PAYABLE_EMAIL[image.source_start : image.source_end] + assert restored == image.chunk_text + assert all(isinstance(chunk, MeaningUnitChunk) for chunk in chunks) + assert all(chunk.source_start < chunk.source_end or chunk.chunk_text == "" for chunk in chunks) + + +def test_oversized_sentence_falls_back_without_mixing_neighbors() -> None: + """A single oversized sentence may token-split; neighbors stay intact.""" + text = ( + "Pay SKU-77 now.\n\n" + + ("word " * 40).strip() + + ".\n\n" + + "Invoice 1042 remains open." + ) + chunks = split_meaning_units( + text, + model="text-embedding-test", + max_tokens=8, + max_chars=240_000, + count_tokens=_count_tokens, + ) + sku = next(chunk for chunk in chunks if "SKU-77" in chunk.chunk_text) + invoice = next(chunk for chunk in chunks if "Invoice 1042" in chunk.chunk_text) + assert "Invoice 1042" not in sku.chunk_text + assert "SKU-77" not in invoice.chunk_text + + +def test_batch_embeddings_preserve_meaning_units_before_backend() -> None: + """The batch path must forward meaning units, not word-packed mixes.""" + orchestrator = TaskOrchestrator( + [ModelAgent(id="mock_worker", model="mock-a", base_url="mock://a", tags=("reasoning",))] + ) + config = InMemoryConfigStore() + config.set("routing", "embedding_max_tokens_per_request", 16) + backend = _RecordingEmbeddingBackend() + coordinator = CostRoutingCoordinator( + orchestrator, + config, + token_counter=HeuristicTokenCounter(tokens_per_word=1.0), + embedding_batch_backend=backend, + ) + document = coordinator.complete_embeddings_batch( + [ACCOUNTS_PAYABLE_EMAIL], + model="text-embedding-test", + attribution={"provider": "acme-provider"}, + ) + texts = [request.input_text for request in backend.requests] + due_date = [text for text in texts if "Invoice 1042" in text] + sku = [text for text in texts if "SKU-77" in text] + assert due_date and sku + assert all("SKU-77" not in text for text in due_date) + assert all(getattr(request, "unit_kind", "") for request in backend.requests) + assert document["input_part_counts"][0] == len(backend.requests) + + +def test_paper_docs_cite_passage_and_late_chunking() -> None: + """Doctoring must name the retrieval papers that justify meaning units.""" + papers = (Path(__file__).resolve().parents[1] / "docs" / "papers" / "README.md").read_text( + encoding="utf-8" + ) + assert "Karpukhin" in papers + assert "2004.04906" in papers + assert "Late chunking" in papers or "Late Chunking" in papers + assert "2409.04701" in papers + + +if __name__ == "__main__": + test_invoice_paragraphs_stay_separable_for_sku_search() + test_sender_block_is_its_own_meaning_unit() + test_embedded_image_keeps_source_offsets() + test_oversized_sentence_falls_back_without_mixing_neighbors() + test_batch_embeddings_preserve_meaning_units_before_backend() + test_paper_docs_cite_passage_and_late_chunking() + print("ok")