"""Layout-aware document chunker using Azure Document Intelligence and LangChain.
Extracts Markdown from documents using Document Intelligence prebuilt-layout,
protects HTML tables from mid-table splits, preserves page numbers from PageBreak
comments, splits on Markdown headings, and chunks sections to profile limits.
Production exception contract
-------------------------------
ConfigurationError – bad or missing environment / profile configuration.
AnalysisError – non-retryable failure from the Document Intelligence service.
ThrottleError – all retry attempts exhausted (HTTP 429).
ChunkError – unexpected state during chunking (e.g. corrupted cache token).
All other exceptions propagate unchanged.
"""
import hashlib
import io
import json
import logging
import os
import re
import time
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Union
from urllib.parse import urlparse
from azure.ai.documentintelligence import DocumentIntelligenceClient
from azure.ai.documentintelligence.models import (
AnalyzeDocumentRequest,
DocumentAnalysisFeature,
DocumentContentFormat,
)
from azure.core.credentials import AzureKeyCredential
from azure.core.exceptions import HttpResponseError, ServiceRequestError
from azure.identity import DefaultAzureCredential
from langchain_core.documents import Document
from langchain_text_splitters import (
MarkdownHeaderTextSplitter,
RecursiveCharacterTextSplitter,
)
log = logging.getLogger(__name__)
# ---------------------------------------------------------------------------
# Compiled regex patterns
# ---------------------------------------------------------------------------
TABLE_PATTERN = re.compile(r"
", re.DOTALL | re.IGNORECASE)
PAGE_META_PATTERN = re.compile(r"", re.DOTALL)
PAGE_BREAK_PATTERN = re.compile(r"")
PAGE_TOKEN_PATTERN = re.compile(r"@@P(\d+)@@")
TABLE_TOKEN_PATTERN = re.compile(r"TBL(\d+)X")
REQUIRED_PROFILE_KEYS = {
"profile_id",
"header_levels",
"chunk_size",
"chunk_overlap",
"prefix_heading",
}
# Retry schedule in seconds for HTTP 429 (Too Many Requests).
# Fibonacci-like spacing; matches the service-limits guidance.
_RETRY_WAITS = (2, 5, 13, 34)
# Minimum length of the custom-subdomain endpoint to catch obvious mis-configs.
_MIN_ENDPOINT_LEN = 20
# ---------------------------------------------------------------------------
# Custom exceptions
# ---------------------------------------------------------------------------
class ConfigurationError(Exception):
"""Raised when environment variables or profile values are invalid."""
class AnalysisError(Exception):
"""Raised on a non-retryable HTTP error from Document Intelligence."""
def __init__(self, message: str, status_code: Optional[int] = None) -> None:
super().__init__(message)
self.status_code = status_code
class ThrottleError(Exception):
"""Raised when all retry attempts are exhausted due to HTTP 429."""
class ChunkError(Exception):
"""Raised when an internal chunking invariant is violated."""
# ---------------------------------------------------------------------------
# Client factory
# ---------------------------------------------------------------------------
def make_client() -> DocumentIntelligenceClient:
"""Create a Document Intelligence client using Entra ID or API key.
Reads DOCINTEL_ENDPOINT from the environment. If DOCINTEL_KEY is also set,
uses AzureKeyCredential (sandbox only). Otherwise uses DefaultAzureCredential
(Entra ID — preferred for production). Requires a custom-subdomain endpoint
(https://${DOCINTEL_NAME}.cognitiveservices.azure.com/); regional endpoints do not
support Entra ID authentication.
Raises:
ConfigurationError: DOCINTEL_ENDPOINT is missing or obviously malformed.
"""
endpoint = os.environ.get("DOCINTEL_ENDPOINT", "").strip()
if not endpoint:
raise ConfigurationError(
"DOCINTEL_ENDPOINT is not set. "
"Export https://${DOCINTEL_NAME}.cognitiveservices.azure.com/ before running."
)
if len(endpoint) < _MIN_ENDPOINT_LEN or not endpoint.startswith("https://"):
raise ConfigurationError(
f"DOCINTEL_ENDPOINT looks malformed: {endpoint!r}. "
"Use a custom-subdomain endpoint starting with https://."
)
key = os.environ.get("DOCINTEL_KEY", "").strip()
if key:
log.debug("Authenticating with API key (sandbox mode).")
return DocumentIntelligenceClient(endpoint, AzureKeyCredential(key))
log.debug("Authenticating with DefaultAzureCredential (Entra ID).")
return DocumentIntelligenceClient(endpoint, DefaultAzureCredential())
# ---------------------------------------------------------------------------
# Markdown extraction
# ---------------------------------------------------------------------------
def to_markdown(
client: DocumentIntelligenceClient,
source: Union[str, Path],
cache_dir: Union[str, Path] = "cache",
**options: Any,
) -> str:
"""Analyse a document and return layout-aware Markdown.
Caches the result by SHA-256 hash of the file bytes and options so that
re-runs are free. Retries automatically on HTTP 429 with exponential back-off.
Args:
client: Authenticated DocumentIntelligenceClient.
source: Local file path or https:// URL.
cache_dir: Directory for cached Markdown files. Created if absent.
**options: Extra keyword arguments forwarded to begin_analyze_document
(e.g. pages="1-3", locale="en-US", features=[...]).
Returns:
Markdown string from Document Intelligence prebuilt-layout.
Raises:
ConfigurationError: source is neither a readable file nor an https URL.
AnalysisError: Non-retryable HTTP error from the service.
ThrottleError: All retry attempts exhausted (HTTP 429).
OSError: File I/O failure (file unreadable, disk full, etc.).
"""
source_str = str(source).strip()
is_url = source_str.lower().startswith("https://")
# ------------------------------------------------------------------ input
if is_url:
parsed = urlparse(source_str)
if not parsed.netloc:
raise ConfigurationError(
f"URL appears malformed (no host): {source_str!r}"
)
raw_bytes = source_str.encode("utf-8")
# Use the last path segment as the stem; fall back to "document".
stem = Path(parsed.path.rstrip("/")).stem or "document"
else:
file_path = Path(source_str)
if not file_path.exists():
raise ConfigurationError(f"Source file not found: {file_path}")
if not file_path.is_file():
raise ConfigurationError(f"Source path is not a file: {file_path}")
try:
raw_bytes = file_path.read_bytes()
except OSError as exc:
raise OSError(f"Cannot read source file {file_path}: {exc}") from exc
if not raw_bytes:
raise ConfigurationError(f"Source file is empty: {file_path}")
stem = file_path.stem
# --------------------------------------------------------------- cache key
options_key = repr(sorted(options.items())).encode("utf-8")
cache_key = hashlib.sha256(raw_bytes + options_key).hexdigest()[:16]
cache_path = Path(cache_dir) / f"{stem}.{cache_key}.md"
if cache_path.exists():
log.debug("Cache hit: %s", cache_path)
try:
return cache_path.read_text(encoding="utf-8")
except OSError as exc:
# Cache file unreadable (permissions, corruption). Re-analyse.
log.warning("Cache file unreadable (%s); re-analysing: %s", exc, source_str)
# ----------------------------------------------------------------- analyse
log.info("Analysing document: %s", source_str)
last_exc: Optional[HttpResponseError] = None
for attempt, wait_seconds in enumerate(_RETRY_WAITS, start=1):
# Re-create BytesIO on every attempt — stream is at EOF after first use.
body: Union[AnalyzeDocumentRequest, io.BytesIO]
if is_url:
body = AnalyzeDocumentRequest(url_source=source_str)
else:
body = io.BytesIO(raw_bytes)
try:
poller = client.begin_analyze_document(
"prebuilt-layout",
body=body,
output_content_format=DocumentContentFormat.MARKDOWN,
**options,
)
result = poller.result()
break
except HttpResponseError as exc:
if exc.status_code == 429:
retry_after = _parse_retry_after(exc) or wait_seconds
log.warning(
"HTTP 429 on attempt %d/%d for %s. Waiting %ds.",
attempt,
len(_RETRY_WAITS),
source_str,
retry_after,
)
last_exc = exc
time.sleep(retry_after)
continue
# Non-retryable HTTP error (400, 401, 403, 404, 413, 500, …)
raise AnalysisError(
f"Document Intelligence returned HTTP {exc.status_code} "
f"for {source_str!r}: {exc.message}",
status_code=exc.status_code,
) from exc
except ServiceRequestError as exc:
# Network-level failure (DNS, TLS, connection reset).
raise AnalysisError(
f"Network error while analysing {source_str!r}: {exc}"
) from exc
else:
raise ThrottleError(
f"Document analysis throttled after {len(_RETRY_WAITS)} retries "
f"for {source_str!r}. Last Retry-After: "
f"{_parse_retry_after(last_exc) if last_exc else 'unknown'}s. "
"Request a quota increase or reduce batch concurrency."
)
# Defensive: service should always return content for a succeeded result.
markdown_content = result.content or ""
if not markdown_content:
log.warning(
"Document Intelligence returned empty content for %s. "
"Check document validity and tier (F0 analyses only 2 pages).",
source_str,
)
# ------------------------------------------------------------------- cache
cache_path.parent.mkdir(parents=True, exist_ok=True)
try:
cache_path.write_text(markdown_content, encoding="utf-8")
log.debug("Cached Markdown to %s", cache_path)
except OSError as exc:
# Write failure must not discard the result — log and continue.
log.warning("Could not write cache file %s: %s", cache_path, exc)
return markdown_content
def _parse_retry_after(exc: Optional[HttpResponseError]) -> Optional[int]:
"""Extract the Retry-After header value in seconds, if present."""
if exc is None:
return None
try:
header = (exc.response.headers or {}).get("Retry-After", "")
return int(header) if header else None
except (AttributeError, ValueError):
return None
# ---------------------------------------------------------------------------
# Chunker
# ---------------------------------------------------------------------------
def chunk(
md: str,
source: str,
profile: Dict[str, Any],
measure: Callable[[str], int] = len,
) -> List[Dict[str, Any]]:
"""Chunk Markdown while preserving tables and attaching heading/page metadata.
Args:
md: Markdown string from Document Intelligence prebuilt-layout.
source: Identifier or file path used for citations and chunk IDs.
profile: Chunking configuration dict (see REQUIRED_PROFILE_KEYS).
measure: Length function. Defaults to character count (len). Pass a
token counter (e.g. tokenizer.encode) to size chunks in tokens.
Returns:
List of chunk dicts, each with: chunk_id, source, page, heading_path,
seq, profile_id, content, embed_text.
Raises:
ConfigurationError: Profile is missing required keys or has invalid values.
ChunkError: Internal token mismatch (e.g. corrupted cache file).
"""
# --------------------------------------------------------- validate profile
missing_keys = REQUIRED_PROFILE_KEYS - set(profile.keys())
if missing_keys:
raise ConfigurationError(
f"Profile is missing required keys: {sorted(missing_keys)}"
)
try:
chunk_size = int(profile["chunk_size"])
chunk_overlap = int(profile["chunk_overlap"])
header_levels = int(profile["header_levels"])
except (TypeError, ValueError) as exc:
raise ConfigurationError(
f"Profile values for chunk_size, chunk_overlap, header_levels "
f"must be integers: {exc}"
) from exc
try:
prefix_heading = bool(profile["prefix_heading"])
profile_id = str(profile["profile_id"])
keep_tables_whole = bool(profile.get("keep_tables_whole", True))
except (TypeError, ValueError) as exc:
raise ConfigurationError(
f"Profile values for prefix_heading, profile_id, or keep_tables_whole are invalid: {exc}"
) from exc
if chunk_size <= 0:
raise ConfigurationError(
f"chunk_size must be a positive integer; got {chunk_size}."
)
if chunk_overlap < 0:
raise ConfigurationError(
f"chunk_overlap must be >= 0; got {chunk_overlap}."
)
if chunk_overlap >= chunk_size:
raise ConfigurationError(
f"chunk_overlap ({chunk_overlap}) must be less than "
f"chunk_size ({chunk_size})."
)
if header_levels < 0:
raise ConfigurationError(
f"header_levels must be >= 0; got {header_levels}."
)
if header_levels > 6:
raise ConfigurationError(
f"header_levels must be <= 6 (Markdown has six heading levels); "
f"got {header_levels}."
)
if not source:
raise ConfigurationError("source must be a non-empty string.")
if not md:
log.warning("chunk() received empty Markdown for source %r.", source)
return []
# ------------------------------------------------ table masking / page tracking
tables: List[str] = []
page_counter = 0
def store_table(table_html: str) -> str:
tables.append(table_html)
return f"TBL{len(tables) - 1}X"
def mask_table(match: re.Match) -> str:
table_html = match.group(0)
if measure(table_html) <= chunk_size:
return store_table(table_html)
rows = ROW_PATTERN.findall(table_html)
if not rows:
log.debug("Table with no
rows encountered; treating as atomic.")
return store_table(table_html)
# First row is treated as the header row regardless of tag type.
header_row = rows[0]
data_rows = rows[1:]
if not data_rows:
# Single-row table: store as-is even if oversized.
log.debug("Table has only one row; storing atomically.")
return store_table(table_html)
groups: List[List[str]] = []
current_group: List[str] = []
for row in data_rows:
candidate = f"
{header_row}{''.join(current_group)}{row}
"
if current_group and measure(candidate) > chunk_size:
groups.append(current_group)
current_group = []
current_group.append(row)
if current_group:
groups.append(current_group)
table_chunks = [
f"
{header_row}{''.join(group)}
" for group in groups
]
return "\n\n".join(store_table(tbl) for tbl in table_chunks)
def mark_page_break(_: re.Match) -> str:
nonlocal page_counter
page_counter += 1
return f"\n@@P{page_counter + 1}@@\n"
cleaned_md = PAGE_META_PATTERN.sub("", md)
paged_md = PAGE_BREAK_PATTERN.sub(mark_page_break, cleaned_md)
if keep_tables_whole:
masked_md = TABLE_PATTERN.sub(mask_table, paged_md)
else:
masked_md = paged_md
def restore_tables(text: str) -> str:
"""Replace TBL{n}X tokens with the stored HTML. Raises ChunkError on bad index."""
def _restore(m: re.Match) -> str:
idx = int(m.group(1))
if idx >= len(tables):
raise ChunkError(
f"Table token TBL{idx}X references index {idx} but only "
f"{len(tables)} table(s) were masked. "
"The Markdown may be corrupted or from a mismatched cache."
)
return tables[idx]
return TABLE_TOKEN_PATTERN.sub(_restore, text)
# --------------------------------------------------------- header splitting
if header_levels > 0:
header_keys = [f"h{i}" for i in range(1, header_levels + 1)]
headers_to_split = [("#" * (i + 1), k) for i, k in enumerate(header_keys)]
header_splitter = MarkdownHeaderTextSplitter(
headers_to_split_on=headers_to_split,
strip_headers=True,
)
sections = header_splitter.split_text(masked_md)
else:
header_keys = []
sections = [Document(page_content=masked_md, metadata={})]
if not sections:
log.warning("Header splitter produced no sections for source %r.", source)
return []
# ------------------------------------------------- recursive size splitting
recursive_splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
length_function=lambda s: measure(restore_tables(s)),
)
documents = recursive_splitter.split_documents(sections)
# ---------------------------------------------------------- assemble output
chunks: List[Dict[str, Any]] = []
current_page = 1
for seq, doc in enumerate(documents):
page_matches = [int(p) for p in PAGE_TOKEN_PATTERN.findall(doc.page_content)]
if page_matches and re.match(r"^\s*@@P", doc.page_content):
assigned_page = page_matches[0]
else:
assigned_page = current_page
if page_matches:
current_page = page_matches[-1]
raw_content = PAGE_TOKEN_PATTERN.sub("", doc.page_content)
final_content = restore_tables(raw_content).strip()
if not final_content:
continue
heading_breadcrumbs = [
str(doc.metadata[k]) for k in header_keys if k in doc.metadata
]
heading_path = " > ".join(heading_breadcrumbs)
chunk_hash = hashlib.sha256(
f"{source}|{seq}|{final_content}".encode("utf-8")
).hexdigest()[:16]
if prefix_heading and heading_path:
embed_text = f"{heading_path}\n\n{final_content}"
else:
embed_text = final_content
chunks.append(
{
"chunk_id": chunk_hash,
"source": str(source),
"page": assigned_page,
"heading_path": heading_path,
"seq": seq,
"profile_id": profile_id,
"content": final_content,
"embed_text": embed_text,
}
)
log.info(
"Chunked %r into %d chunks (profile: %s, chunk_size: %d).",
source,
len(chunks),
profile_id,
chunk_size,
)
return chunks
# ---------------------------------------------------------------------------
# Quality report
# ---------------------------------------------------------------------------
def report(
chunks: List[Dict[str, Any]],
profile: Dict[str, Any],
measure: Callable[[str], int] = len,
) -> Dict[str, int]:
"""Return summary quality statistics for a set of chunks.
Args:
chunks: Output of chunk().
profile: Must contain chunk_size (int).
measure: Same length function used when chunking.
Returns:
Dict with keys: chunks, over_max, no_heading, broken_tables.
Raises:
ConfigurationError: profile is missing chunk_size or it is not an int.
"""
if "chunk_size" not in profile:
raise ConfigurationError(
"report() requires profile['chunk_size']. "
"Pass the same profile used during chunking."
)
try:
max_size = int(profile["chunk_size"])
except (TypeError, ValueError) as exc:
raise ConfigurationError(
f"profile['chunk_size'] must be an integer: {exc}"
) from exc
return {
"chunks": len(chunks),
"over_max": sum(measure(c["content"]) > max_size for c in chunks),
"no_heading": sum(not c["heading_path"] for c in chunks),
"broken_tables": sum(
c["content"].count("
") != c["content"].count("
")
for c in chunks
),
}
# ---------------------------------------------------------------------------
# Self-test (run locally without Azure credentials)
# ---------------------------------------------------------------------------
def sample_run() -> None:
"""Execute a self-check against illustrative Markdown (no Azure calls)."""
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
sample_profile = {
"profile_id": "${PROFILE_ID}",
"header_levels": 3,
"chunk_size": 150,
"chunk_overlap": 0,
"prefix_heading": True,
"keep_tables_whole": True,
}
sample_md = (
"# Pump Manual\n"
"## 1. Safety\n"
"Isolate the pump before any work. Confirm pressure is zero.\n"
"## 2. Service intervals\n"
"