77 lines
1.9 KiB
Python
77 lines
1.9 KiB
Python
|
|
"""Document chunking for RAG indexing.
|
||
|
|
|
||
|
|
Splits extracted text into overlapping chunks suitable for embedding
|
||
|
|
generation and vector search at the chunk level.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import hashlib
|
||
|
|
import logging
|
||
|
|
|
||
|
|
logger = logging.getLogger(__name__)
|
||
|
|
|
||
|
|
|
||
|
|
def chunk_text(
|
||
|
|
text: str,
|
||
|
|
chunk_size: int = 1000,
|
||
|
|
overlap: int = 200,
|
||
|
|
) -> list[dict]:
|
||
|
|
"""Split text into overlapping chunks.
|
||
|
|
|
||
|
|
Args:
|
||
|
|
text: Input text to chunk.
|
||
|
|
chunk_size: Maximum characters per chunk.
|
||
|
|
overlap: Number of overlapping characters between consecutive chunks.
|
||
|
|
|
||
|
|
Returns:
|
||
|
|
List of chunk dicts with keys:
|
||
|
|
- chunk_index: zero-based index
|
||
|
|
- chunk_text: the chunk content
|
||
|
|
- chunk_hash: sha256 hex digest of chunk_text
|
||
|
|
|
||
|
|
Edge cases:
|
||
|
|
- Empty text returns an empty list.
|
||
|
|
- Text shorter than chunk_size returns a single chunk.
|
||
|
|
"""
|
||
|
|
if not text or not text.strip():
|
||
|
|
return []
|
||
|
|
|
||
|
|
# Normalise whitespace to avoid degenerate chunks
|
||
|
|
cleaned = " ".join(text.split())
|
||
|
|
if not cleaned:
|
||
|
|
return []
|
||
|
|
|
||
|
|
chunks: list[dict] = []
|
||
|
|
start = 0
|
||
|
|
idx = 0
|
||
|
|
text_len = len(cleaned)
|
||
|
|
|
||
|
|
while start < text_len:
|
||
|
|
end = min(start + chunk_size, text_len)
|
||
|
|
chunk = cleaned[start:end]
|
||
|
|
|
||
|
|
if chunk.strip():
|
||
|
|
chunks.append(
|
||
|
|
{
|
||
|
|
"chunk_index": idx,
|
||
|
|
"chunk_text": chunk,
|
||
|
|
"chunk_hash": hashlib.sha256(chunk.encode("utf-8")).hexdigest(),
|
||
|
|
}
|
||
|
|
)
|
||
|
|
idx += 1
|
||
|
|
|
||
|
|
# If we've reached the end, stop
|
||
|
|
if end >= text_len:
|
||
|
|
break
|
||
|
|
|
||
|
|
# Advance by chunk_size - overlap
|
||
|
|
step = chunk_size - overlap
|
||
|
|
if step <= 0:
|
||
|
|
# Prevent infinite loop if overlap >= chunk_size
|
||
|
|
step = chunk_size
|
||
|
|
start += step
|
||
|
|
|
||
|
|
logger.debug("Chunked text (len=%d) into %d chunks", text_len, len(chunks))
|
||
|
|
return chunks
|