diff --git a/.env.example b/.env.example index d1fb4f4..bc1726a 100644 --- a/.env.example +++ b/.env.example @@ -48,5 +48,11 @@ OPENAI_VISION_MODEL= # --- chunk settings --- CHUNK_SIZE= CHUNK_OVERLAP= +# Character-count ceiling (proxy for a token limit — no tokenizer in this +# project) for text/PDF content sent to the LLM-based context-aware chunker +# (POST /api/v1/ingest, semantic_boundaries/contextualize flags). Above this, +# ingestion falls back to the plain chunk_text strategy. Leave empty for the +# default (24000). +CONTEXT_AWARE_CHUNKING_MAX_CHARS= diff --git a/.gitignore b/.gitignore index 8debb71..60ca2b9 100644 --- a/.gitignore +++ b/.gitignore @@ -220,4 +220,6 @@ __marimo__/ .streamlit/secrets.toml -data/ \ No newline at end of file +data/ + +opencode.json \ No newline at end of file diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000..d642615 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,94 @@ +# SprintStart AI Service + +FastAPI service that ingests project artifacts, indexes them via RAG, and +answers questions / generates personalized onboarding paths. + +## Setup + +```bash +uv sync +cp .env.example .env # fill in values +uv run python -m src.main # dev server on :8000, docs at /docs +``` + +Docker: `docker-compose up --build` (overrides `OLLAMA_BASE_URL` to +`http://host.docker.internal:11434` automatically). + +Production Docker CMD (Dockerfile): `uvicorn api.app:app --app-dir src`. + +## Commands + +| Purpose | Command | +|---|---| +| Run tests | `uv run pytest` | +| Tests + coverage | `uv run pytest --cov=src --cov-report=term-missing` | +| Integration tests (needs Ollama) | `uv run pytest -m integration` | +| Lint | `uv run ruff check .` | +| Format check | `uv run ruff format --check .` | +| Auto-format | `uv run ruff format .` | +| Type-check | `uv run pyright src/` | + +Run **lint → format check → type-check → tests** before considering a change +done. CI enforces this order: +`gitleaks → quality (ruff lint + ruff format check + pyright) → pytest`. + +## Branch protection + +PRs to `main` must come from `dev` or `hotfix/*`. Push to `main`/`dev` +triggers Docker image publish to `ghcr.io/sprintstartproject/sprintstart-ai`. + +## Architecture + +`src/` layout: + +- **`api/`** — FastAPI app (`app.py`), DI (`dependencies.py`), `schemas.py`, + SSE helpers (`sse.py`), and route modules under `routes/`. +- **`agents/`** — `Agent` base class (`base.py`) with a `decision_role` / + `answer_system` prompt pair and `max_steps` cap (default 5). + `OrchestratorAgent` decides whether to delegate to sub-agents or answer + directly. Tools live in `agents/tools/`, registered via `ToolRegistry`. + `ChatOrchestrator` wraps the agent run into SSE events for `/api/v1/chat`. +- **`ingestion/`** — Per-filetype parsers (`text_parser`, `pdf_parser`, + `code_parser`, `image_parser`) behind `parser.py`, then `chunker.py` and + `metadata_store.py`. +- **`rag/`** — `retriever.py`, `hybrid.py` (BM25 + vector hybrid retrieval + with RRF fusion), `citation.py`, `query_expansion.py`, `prompt.py`. +- **`llm/`** — `LLMClient` protocol (`base.py`) with implementations: + `ollama_client.py`, `openai_client.py`, `anthropic_client.py`. + `SplitLLMClient` lets chat and embeddings use different backends + (`LLM_BACKEND` vs `EMBED_BACKEND`). +- **`store/`** — `VectorStore` protocol (`base.py`), `chroma_store.py`. +- **`onboarding/`** — Deterministic staged pipeline (not agentic): + `select → filter → retrieve → synthesize → validate → emit`, yielding + `StageProgress` markers. Blueprints are owned by the backend and passed in + on each request — the service is stateless. + +## Conventions + +- `src` is on `pythonpath` for tests — import as `from agents.base import + Agent`, **not** `from src.agents...`. +- Don't assume Ollama-only: check `llm/base.py`'s `LLMClient` protocol when + touching LLM calls. Backend is configurable per deployment via + `LLM_BACKEND` / `EMBED_BACKEND`. +- New sub-agents should follow the `SynthesisAgent` / `OrchestratorAgent` + pattern (`Agent` base + `ToolRegistry`/`AgentTool`), not a bespoke loop. +- `AGENT_DEBUG=1` logs each agent's reasoning (LLM text + tool calls) to + stderr — useful when debugging agent behavior. +- The onboarding pipeline is intentionally deterministic/staged rather than + agentic: a bad LLM output degrades to a blueprint-only path instead of + breaking. Keep that property when modifying it. +- `data/` is gitignored — don't commit local ChromaDB state or test fixtures + there. + +## Tests + +`tests/` mirrors `src/`. Key utilities in `conftest.py`: + +- `llm_required` / `vision_required` markers skip tests when Ollama is + unreachable or vision model unconfigured. +- `parse_sse_events()` parses SSE streams into `dict` lists. +- `clear_dependency_caches` (autouse fixture) resets `lru_cache` on + `get_llm`, `get_store`, `get_ingestion_metadata_store` before each test. + +Reuse the fakes in `tests/stubs/llm.py` (`StubLLMClient`, `ScriptedLLMClient`) +and `tests/stubs/store.py` (`StubVectorStore`) instead of hand-rolling mocks. \ No newline at end of file diff --git a/README.md b/README.md index 8f44c68..f7ec35a 100644 --- a/README.md +++ b/README.md @@ -63,13 +63,14 @@ The service runs on port `8000`. | `AGENT_DEBUG` | `0` | When set to a truthy value, logs each agent's reasoning step (LLM text and tool calls) to stderr. Disabled by default and for `0`/`false`/`no`/`off`/empty. | | `CHUNK_SIZE` | `512` | Maximum number of characters per chunk | | `CHUNK_OVERLAP` | `64` | Number of characters reused between consecutive chunks to preserve context when splitting large chunks | +| `CONTEXT_AWARE_CHUNKING_MAX_CHARS` | `24000` | Character-count ceiling (a proxy for a token limit — there is no tokenizer in this project) for text/PDF content sent to the LLM-based context-aware chunker. Above this, ingestion falls back to the plain `chunk_text` strategy. See [Context-aware chunking](#context-aware-chunking). | ## API Endpoints | Method | Path | Description | |---|---|---| | `GET` | `/api/v1/health` | Reports service health including LLM backend status. Returns `503` if Ollama is unreachable. | -| `POST` | `/api/v1/ingest` | Parses, chunks, and embeds a document and stores it in the vector store. Re-ingesting the same `artifact_id` replaces existing chunks. Supports text files and images (send image content as base64; requires `OLLAMA_VISION_MODEL`). | +| `POST` | `/api/v1/ingest` | Parses, chunks, and embeds a document and stores it in the vector store. Re-ingesting the same `artifact_id` replaces existing chunks. Supports text files and images (send image content as base64; requires `OLLAMA_VISION_MODEL`). Text/PDF content can optionally be chunked using the LLM-based [context-aware chunker](#context-aware-chunking) (opt-in, off by default). | | `POST` | `/api/v1/chat` | Retrieves relevant chunks and streams a generated answer as Server-Sent Events (SSE). | | `POST` | `/api/v1/title` | Generates a short descriptive title from a user prompt. | | `POST` | `/api/v1/onboarding/path` | Generates a personalized onboarding path (SSE). Blueprints are passed in by the backend on each request — the service is stateless. | @@ -87,6 +88,26 @@ The `/api/v1/chat` endpoint streams newline-delimited JSON events: | `done` | Signals the end of the stream | | `error` | Emitted on failure instead of the above | +## Context-aware chunking + +`POST /api/v1/ingest` can optionally chunk text and PDF content using an LLM instead of the plain character-length `chunk_text` strategy (`ingestion/chunker.py`). This is opt-in per request via two independently toggleable flags on `IngestRequest`, both `false` by default: + +| Field | Default | Description | +|---|---|---| +| `semantic_boundaries` | `false` | Let the LLM choose chunk boundaries based on topic shifts / section boundaries instead of accumulating paragraphs up to `CHUNK_SIZE`. | +| `contextualize` | `false` | Let the LLM selectively prepend a short situating "context block" (Anthropic-style *Contextual Retrieval*) to chunks that would otherwise lack context; self-contained chunks are left untouched. | + +Both flags share a single LLM call (see `ingestion/context_aware_chunker.py`). Leaving both at `false` (the default) skips the LLM call entirely and behaves exactly like the legacy chunker — existing callers are unaffected unless they explicitly opt in. When either is enabled, the strategy still always falls back to `chunk_text` — no request ever fails because of it — when: + +- the content exceeds `CONTEXT_AWARE_CHUNKING_MAX_CHARS`, +- the LLM backend is unreachable, or +- the LLM response is missing, malformed, or fails validation (e.g. out-of-range boundaries). + +Resulting chunks carry extra metadata so context/overlap portions of a chunk's content are identifiable: `has_context_block`, `context_block_range`, `has_overlap`, `overlap_range` (character ranges as `"start:end"` strings, empty when not applicable). + +**Scope:** only text and PDF content are affected; code and image parsing are unchanged. For PDFs, the LLM is invoked **per page** (pages are already parsed independently), so semantic boundaries and context blocks are page-scoped, not whole-document-scoped. The batch endpoint `POST /api/v1/ingest/sync` (GitHub run sync) does not expose these flags and always uses the plain chunker — changing that would mean changing the batch wire contract with the backend, which is a separate decision. + + ## AI-proposed onboarding blueprints Onboarding paths are assembled from versioned **blueprints** scoped by `global` or `area:`. Blueprints carry a `source`: `authored` (human-written) or `generated` (drafted by the AI from the ingested corpus). The backend owns blueprint persistence, versioning, and rollback — the AI service is stateless and only returns generated data. diff --git a/src/api/routes/ingest.py b/src/api/routes/ingest.py index 8b6bc35..eaf0317 100644 --- a/src/api/routes/ingest.py +++ b/src/api/routes/ingest.py @@ -148,7 +148,15 @@ def ingest( metadata_store.mark_failed(body.artifact_id, detail, _utc_now()) raise HTTPException(status_code=413, detail=detail) - parsed_chunks = parse(body.filename, content_bytes) + use_context_aware_chunking = body.semantic_boundaries or body.contextualize + + parsed_chunks = parse( + body.filename, + content_bytes, + llm=llm if use_context_aware_chunking else None, + semantic_boundaries=body.semantic_boundaries, + contextualize=body.contextualize, + ) if not parsed_chunks: try: diff --git a/src/api/schemas.py b/src/api/schemas.py index 189355b..fdc594c 100644 --- a/src/api/schemas.py +++ b/src/api/schemas.py @@ -46,6 +46,29 @@ class IngestRequest(BaseModel): ), examples=["primary"], ) + semantic_boundaries: bool = Field( + default=False, + description=( + "Only affects text and PDF content. When true, an LLM chooses " + "chunk boundaries based on semantic coherence (topic shifts, " + "section boundaries) instead of the default character-length " + "accumulation. Falls back to the default chunker if the " + "content is too large for the LLM or the LLM output is " + "invalid. Independently toggleable from 'contextualize'." + ), + ) + contextualize: bool = Field( + default=False, + description=( + "Only affects text and PDF content. When true, an LLM flags " + "which chunks would benefit from a short situating context " + "block (Anthropic-style Contextual Retrieval) and prepends it " + "to their content; self-contained chunks are left untouched. " + "Falls back to the default chunker if the content is too " + "large for the LLM or the LLM output is invalid. " + "Independently toggleable from 'semantic_boundaries'." + ), + ) @field_validator("filename") @classmethod diff --git a/src/ingestion/chunker.py b/src/ingestion/chunker.py index 259eada..e19631f 100644 --- a/src/ingestion/chunker.py +++ b/src/ingestion/chunker.py @@ -20,6 +20,10 @@ def to_parsed_chunk( filename: str, chunk_index: int, total_chunks_amount: int, + has_context_block: bool = False, + context_block_range: tuple[int, int] | None = None, + has_overlap: bool = False, + overlap_range: tuple[int, int] | None = None, ): """Create a ParsedChunk with standard metadata. @@ -39,6 +43,24 @@ def to_parsed_chunk( total_chunks_amount (int): Total number of chunks produced for the source file. + has_context_block (bool, optional): + Whether a situating context block was prepended to the chunk + content. Defaults to ``False``. + + context_block_range (tuple[int, int] | None, optional): + ``(start, end)`` character range of the context block within + ``chunk_content``, if ``has_context_block`` is ``True``. + Defaults to ``None``. + + has_overlap (bool, optional): + Whether the chunk contains character/paragraph overlap carried + over from a neighboring chunk. Defaults to ``False``. + + overlap_range (tuple[int, int] | None, optional): + ``(start, end)`` character range of the overlap within + ``chunk_content``, if ``has_overlap`` is ``True``. Defaults to + ``None``. + Returns: ParsedChunk: Chunk instance with content, type and metadata. @@ -50,33 +72,60 @@ def to_parsed_chunk( **build_metadata(Path(filename)), "chunk_index": str(chunk_index), "total_chunks": str(total_chunks_amount), + "has_context_block": str(has_context_block).lower(), + "context_block_range": ( + f"{context_block_range[0]}:{context_block_range[1]}" + if context_block_range is not None + else "" + ), + "has_overlap": str(has_overlap).lower(), + "overlap_range": ( + f"{overlap_range[0]}:{overlap_range[1]}" + if overlap_range is not None + else "" + ), }, ) -def chunk_text( - filename: str, - text: str, +def split_into_paragraphs(text: str) -> list[str]: + """Split text into non-empty, whitespace-stripped paragraphs. + + Paragraphs are separated by double newlines (``\\n\\n``). + + Args: + text (str): + Text content to split. + + Returns: + list[str]: + Paragraphs in document order. + """ + return [paragraph.strip() for paragraph in text.split("\n\n") if paragraph.strip()] + + +def group_paragraphs_into_chunks( + paragraphs: list[str], chunk_size: int = chunk_size, chunk_overlap: int = chunk_overlap, -) -> list[ParsedChunk]: - """Split text into paragraph-aware chunks. +) -> list[str]: + """Accumulate paragraphs into character-limited, overlap-aware chunks. - The function preserves paragraph boundaries by splitting on - double newlines (``\\n\\n``). Paragraphs are accumulated until - adding another paragraph would exceed the configured chunk size. + Paragraphs are accumulated until adding another paragraph would exceed + ``chunk_size``. When a chunk boundary is reached, the last paragraph is + carried into the next chunk as overlap context. Paragraphs that exceed + ``chunk_size`` on their own are hard-split into overlapping character + chunks. - When a chunk boundary is reached, the last paragraph is carried - into the next chunk as overlap context. Paragraphs that exceed - ``chunk_size`` on their own are split into overlapping character - chunks by the given chunk_overlap. + This grouping logic is shared by :func:`chunk_text` and the LLM-based + context-aware chunker (when it is asked to only contextualize + already-fixed chunk boundaries rather than choose semantic boundaries + itself). Args: - filename (str): - Name of the source file. - - text (str): - Text content to split. + paragraphs (list[str]): + Paragraphs in document order, e.g. from + :func:`split_into_paragraphs`. chunk_size (int, optional): Maximum chunk size in characters. @@ -87,15 +136,10 @@ def chunk_text( Defaults to the value configured via ``CHUNK_OVERLAP``. Returns: - list[ParsedChunk]: - Paragraph-aware text chunks with metadata. + list[str]: + Raw chunk contents (paragraphs joined by ``\\n\\n``). """ - raw_chunks_content: list[str] = [] - - paragraphs: list[str] = [ - paragraph.strip() for paragraph in text.split("\n\n") if paragraph.strip() - ] current_chunk_content: list[str] = [] for paragraph in paragraphs: @@ -140,6 +184,50 @@ def chunk_text( if current_chunk_content: raw_chunks_content.append("\n\n".join(current_chunk_content)) + return raw_chunks_content + + +def chunk_text( + filename: str, + text: str, + chunk_size: int = chunk_size, + chunk_overlap: int = chunk_overlap, +) -> list[ParsedChunk]: + """Split text into paragraph-aware chunks. + + The function preserves paragraph boundaries by splitting on + double newlines (``\\n\\n``). Paragraphs are accumulated until + adding another paragraph would exceed the configured chunk size. + + When a chunk boundary is reached, the last paragraph is carried + into the next chunk as overlap context. Paragraphs that exceed + ``chunk_size`` on their own are split into overlapping character + chunks by the given chunk_overlap. + + Args: + filename (str): + Name of the source file. + + text (str): + Text content to split. + + chunk_size (int, optional): + Maximum chunk size in characters. + Defaults to the value configured via ``CHUNK_SIZE``. + + chunk_overlap (int, optional): + Overlap used when hard-splitting oversized paragraphs. + Defaults to the value configured via ``CHUNK_OVERLAP``. + + Returns: + list[ParsedChunk]: + Paragraph-aware text chunks with metadata. + """ + paragraphs = split_into_paragraphs(text) + raw_chunks_content = group_paragraphs_into_chunks( + paragraphs, chunk_size, chunk_overlap + ) + total_chunks_amount = len(raw_chunks_content) return [ diff --git a/src/ingestion/context_aware_chunker.py b/src/ingestion/context_aware_chunker.py new file mode 100644 index 0000000..1e2d841 --- /dev/null +++ b/src/ingestion/context_aware_chunker.py @@ -0,0 +1,445 @@ +"""LLM-based context-aware chunking strategy for text content. + +This is an opt-in, additive alternative to the paragraph-boundary +``chunk_text`` strategy in :mod:`ingestion.chunker`. Instead of splitting on +character length alone, an LLM is given the full text (split into +paragraphs) and asked to choose semantically coherent chunk boundaries and, +optionally, a short situating "context block" for chunks that would benefit from one. + +Both behaviors share a single LLM call but are independently toggleable: + +- ``semantic_boundaries``: let the LLM choose paragraph groupings instead of + the character-length accumulation used by ``chunk_text``. +- ``contextualize``: let the LLM flag chunks that need a short context block + and prepend it to their content. + +If either the text is too large for the LLM, or the LLM output cannot be +parsed, the whole call falls back to :func:`ingestion.chunker.chunk_text`. + +The LLM is asked to reason over small integer *indices*, not character +offsets — LLMs are unreliable at counting characters, but reliably reproduce +small integers that are put in front of each segment in the prompt. + +Note on ``semantic_boundaries=False``: the caller has already grouped the +text into its final chunks (e.g. via ``chunk_text``'s paragraph-accumulation +logic) *before* calling this module, and passes those finished chunks in as +``segments``. This way the numbering the LLM sees lines up exactly with the +chunk order index that ``context_blocks`` must key off — there is no +separate "boundaries" step to reconcile. When ``semantic_boundaries=True``, +``segments`` are the raw paragraphs instead, and the LLM's own ``boundaries`` +determine the resulting chunk grouping (and thus what order index each +chunk has). +""" + +import json +import logging +import os + +from pydantic import BaseModel, Field, ValidationError + +from ingestion.chunker import ( + chunk_overlap as default_chunk_overlap, +) +from ingestion.chunker import ( + chunk_size as default_chunk_size, +) +from ingestion.chunker import ( + chunk_text, + group_paragraphs_into_chunks, + split_into_paragraphs, + to_parsed_chunk, +) +from ingestion.models import ParsedChunk +from llm.base import LLMClient, Message +from llm.errors import LLMUnavailableError +from llm.parsing import extract_json_object + +logger = logging.getLogger(__name__) + +try: + # There is no tokenizer in this project (backend is configurable across + # Ollama/OpenAI/Anthropic), so this is a character-count proxy for a + # token-count limit rather than an exact one. ~4 chars/token is a common + # rough estimate for English prose; the default is deliberately + # conservative to leave headroom for the prompt scaffolding itself. + max_chars: int = int(os.getenv("CONTEXT_AWARE_CHUNKING_MAX_CHARS", "24000")) +except ValueError as err: + raise ValueError("CONTEXT_AWARE_CHUNKING_MAX_CHARS must be an integer") from err + + +def exceeds_llm_limit(text: str, max_chars: int = max_chars) -> bool: + """Check whether ``text`` is too large to hand to the LLM as-is. + + This is a cheap character-count proxy for a token-count limit — used to + decide upfront whether to attempt context-aware chunking at all, before + spending an LLM call. Callers should fall back to + :func:`ingestion.chunker.chunk_text` when this returns ``True``. + + Args: + text (str): + The full text that would be sent to the LLM. + + max_chars (int, optional): + Maximum character count considered safe to send. + Defaults to the value configured via + ``CONTEXT_AWARE_CHUNKING_MAX_CHARS``. + + Returns: + bool: + ``True`` if ``text`` exceeds ``max_chars``. + """ + return len(text) > max_chars + + +class ContextAwareChunkingError(Exception): + """Raised when the LLM output cannot be parsed/validated. + + Callers should catch this and fall back to :func:`chunker.chunk_text`. + """ + + +class _Payload(BaseModel): + # Segment indices where a new chunk starts (only meaningful when the LLM + # was asked to choose boundaries itself). Must be sorted, start at 0, and + # contain no duplicates or out-of-range indices. + boundaries: list[int] = Field(default_factory=list[int]) + # Chunk *order* index (0-based, as string) -> short situating sentence. + # Chunks missing from this dict need no context block. + context_blocks: dict[str, str] = Field(default_factory=dict) + + +def _build_prompt( + segments: list[str], + semantic_boundaries: bool, + contextualize: bool, +) -> list[Message]: + """Build the chat prompt for the context-aware chunking LLM call. + + Args: + segments (list[str]): + The text to analyze, in document order. When + ``semantic_boundaries`` is ``True`` these are raw paragraphs + that the LLM groups into chunks itself via ``boundaries``. When + ``semantic_boundaries`` is ``False`` these are already the + caller's finished chunks — the LLM only fills + ``context_blocks`` for them, keyed by their position here. + + semantic_boundaries (bool): + Whether the LLM should choose chunk boundaries itself. When + ``False``, the caller has already fixed the boundaries and the + LLM is only asked to fill ``context_blocks`` for those chunks. + + contextualize (bool): + Whether the LLM should propose context blocks at all. When + ``False``, the LLM is instructed to always leave + ``context_blocks`` empty. + + Returns: + list[Message]: + System + user messages ready to pass to ``LLMClient.generate``. + """ + unit_name = "paragraph" if semantic_boundaries else "chunk" + numbered_segments = "\n\n".join( + f"[{index}] {segment}" for index, segment in enumerate(segments) + ) + + instructions: list[str] = [] + if semantic_boundaries: + instructions.append( + "1. 'boundaries': a sorted list of paragraph indices where a new " + "chunk should start, based on topic shifts, section boundaries, " + "or other semantic breaks. It must start with 0 and contain no " + "duplicates or indices outside the given paragraph range." + ) + else: + instructions.append( + "1. 'boundaries': always return an empty list — chunk " + "boundaries are fixed by the caller and are not part of this " + "task." + ) + + if contextualize: + chunk_reference = ( + "the resulting chunk's order index (0-based, counted over the " + "chunks formed by your own 'boundaries' — NOT the paragraph " + "index)" + if semantic_boundaries + else "the chunk's index as given above" + ) + instructions.append( + "2. 'context_blocks': a mapping from " + chunk_reference + " (as " + "a string) to a short (one to two sentence) situating sentence " + "that gives the chunk context it would otherwise lack (e.g. " + "what document/section it is part of, what it is about). Only " + "include chunks that genuinely benefit from this — " + "self-contained chunks should be omitted from the mapping " + "entirely." + ) + else: + instructions.append( + "2. 'context_blocks': always return an empty object — " + "contextualization is not part of this task." + ) + + system = ( + f"You analyze a document that has been split into numbered " + f"{unit_name}s, in order to prepare it for chunking in a retrieval " + "system.\n\n" + "Return STRICT JSON only (no prose, no markdown fences) with " + "exactly two keys:\n" + "\n".join(instructions) + "\n\n" + 'JSON schema: {"boundaries": [int], "context_blocks": ' + '{"": str}}' + ) + user = f"Numbered {unit_name}s:\n\n{numbered_segments}" + + return [ + Message(role="system", content=system), + Message(role="user", content=user), + ] + + +def _request_chunking_plan( + segments: list[str], + llm: LLMClient, + semantic_boundaries: bool, + contextualize: bool, +) -> _Payload: + """Call the LLM and parse its response into a validated ``_Payload``. + + Raises: + ContextAwareChunkingError: + If the LLM output is not valid/parseable JSON matching the + expected schema. + """ + messages = _build_prompt(segments, semantic_boundaries, contextualize) + raw = llm.generate(messages) + + try: + payload = _Payload.model_validate_json(extract_json_object(raw)) + except (ValidationError, json.JSONDecodeError, ValueError) as exc: + raise ContextAwareChunkingError( + f"invalid context-aware chunking output: {exc}" + ) from exc + + num_segments = len(segments) + if any( + index < 0 or index >= num_segments for index in payload.boundaries + ) or payload.boundaries != sorted(set(payload.boundaries)): + raise ContextAwareChunkingError( + "LLM returned invalid boundaries: " + f"{payload.boundaries!r} for {num_segments} segments" + ) + if payload.boundaries and payload.boundaries[0] != 0: + raise ContextAwareChunkingError( + f"LLM boundaries must start at 0, got {payload.boundaries!r}" + ) + + return payload + + +def _hard_split_with_context( + context_block: str, + body: str, + chunk_size: int, + chunk_overlap: int, +) -> list[tuple[str, tuple[int, int] | None, bool, tuple[int, int] | None]]: + """Split an oversized ``context_block + body`` chunk into sub-chunks. + + The context block is re-prepended to every resulting sub-chunk (it is + small and self-contained, so duplicating it preserves the situating + context each sub-chunk would otherwise lose). The remaining budget + (``chunk_size`` minus the context block) is then hard-split the same + way :func:`ingestion.chunker.group_paragraphs_into_chunks` hard-splits + oversized paragraphs — overlapping character windows of ``chunk_size - + chunk_overlap`` stride. + + Args: + context_block (str): + The situating context block text, already including its + trailing separator (empty string if there is none). + + body (str): + The chunk's own content, without the context block. + + chunk_size (int): + Maximum size of each resulting sub-chunk, in characters, + including the re-prepended context block. + + chunk_overlap (int): + Overlap between consecutive sub-chunk bodies. + + Returns: + list[tuple[str, tuple[int, int] | None, bool, tuple[int, int] | None]]: + One tuple per sub-chunk: ``(content, context_block_range, + has_overlap, overlap_range)``. + """ + results: list[tuple[str, tuple[int, int] | None, bool, tuple[int, int] | None]] = [] + context_len = len(context_block) + body_budget = max(chunk_size - context_len, 1) + body_overlap = min(chunk_overlap, body_budget - 1) if body_budget > 1 else 0 + + start = 0 + first = True + while start < len(body) or first: + body_slice = body[start : start + body_budget] + content = context_block + body_slice + + context_range = (0, context_len) if context_len else None + has_overlap = not first and body_overlap > 0 + overlap_range = ( + (context_len, context_len + body_overlap) if has_overlap else None + ) + + results.append((content, context_range, has_overlap, overlap_range)) + + if start + body_budget >= len(body): + break + start += body_budget - body_overlap + first = False + + return results + + +def chunk_text_context_aware( + filename: str, + text: str, + llm: LLMClient, + semantic_boundaries: bool = True, + contextualize: bool = True, + chunk_size: int = default_chunk_size, + chunk_overlap: int = default_chunk_overlap, +) -> list[ParsedChunk]: + """Split text into chunks using an LLM for semantic boundaries and/or + selective contextualization. + + This is an opt-in alternative to :func:`ingestion.chunker.chunk_text`. + Falls back to it when: + + - both ``semantic_boundaries`` and ``contextualize`` are ``False`` + (nothing for the LLM to do), + - ``text`` exceeds the configured LLM size limit (see + :func:`exceeds_llm_limit`), + - the LLM backend is unreachable, or + - the LLM output cannot be parsed/validated. + + Args: + filename (str): + Name of the source file. + + text (str): + Text content to split. + + llm (LLMClient): + Client used to request the chunking plan. + + semantic_boundaries (bool, optional): + Let the LLM choose chunk boundaries based on semantic coherence + instead of ``chunk_text``'s character-length accumulation. + Defaults to ``True``. + + contextualize (bool, optional): + Let the LLM flag chunks that would benefit from a short + situating context block and prepend it to their content. + Defaults to ``True``. + + chunk_size (int, optional): + Maximum chunk size in characters, including any prepended + context block. Defaults to the value configured via + ``CHUNK_SIZE``. + + chunk_overlap (int, optional): + Overlap used when hard-splitting oversized chunks. + Defaults to the value configured via ``CHUNK_OVERLAP``. + + Returns: + list[ParsedChunk]: + Context-aware text chunks with metadata, or the + ``chunk_text`` fallback result. + """ + if not semantic_boundaries and not contextualize: + return chunk_text(filename, text, chunk_size, chunk_overlap) + + if exceeds_llm_limit(text): + logger.warning( + "Text for %s exceeds context-aware chunking size limit " + "(%d chars) — falling back to chunk_text", + filename, + len(text), + ) + return chunk_text(filename, text, chunk_size, chunk_overlap) + + paragraphs = split_into_paragraphs(text) + if not paragraphs: + return chunk_text(filename, text, chunk_size, chunk_overlap) + + segments = ( + paragraphs + if semantic_boundaries + else group_paragraphs_into_chunks(paragraphs, chunk_size, chunk_overlap) + ) + + try: + payload = _request_chunking_plan( + segments, llm, semantic_boundaries, contextualize + ) + except (ContextAwareChunkingError, LLMUnavailableError) as exc: + logger.warning( + "Context-aware chunking failed for %s (%s) — falling back to chunk_text", + filename, + exc, + ) + return chunk_text(filename, text, chunk_size, chunk_overlap) + + if semantic_boundaries: + boundaries = payload.boundaries or [0] + raw_chunks = [ + "\n\n".join(paragraphs[start:end]) + for start, end in zip( + boundaries, [*boundaries[1:], len(paragraphs)], strict=True + ) + ] + else: + raw_chunks = segments + + final_chunks: list[ + tuple[str, tuple[int, int] | None, bool, tuple[int, int] | None] + ] = [] + for order_index, chunk_body in enumerate(raw_chunks): + context_text = ( + payload.context_blocks.get(str(order_index)) if contextualize else None + ) + context_block = f"{context_text}\n\n" if context_text else "" + combined = context_block + chunk_body + + if len(combined) <= chunk_size: + context_range = (0, len(context_block)) if context_block else None + final_chunks.append((combined, context_range, False, None)) + continue + + final_chunks.extend( + _hard_split_with_context( + context_block, chunk_body, chunk_size, chunk_overlap + ) + ) + + total_chunks_amount = len(final_chunks) + + return [ + to_parsed_chunk( + content, + "text", + filename, + chunk_index, + total_chunks_amount, + has_context_block=context_range is not None, + context_block_range=context_range, + has_overlap=has_overlap, + overlap_range=overlap_range, + ) + for chunk_index, ( + content, + context_range, + has_overlap, + overlap_range, + ) in enumerate(final_chunks) + ] diff --git a/src/ingestion/parser.py b/src/ingestion/parser.py index ac0cb9a..ddd6cf8 100644 --- a/src/ingestion/parser.py +++ b/src/ingestion/parser.py @@ -7,6 +7,7 @@ from ingestion.models import ParsedChunk from ingestion.pdf_parser import parse_pdf from ingestion.text_parser import parse_text +from llm.base import LLMClient logger = logging.getLogger(__name__) @@ -17,7 +18,13 @@ IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"} -def parse(filename: str, content: bytes) -> list[ParsedChunk]: +def parse( + filename: str, + content: bytes, + llm: LLMClient | None = None, + semantic_boundaries: bool = False, + contextualize: bool = False, +) -> list[ParsedChunk]: """Parse a file into structured chunks based on its file extension. The parser acts as a dispatcher and forwards the file content @@ -32,10 +39,30 @@ def parse(filename: str, content: bytes) -> list[ParsedChunk]: There is no "unsupported" type: an unknown extension falls back to the plain-text parser rather than raising. + ``llm``, ``semantic_boundaries``, and ``contextualize`` only affect the + plain-text and PDF paths (see :func:`ingestion.text_parser.parse_text` + and :func:`ingestion.pdf_parser.parse_pdf`); code and image parsing are + unaffected. Passing no ``llm`` preserves the exact prior behavior. + Args: filename (str): Name of the uploaded file including extension. content (bytes): Raw file content as bytes. + llm (LLMClient | None, optional): + Client used for context-aware chunking of text/PDF content. + When ``None`` (the default), plain paragraph/character-based + chunking is used. Defaults to ``None``. + + semantic_boundaries (bool, optional): + Let the LLM choose chunk boundaries based on semantic + coherence. Only relevant when ``llm`` is provided. Defaults to + ``True``. + + contextualize (bool, optional): + Let the LLM prepend a short situating context block to chunks + that would benefit from one. Only relevant when ``llm`` is + provided. Defaults to ``True``. + Returns: list[ParsedChunk]: A list of ParsedChunk objects extracted from the file. """ @@ -48,10 +75,22 @@ def parse(filename: str, content: bytes) -> list[ParsedChunk]: return parse_code(filename, content) elif file_suffix in PDF_EXTENSION: - return parse_pdf(filename, content) + return parse_pdf( + filename, + content, + llm=llm, + semantic_boundaries=semantic_boundaries, + contextualize=contextualize, + ) elif file_suffix in IMAGE_EXTENSIONS: return parse_image(filename, content) else: - return parse_text(filename, content) + return parse_text( + filename, + content, + llm=llm, + semantic_boundaries=semantic_boundaries, + contextualize=contextualize, + ) diff --git a/src/ingestion/pdf_parser.py b/src/ingestion/pdf_parser.py index 3072939..f7466ea 100644 --- a/src/ingestion/pdf_parser.py +++ b/src/ingestion/pdf_parser.py @@ -4,17 +4,27 @@ from pypdf import PdfReader from ingestion.chunker import chunk_text +from ingestion.context_aware_chunker import chunk_text_context_aware from ingestion.models import ParsedChunk +from llm.base import LLMClient logger = logging.getLogger(__name__) -def parse_pdf(filename: str, content: bytes) -> list[ParsedChunk]: +def parse_pdf( + filename: str, + content: bytes, + llm: LLMClient | None = None, + semantic_boundaries: bool = False, + contextualize: bool = False, +) -> list[ParsedChunk]: """ Parse a PDF file into structured text chunks. Each PDF page is processed independently. Extracted text is split - into fixed-size chunks using `chunk_text` (max 512 characters per chunk). + into fixed-size chunks using `chunk_text` (max 512 characters per chunk), + or, when ``llm`` is provided, using the LLM-based context-aware chunker + (see :func:`ingestion.context_aware_chunker.chunk_text_context_aware`). Each resulting chunk is enriched with PDF-specific metadata. @@ -31,6 +41,15 @@ def parse_pdf(filename: str, content: bytes) -> list[ParsedChunk]: - global_pdf_chunk_index: Sequential global index of the chunk across the entire PDF document + Note on context-aware chunking scope: the LLM is invoked once **per + page** (not once for the whole document), since pages are already + processed independently here. This means semantic boundaries and + context blocks are chosen with only a single page's text in view, + not the full document — context blocks are therefore page-scoped + rather than document-scoped. Widening this to whole-document context + (e.g. by threading a document-level summary into each page's call) is + a possible future refinement, not implemented here. + Args: filename (str): Name of the source PDF file. @@ -38,6 +57,21 @@ def parse_pdf(filename: str, content: bytes) -> list[ParsedChunk]: content (bytes): Raw binary content of the uploaded PDF file. + llm (LLMClient | None, optional): + Client used for context-aware chunking. When ``None`` (the + default), plain paragraph/character-based chunking is used for + every page. Defaults to ``None``. + + semantic_boundaries (bool, optional): + Let the LLM choose chunk boundaries per page based on semantic + coherence. Only relevant when ``llm`` is provided. Defaults to + ``False``. + + contextualize (bool, optional): + Let the LLM prepend a short situating context block to chunks + that would benefit from one. Only relevant when ``llm`` is + provided. Defaults to ``False``. + Returns: list[ParsedChunk]: A flat list of all extracted and chunked PDF content, @@ -68,7 +102,18 @@ def parse_pdf(filename: str, content: bytes) -> list[ParsedChunk]: logger.warning("Skipping empty page %s in %s", page_number, filename) continue - chunks: list[ParsedChunk] = chunk_text(filename, text) + use_context_aware_chunking = semantic_boundaries or contextualize + chunks: list[ParsedChunk] = ( + chunk_text(filename, text) + if not use_context_aware_chunking or llm is None + else chunk_text_context_aware( + filename, + text, + llm, + semantic_boundaries=semantic_boundaries, + contextualize=contextualize, + ) + ) for chunk in chunks: chunk.kind = "pdf" chunk.metadata["global_pdf_chunk_index"] = str(global_pdf_chunk_index) diff --git a/src/ingestion/text_parser.py b/src/ingestion/text_parser.py index fdef2f5..54c8efb 100644 --- a/src/ingestion/text_parser.py +++ b/src/ingestion/text_parser.py @@ -1,17 +1,57 @@ from ingestion.chunker import chunk_text +from ingestion.context_aware_chunker import chunk_text_context_aware from ingestion.models import ParsedChunk +from llm.base import LLMClient -def parse_text(filename: str, content: bytes) -> list[ParsedChunk]: - """Parse plain text-based files into a single text chunk. +def parse_text( + filename: str, + content: bytes, + llm: LLMClient | None = None, + semantic_boundaries: bool = False, + contextualize: bool = False, +) -> list[ParsedChunk]: + """Parse plain text-based files into paragraph-aware text chunks. + + When ``llm`` is provided, the LLM-based context-aware chunking strategy + (see :func:`ingestion.context_aware_chunker.chunk_text_context_aware`) + is used instead of the plain :func:`ingestion.chunker.chunk_text`. This + is opt-in: without an ``llm``, behavior is unchanged regardless of the + ``semantic_boundaries``/``contextualize`` flags, since there is no LLM + to call. Args: filename (str): Name of the source file. content (bytes): Raw file content as bytes. + llm (LLMClient | None, optional): + Client used for context-aware chunking. When ``None`` (the + default), plain paragraph/character-based chunking is used. + Defaults to ``None``. + + semantic_boundaries (bool, optional): + Let the LLM choose chunk boundaries based on semantic + coherence. Only relevant when ``llm`` is provided. Defaults to + ``False``. + + contextualize (bool, optional): + Let the LLM prepend a short situating context block to chunks + that would benefit from one. Only relevant when ``llm`` is + provided. Defaults to ``False``. + Returns: - list[ParsedChunk]: A list containing a single ParsedChunk with kind='text'. + list[ParsedChunk]: A list of ParsedChunk objects, each with kind='text'. """ text = content.decode(encoding="utf-8", errors="replace") - return chunk_text(filename, text) + use_context_aware_chunking = semantic_boundaries or contextualize + if not use_context_aware_chunking or llm is None: + return chunk_text(filename, text) + + return chunk_text_context_aware( + filename, + text, + llm, + semantic_boundaries=semantic_boundaries, + contextualize=contextualize, + ) diff --git a/tests/api/test_ingest.py b/tests/api/test_ingest.py index 0405740..4f2820d 100644 --- a/tests/api/test_ingest.py +++ b/tests/api/test_ingest.py @@ -1,3 +1,4 @@ +import json from collections.abc import Iterable from pathlib import Path @@ -7,6 +8,7 @@ from api.app import app from api.dependencies import get_ingestion_metadata_store, get_llm, get_store from ingestion.metadata_store import IngestionMetadataStore +from llm.base import Message from llm.errors import LLMUnavailableError from tests.stubs.llm import StubLLMClient from tests.stubs.store import StubVectorStore @@ -17,6 +19,18 @@ def embed(self, text: str) -> list[float]: raise LLMUnavailableError("embedding backend unavailable") +class RecordingLLMClient(StubLLMClient): + """Records every prompt passed to ``generate`` for assertion.""" + + def __init__(self, generate_response: str = "stub answer") -> None: + super().__init__(generate_response=generate_response) + self.generate_calls: list[list[Message]] = [] + + def generate(self, messages: list[Message]) -> str: + self.generate_calls.append(messages) + return super().generate(messages) + + @pytest.fixture def vector_store() -> StubVectorStore: return StubVectorStore() @@ -216,3 +230,118 @@ def test_ingest_defaults_to_primary_for_normal_files( chunks = vector_store.all_chunks() assert chunks assert all(chunk.source_role == "primary" for chunk in chunks) + + +def test_ingest_without_flags_skips_llm_call_by_default( + vector_store: StubVectorStore, + metadata_store: IngestionMetadataStore, +) -> None: + # Defaults (semantic_boundaries=False, contextualize=False) must not + # trigger an LLM call — the feature is opt-in. + llm = RecordingLLMClient() + app.dependency_overrides[get_store] = lambda: vector_store + app.dependency_overrides[get_llm] = lambda: llm + app.dependency_overrides[get_ingestion_metadata_store] = lambda: metadata_store + client = TestClient(app) + + response = client.post( + "/api/v1/ingest", + json={ + "artifact_id": "artifact-default-flags", + "filename": "notes.txt", + "content": "First paragraph.\n\nSecond paragraph.", + }, + ) + + app.dependency_overrides.clear() + + assert response.status_code == 200 + assert llm.generate_calls == [] + + +def test_ingest_calls_llm_when_flags_explicitly_enabled( + vector_store: StubVectorStore, + metadata_store: IngestionMetadataStore, +) -> None: + # Explicitly opting in should trigger an LLM call for text content, + # even though the stub's non-JSON response makes the chunker fall back + # to chunk_text. + llm = RecordingLLMClient() + app.dependency_overrides[get_store] = lambda: vector_store + app.dependency_overrides[get_llm] = lambda: llm + app.dependency_overrides[get_ingestion_metadata_store] = lambda: metadata_store + client = TestClient(app) + + response = client.post( + "/api/v1/ingest", + json={ + "artifact_id": "artifact-explicit-flags", + "filename": "notes.txt", + "content": "First paragraph.\n\nSecond paragraph.", + "semantic_boundaries": True, + "contextualize": True, + }, + ) + + app.dependency_overrides.clear() + + assert response.status_code == 200 + assert len(llm.generate_calls) == 1 + + +def test_ingest_with_both_chunking_flags_disabled_skips_llm_call( + vector_store: StubVectorStore, + metadata_store: IngestionMetadataStore, +) -> None: + llm = RecordingLLMClient() + app.dependency_overrides[get_store] = lambda: vector_store + app.dependency_overrides[get_llm] = lambda: llm + app.dependency_overrides[get_ingestion_metadata_store] = lambda: metadata_store + client = TestClient(app) + + response = client.post( + "/api/v1/ingest", + json={ + "artifact_id": "artifact-no-llm-chunking", + "filename": "notes.txt", + "content": "First paragraph.\n\nSecond paragraph.", + "semantic_boundaries": False, + "contextualize": False, + }, + ) + + app.dependency_overrides.clear() + + assert response.status_code == 200 + assert llm.generate_calls == [] + + +def test_ingest_contextualize_prepends_context_block( + vector_store: StubVectorStore, + metadata_store: IngestionMetadataStore, +) -> None: + plan = json.dumps( + {"boundaries": [], "context_blocks": {"0": "Context: onboarding notes."}} + ) + llm = RecordingLLMClient(generate_response=plan) + app.dependency_overrides[get_store] = lambda: vector_store + app.dependency_overrides[get_llm] = lambda: llm + app.dependency_overrides[get_ingestion_metadata_store] = lambda: metadata_store + client = TestClient(app) + + response = client.post( + "/api/v1/ingest", + json={ + "artifact_id": "artifact-contextualized", + "filename": "notes.txt", + "content": "Just one short paragraph.", + "semantic_boundaries": False, + "contextualize": True, + }, + ) + + app.dependency_overrides.clear() + + assert response.status_code == 200 + body = response.json() + assert body["chunks"][0]["text"].startswith("Context: onboarding notes.") diff --git a/tests/ingestion/test_context_aware_chunker.py b/tests/ingestion/test_context_aware_chunker.py new file mode 100644 index 0000000..a1cec58 --- /dev/null +++ b/tests/ingestion/test_context_aware_chunker.py @@ -0,0 +1,188 @@ +import json + +import pytest +from pytest import MonkeyPatch + +from ingestion.context_aware_chunker import ( + chunk_text_context_aware, + exceeds_llm_limit, +) +from llm.base import Message +from llm.errors import LLMUnavailableError +from tests.stubs.llm import StubLLMClient + + +class _RaisingLLMClient(StubLLMClient): + """Raises LLMUnavailableError instead of returning a response.""" + + def generate(self, messages: list[Message]) -> str: + raise LLMUnavailableError(host="http://stub") + + +def test_both_flags_disabled_skips_llm_and_matches_chunk_text(): + text = "Paragraph A\n\nParagraph B\n\nParagraph C" + llm = StubLLMClient(generate_response="should never be used") + + chunks = chunk_text_context_aware( + "file.txt", + text, + llm, + semantic_boundaries=False, + contextualize=False, + chunk_size=25, + ) + + assert len(chunks) == 2 + assert all(chunk.metadata["has_context_block"] == "false" for chunk in chunks) + assert all(chunk.metadata["has_overlap"] == "false" for chunk in chunks) + + +def test_empty_text_falls_back_to_chunk_text(): + llm = StubLLMClient(generate_response="should never be used") + + chunks = chunk_text_context_aware("file.txt", "", llm) + + assert chunks == [] + + +def test_text_exceeding_llm_limit_falls_back(monkeypatch: MonkeyPatch): + monkeypatch.setattr("ingestion.context_aware_chunker.max_chars", 10, raising=False) + llm = StubLLMClient(generate_response="should never be used") + text = "Paragraph A\n\nParagraph B\n\nParagraph C" + + chunks = chunk_text_context_aware("file.txt", text, llm, chunk_size=25) + + # Falls back to plain chunk_text grouping (paragraph accumulation). + assert len(chunks) == 2 + assert "Paragraph A" in chunks[0].content + + +def test_exceeds_llm_limit_helper(): + assert exceeds_llm_limit("a" * 100, max_chars=50) is True + assert exceeds_llm_limit("a" * 10, max_chars=50) is False + + +def test_semantic_boundaries_group_paragraphs_per_llm_plan(): + text = "Intro paragraph.\n\nDetail paragraph.\n\nUnrelated paragraph." + # boundaries=[0, 2] -> chunk 0 = paragraphs[0:2], chunk 1 = paragraphs[2:3] + plan = json.dumps({"boundaries": [0, 2], "context_blocks": {}}) + llm = StubLLMClient(generate_response=plan) + + chunks = chunk_text_context_aware( + "file.txt", + text, + llm, + semantic_boundaries=True, + contextualize=False, + chunk_size=1000, + ) + + assert len(chunks) == 2 + assert "Intro paragraph." in chunks[0].content + assert "Detail paragraph." in chunks[0].content + assert "Unrelated paragraph." in chunks[1].content + assert chunks[0].metadata["has_context_block"] == "false" + + +def test_contextualize_only_prepends_context_block_to_fixed_chunks(): + text = "Paragraph A\n\nParagraph B\n\nParagraph C" + plan = json.dumps( + {"boundaries": [], "context_blocks": {"0": "This chunk is about A and B."}} + ) + llm = StubLLMClient(generate_response=plan) + + chunks = chunk_text_context_aware( + "file.txt", + text, + llm, + semantic_boundaries=False, + contextualize=True, + chunk_size=1000, + ) + + assert len(chunks) == 1 + assert chunks[0].content.startswith("This chunk is about A and B.") + assert chunks[0].metadata["has_context_block"] == "true" + assert chunks[0].metadata["context_block_range"] != "" + + +def test_invalid_llm_output_falls_back_to_chunk_text(): + text = "Paragraph A\n\nParagraph B\n\nParagraph C" + llm = StubLLMClient(generate_response="not valid json at all") + + chunks = chunk_text_context_aware( + "file.txt", text, llm, semantic_boundaries=True, chunk_size=25 + ) + + assert len(chunks) == 2 + assert all(chunk.metadata["has_context_block"] == "false" for chunk in chunks) + + +def test_out_of_range_boundaries_falls_back_to_chunk_text(): + text = "Paragraph A\n\nParagraph B\n\nParagraph C" + plan = json.dumps({"boundaries": [0, 99], "context_blocks": {}}) + llm = StubLLMClient(generate_response=plan) + + chunks = chunk_text_context_aware( + "file.txt", text, llm, semantic_boundaries=True, chunk_size=25 + ) + + assert len(chunks) == 2 # chunk_text fallback result + + +def test_llm_unavailable_falls_back_to_chunk_text(): + text = "Paragraph A\n\nParagraph B\n\nParagraph C" + llm = _RaisingLLMClient() + + chunks = chunk_text_context_aware( + "file.txt", text, llm, semantic_boundaries=True, chunk_size=25 + ) + + assert len(chunks) == 2 + + +def test_context_block_causing_overflow_is_hard_split_and_duplicated(): + # A single paragraph forming one chunk; the context block pushes the + # combined chunk over chunk_size, triggering the hard-split path. + text = "B" * 30 + context = "Context sentence." + plan = json.dumps({"boundaries": [0], "context_blocks": {"0": context}}) + llm = StubLLMClient(generate_response=plan) + + chunks = chunk_text_context_aware( + "file.txt", + text, + llm, + semantic_boundaries=True, + contextualize=True, + chunk_size=40, + chunk_overlap=5, + ) + + assert len(chunks) >= 2 + for chunk in chunks: + assert len(chunk.content) <= 40 + assert chunk.content.startswith(context) + assert chunk.metadata["has_context_block"] == "true" + # sub-chunks after the first carry overlap of the body text + assert chunks[1].metadata["has_overlap"] == "true" + + +@pytest.mark.parametrize("semantic_boundaries", [True, False]) +def test_chunks_without_needs_context_are_left_untouched( + semantic_boundaries: bool, +): + text = "Paragraph A\n\nParagraph B" + plan = json.dumps({"boundaries": [0, 1], "context_blocks": {}}) + llm = StubLLMClient(generate_response=plan) + + chunks = chunk_text_context_aware( + "file.txt", + text, + llm, + semantic_boundaries=semantic_boundaries, + contextualize=True, + chunk_size=1000, + ) + + assert all(chunk.metadata["has_context_block"] == "false" for chunk in chunks)