# Copyright © 2025 Dr.-Ing. Paul Wilhelm <paul@wilhelm.dev>
# This file is part of Archive Agent. See LICENSE for details.
from logging import Logger
import re
from concurrent.futures import ProcessPoolExecutor
import time
from typing import List, Tuple, Optional, Callable
from dataclasses import dataclass
import spacy
from spacy.language import Language
from spacy.tokens import Doc
import bisect
from archive_agent.ai.AiManager import AiManager
from archive_agent.ai.AiManagerFactory import AiManagerFactory
from archive_agent.ai.chunk.AiChunk import ChunkSchema
from archive_agent.data.DocumentContent import DocumentContent, ReferenceList
from archive_agent.util.format import format_file
from archive_agent.util.text_util import splitlines_exact
from archive_agent.core.ProgressManager import ProgressInfo
SentenceRange = Tuple[int, int]
SentenceBoundary = Tuple[str, int, int] # (text, start_char, end_char)
@dataclass
class SentenceWithRange:
text: str
reference_range: SentenceRange
@dataclass
class ChunkWithRange:
text: str
reference_range: SentenceRange
def _normalize_inline_whitespace(text: str) -> str:
"""
Normalize inline whitespace inside a string:
- Replace tabs and non-breaking spaces with normal space
- Collapse multiple spaces into one
- Strip leading and trailing whitespace
:param text: Single paragraph string
:return: Cleaned inline text
"""
text = text.replace("\t", " ").replace("\xa0", " ")
text = re.sub(r" {2,}", " ", text)
return text.strip()
def _extract_paragraph_sentences_with_reference_ranges(
paragraph_doc_content: DocumentContent,
sentence_boundaries: List[SentenceBoundary],
) -> List[SentenceWithRange]:
"""
Map pre-computed sentence boundaries to reference ranges.
:param paragraph_doc_content: Document content (single paragraph).
:param sentence_boundaries: Sentence boundaries from spaCy subprocess.
:return: List of SentenceWithRange for the block.
"""
paragraph_lines = paragraph_doc_content.lines
paragraph_references = paragraph_doc_content.get_per_line_references()
normalized_lines = [_normalize_inline_whitespace(line) for line in paragraph_lines]
line_starts: Optional[List[int]] = None
if paragraph_references:
line_starts = [0]
for norm_line in normalized_lines[:-1]:
line_starts.append(line_starts[-1] + len(norm_line) + 1)
block_sentences: List[SentenceWithRange] = []
for sent_text, start_char, end_char in sentence_boundaries:
sentence_text = sent_text.strip()
if sentence_text and paragraph_references and line_starts:
start_line_idx = bisect.bisect_right(line_starts, start_char) - 1
start_line_idx = min(start_line_idx, len(paragraph_lines) - 1)
end_line_idx = bisect.bisect_right(line_starts, end_char - 1) - 1
end_line_idx = min(end_line_idx, len(paragraph_lines) - 1)
sentence_refs = paragraph_references[start_line_idx: end_line_idx + 1]
min_ref = min(sentence_refs) if sentence_refs else 0
max_ref = max(sentence_refs) if sentence_refs else 0
else:
min_ref = max_ref = 0
block_sentences.append(SentenceWithRange(sentence_text, (min_ref, max_ref)))
return block_sentences
@Language.component("_spacy_markdown_sentence_fixer")
def _spacy_markdown_sentence_fixer(doc: Doc) -> Doc:
"""
Adjust sentence segmentation to better handle Markdown syntax.
Prevents sentence breaks after Markdown-specific tokens like headers, list items, and code blocks.
"""
skip_next = False
for i, token in enumerate(doc[:-1]):
if skip_next:
doc[i + 1].is_sent_start = False
skip_next = False
continue
# Inline code or triple backticks
if token.text in {"`", "```"}:
doc[i + 1].is_sent_start = False
continue
# Markdown headers
if token.text.startswith("#"):
doc[i + 1].is_sent_start = False
continue
# Bullet list
if token.text in {"-", "*"} and (i == 0 or doc[i - 1].text == "\n"):
doc[i + 1].is_sent_start = False
continue
# Avoid break inside code block fences
if token.text == "```":
skip_next = True
return doc
_nlp_cache: Optional[Language] = None
def _get_nlp() -> Language:
"""
Get NLP model for sentence splitting (cached).
:return: spaCy language model.
"""
global _nlp_cache
if _nlp_cache is not None:
return _nlp_cache
# Use blank English pipeline - only tokenizer + sentencizer needed.
nlp = spacy.blank("en")
nlp.max_length = 100_000_000
nlp.add_pipe("sentencizer")
nlp.add_pipe("_spacy_markdown_sentence_fixer", after="sentencizer")
_nlp_cache = nlp
return nlp
def _spacy_tokenize_paragraphs(para_texts: List[str]) -> List[List[SentenceBoundary]]:
"""
Run spaCy tokenization in a subprocess with its own GIL.
The Cython tokenizer holds the GIL without releasing it, starving the printer
thread and all other workers. Running in a subprocess isolates the GIL contention.
:param para_texts: Paragraph texts to tokenize.
:return: Sentence boundaries per paragraph: [(text, start_char, end_char), ...].
"""
nlp = _get_nlp()
return [
[(sent.text, sent.start_char, sent.end_char) for sent in nlp(text).sents]
for text in para_texts
]
_spacy_executor: Optional[ProcessPoolExecutor] = None
def _get_spacy_executor() -> ProcessPoolExecutor:
"""
Get or create the subprocess executor for spaCy tokenization.
:return: ProcessPoolExecutor with a single worker subprocess.
"""
global _spacy_executor
if _spacy_executor is None:
_spacy_executor = ProcessPoolExecutor(max_workers=1)
return _spacy_executor
def _extract_paragraphs(doc_content: DocumentContent) -> List[DocumentContent]:
"""
Extract paragraphs.
- Respects empty lines.
- Turns Markdown list items into separate paragraphs to let NLP (spaCy) work properly.
- Turns Markdown headings into separate paragraphs.
:param doc_content: Document content.
:return: List document content (one per paragraph).
"""
per_line_references = doc_content.get_per_line_references()
para_blocks: List[DocumentContent] = []
current_paragraph_lines: List[str] = []
current_reference_list: ReferenceList = []
has_references = bool(per_line_references)
for line_index, line_text in enumerate(doc_content.lines):
next_paragraph = False
discard_line = False
# Empty line starts new paragraph
if not line_text:
next_paragraph = True
discard_line = True
# Markdown item starts new paragraph
if line_text.startswith("- "):
next_paragraph = True
# Markdown heading starts new paragraph
if line_text.startswith("#"):
next_paragraph = True
# Push current paragraph and start next
if next_paragraph:
if current_paragraph_lines:
para_blocks.append(DocumentContent.from_lines(lines=current_paragraph_lines, lines_per_line=current_reference_list))
current_paragraph_lines = []
current_reference_list = []
# Discard line
if discard_line:
continue
# Append text line with page or line number reference
current_paragraph_lines.append(line_text)
if has_references:
current_reference_list.append(per_line_references[line_index])
# If heading, push as single-line paragraph
if line_text.startswith("#"):
para_blocks.append(DocumentContent.from_lines(lines=current_paragraph_lines, lines_per_line=current_reference_list))
current_paragraph_lines = []
current_reference_list = []
# Push last paragraph if non-empty
if current_paragraph_lines:
para_blocks.append(DocumentContent.from_lines(lines=current_paragraph_lines, lines_per_line=current_reference_list))
return para_blocks
def _extract_sentences_with_reference_ranges(
paragraphs_doc_content: List[DocumentContent],
) -> List[SentenceWithRange]:
# Prepare paragraph texts for batch tokenization in subprocess
para_texts: List[str] = []
for paragraph_doc_content in paragraphs_doc_content:
normalized_lines = [_normalize_inline_whitespace(line) for line in paragraph_doc_content.lines]
para_texts.append(" ".join(normalized_lines))
# Run spaCy in subprocess (own GIL - doesn't starve main process)
future = _get_spacy_executor().submit(_spacy_tokenize_paragraphs, para_texts)
all_sentence_boundaries: List[List[SentenceBoundary]] = future.result()
# Map sentence boundaries to reference ranges (main process, GIL-friendly)
sentences_with_reference_ranges: List[SentenceWithRange] = []
for paragraph_index, (paragraph_doc_content, sentence_boundaries) in enumerate(
zip(paragraphs_doc_content, all_sentence_boundaries)
):
if paragraph_index > 0:
# Insert paragraph delimiter: empty sentence with zero range.
# (Unit tests are sensitive to this logic)
sentences_with_reference_ranges.append(SentenceWithRange("", (0, 0)))
paragraph_sentences_with_reference_ranges = _extract_paragraph_sentences_with_reference_ranges(
paragraph_doc_content=paragraph_doc_content,
sentence_boundaries=sentence_boundaries,
)
sentences_with_reference_ranges.extend(paragraph_sentences_with_reference_ranges)
return sentences_with_reference_ranges
def get_sentences_with_reference_ranges(doc_content: DocumentContent) -> List[SentenceWithRange]:
"""
Use preprocessing and NLP (spaCy) to split text into sentences, keeping track of reference ranges.
Processes text with paragraph and sentence handling, inferring ranges from provided references (lines or pages).
spaCy runs in a subprocess to avoid GIL starvation of the printer thread and other workers.
:param doc_content: Document content.
:return: List of SentenceWithRange (text and range pairs).
"""
return _extract_sentences_with_reference_ranges(
paragraphs_doc_content=_extract_paragraphs(doc_content),
)
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
def _group_blocks_of_sentences(sentences: List[str], sentences_per_block: int) -> List[List[str]]:
"""
Group sentences into blocks of a specified size.
:param sentences: List of all sentences.
:param sentences_per_block: Number of sentences per block.
:return: List of sentence blocks.
"""
return [sentences[i:i + sentences_per_block] for i in range(0, len(sentences), sentences_per_block)]
def _chunk_start_to_ranges(start_lines: List[int], total_lines: int) -> List[SentenceRange]:
"""
Convert a list of chunk start lines into ranges of sentence indices.
:param start_lines: List of starting line numbers for chunks.
:param total_lines: Total number of lines in the block.
:return: List of (start, end) ranges for each chunk.
"""
extended = start_lines + [total_lines + 1]
return list(zip(extended[:-1], extended[1:]))
def _extract_chunks_and_carry(sentences: List[str], ranges: List[SentenceRange]) -> Tuple[List[str], str | None]:
"""
Extract chunks and any carry-over text from sentences based on ranges.
:param sentences: List of sentences in the block.
:param ranges: List of (start, end) ranges for chunks.
:return: Tuple of (list of chunk texts, carry-over text or None).
"""
if not ranges:
return [], None
if len(ranges) == 1:
start, end = ranges[0]
return [], "\n".join(sentences[start - 1:end - 1])
*main, last = ranges
chunks = ["\n".join(sentences[start - 1:end - 1]) for start, end in main]
carry = "\n".join(sentences[last[0] - 1:last[1] - 1])
return chunks, carry
def _aggregate_ranges(reference_ranges: List[SentenceRange]) -> SentenceRange:
"""
Aggregate min and max references from a list of ranges, ignoring zeros.
:param reference_ranges: List of (min, max) reference ranges.
:return: Aggregated (min, max) reference.
"""
valid_mins = [min_r for min_r, _ in reference_ranges if min_r > 0]
valid_maxs = [max_r for _, max_r in reference_ranges if max_r > 0]
r_min = min(valid_mins) if valid_mins else 0
r_max = max(valid_maxs) if valid_maxs else 0
return r_min, r_max
def _format_chunk(_file_path: str, header: str, body: str) -> str:
"""
Format chunk.
:param _file_path: File path.
:param header: Chunk header.
:param body: Chunk body.
:return: Formatted chunk text.
"""
return "\n".join([
# NOTE: Removed in v11.0.0. Use `tools/qdrant-remove-paths-from-context-headers.py` to remove this from pre-v11.0.0.
# f"# {format_file(file_path)}",
# Context header
f"# {header}",
f"",
# Chunk body
body,
])
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
def get_chunks_with_reference_ranges(
ai_factory: AiManagerFactory,
sentences_with_references: List[SentenceWithRange],
chunk_callback: Callable[[AiManager, List[str]], ChunkSchema],
chunk_lines_block: int,
file_path: str,
progress_info: ProgressInfo,
logger: Logger,
verbose: bool = True,
) -> List[ChunkWithRange]:
"""
Chunkify a list of sentences into AI-determined chunks, carrying over leftover sentences where needed,
and annotate each chunk with the corresponding (min, max) line reference range.
:param ai_factory: AI manager factory.
:param sentences_with_references: List of sentences with their reference ranges.
:param chunk_callback: Chunk callback.
:param chunk_lines_block: Number of sentences per block to be chunked.
:param file_path: Path to the originating file (used for logging and labeling).
:param progress_info: Progress tracking information
:param logger: Logger.
:param verbose: Enable to show additional information.
:return: List of ChunkWithRange objects containing the formatted chunk and its reference range.
"""
sentences = [s.text for s in sentences_with_references]
sentence_reference_ranges = [s.reference_range for s in sentences_with_references]
if len(sentence_reference_ranges) != len(sentences):
raise ValueError(
f"Reference range mismatch: "
f"{len(sentence_reference_ranges)} ranges, "
f"{len(sentences)} sentences "
f"for {format_file(file_path)}"
)
blocks_of_sentences = _group_blocks_of_sentences(sentences, chunk_lines_block)
logger.info(
f"Chunking ({len(sentences)}) sentences in ({len(blocks_of_sentences)}) blocks "
f"for {format_file(file_path)}"
)
# Set progress total based on total sentences (more meaningful than block count)
progress_info.progress_manager.set_total(progress_info.parent_key, len(sentences_with_references))
chunks_with_ranges: List[ChunkWithRange] = []
carry: Optional[str] = None
carry_reference_ranges: Optional[List[SentenceRange]] = None
last_carry_header: Optional[str] = None
failed_blocks = 0
chunk_t0 = time.monotonic()
idx = 0
for block_index, block_of_sentences in enumerate(blocks_of_sentences):
block_len = len(block_of_sentences)
block_sentence_reference_ranges = sentence_reference_ranges[idx: idx + block_len]
idx += block_len
if verbose:
logger.info(f"Chunking block ({block_index + 1}) / ({len(blocks_of_sentences)}) of {format_file(file_path)}")
if carry:
assert carry_reference_ranges is not None
carry_lines = splitlines_exact(carry)
current_block_line_count: int = len(carry_lines) + len(block_of_sentences)
if verbose:
logger.info(f"Carrying over ({len(carry_lines)}) lines; current block has ({current_block_line_count}) lines")
block_of_sentences = carry_lines + block_of_sentences
block_sentence_reference_ranges = carry_reference_ranges + block_sentence_reference_ranges
range_start = block_sentence_reference_ranges[0] if block_sentence_reference_ranges else (0, 0)
range_stop = block_sentence_reference_ranges[-1] if block_sentence_reference_ranges else (0, 0)
logger.debug(f"Chunking block {block_index + 1}: {len(block_of_sentences)} sentences, range {range_start} to {range_stop}")
try:
ai = ai_factory.get_ai()
chunk_result: ChunkSchema = chunk_callback(ai, block_of_sentences)
except Exception as e:
failed_blocks += 1
logger.critical(
f"BLOCK SKIPPED: Chunking failed for block ({block_index + 1}) / ({len(blocks_of_sentences)}) "
f"of {format_file(file_path)} — ({len(block_of_sentences)}) sentences lost: {e}"
)
carry = None
carry_reference_ranges = None
last_carry_header = None
progress_info.progress_manager.update_task(progress_info.parent_key, advance=block_len)
continue
chunk_start_lines = chunk_result.get_chunk_start_lines()
chunk_headers = chunk_result.get_chunk_headers()
ranges = _chunk_start_to_ranges(start_lines=chunk_start_lines, total_lines=len(block_of_sentences))
block_chunks, carry = _extract_chunks_and_carry(block_of_sentences, ranges)
for i, (r_start, r_end) in enumerate(ranges[:-1] if carry else ranges):
r_reference_ranges = block_sentence_reference_ranges[r_start - 1:r_end - 1]
r_range = _aggregate_ranges(r_reference_ranges)
if r_range == (0, 0):
# AI returned `chunk_start_sentences` selecting empty lines only
logger.warning(f"Chunk {len(chunks_with_ranges) + 1}: Discarding empty chunk")
continue
# DEBUG
# logger.info(f"Chunk {len(chunks_with_ranges) + 1}: sentences {r_start}:{r_end}, range {r_range}")
body = "\n".join(block_of_sentences[r_start - 1:r_end - 1])
chunk_text = _format_chunk(
_file_path=file_path,
header=chunk_headers[i],
body=body,
)
chunks_with_ranges.append(ChunkWithRange(chunk_text, r_range))
if carry:
carry_reference_ranges = block_sentence_reference_ranges[ranges[-1][0] - 1:ranges[-1][1] - 1]
last_carry_header = chunk_headers[-1]
# Update progress after processing each block (advance by number of sentences processed)
progress_info.progress_manager.update_task(progress_info.parent_key, advance=block_len)
block_elapsed = time.monotonic() - chunk_t0
logger.info(
f"Chunked block ({block_index + 1}) / ({len(blocks_of_sentences)}) "
f"({len(chunks_with_ranges)} chunks so far, {block_elapsed:.1f}s elapsed)"
)
if failed_blocks > 0:
logger.critical(
f"INCOMPLETE INGESTION: ({failed_blocks}) of ({len(blocks_of_sentences)}) blocks failed chunking "
f"for {format_file(file_path)} — re-run with a different model or smaller block size"
)
if carry:
assert last_carry_header is not None, "Internal error: carry exists but no header was recorded"
assert carry_reference_ranges is not None
carry_range = _aggregate_ranges(carry_reference_ranges)
if carry_range == (0, 0):
logger.warning(f"Chunk {len(chunks_with_ranges) + 1}: Discarding empty chunk (final)")
return chunks_with_ranges
carry_lines = splitlines_exact(carry)
final_chunk_line_count: int = len(carry_lines)
if verbose:
logger.info(f"Appending final carry of ({final_chunk_line_count}) lines; final chunk has ({final_chunk_line_count}) lines")
formatted_carry = _format_chunk(
_file_path=file_path,
header=last_carry_header,
body=carry,
)
chunks_with_ranges.append(ChunkWithRange(formatted_carry, carry_range))
return chunks_with_ranges
