Initial scaffold: .gitignore, config, requirements, data seeds, directory structure
This commit is contained in:
n0mad1k
@@ -0,0 +1,5 @@
|
||||
from .http_limiter import HTTPRateLimiter
|
||||
from .api_limiter import APIRateLimiter
|
||||
from .cache import DownloadCache
|
||||
from .sanitize import sanitize_filename, sanitize_path, validate_content_type
|
||||
from .text_utils import chunk_text, clean_text, estimate_tokens
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Exponential backoff rate limiter with circuit breaker for Anthropic API."""
|
||||
import time
|
||||
import random
|
||||
import threading
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class CircuitBreakerOpen(Exception):
|
||||
"""Raised when circuit breaker is open (too many consecutive failures)."""
|
||||
pass
|
||||
|
||||
|
||||
class APIRateLimiter:
|
||||
"""Exponential backoff with jitter and circuit breaker for API calls."""
|
||||
|
||||
def __init__(self, base_delay: float = 1.0, max_delay: float = 60.0,
|
||||
circuit_threshold: int = 5):
|
||||
"""
|
||||
Args:
|
||||
base_delay: Initial backoff delay in seconds
|
||||
max_delay: Maximum backoff delay in seconds
|
||||
circuit_threshold: Consecutive failures before circuit opens
|
||||
"""
|
||||
self.base_delay = base_delay
|
||||
self.max_delay = max_delay
|
||||
self.circuit_threshold = circuit_threshold
|
||||
self._consecutive_failures = 0
|
||||
self._current_delay = 0.0
|
||||
self._lock = threading.Lock()
|
||||
|
||||
@property
|
||||
def is_circuit_open(self) -> bool:
|
||||
return self._consecutive_failures >= self.circuit_threshold
|
||||
|
||||
def record_success(self):
|
||||
"""Record a successful API call — reset backoff."""
|
||||
with self._lock:
|
||||
self._consecutive_failures = 0
|
||||
self._current_delay = 0.0
|
||||
|
||||
def record_failure(self, error: Optional[str] = None):
|
||||
"""Record a failed API call — increase backoff."""
|
||||
with self._lock:
|
||||
self._consecutive_failures += 1
|
||||
logger.warning(
|
||||
"API failure #%d/%d: %s",
|
||||
self._consecutive_failures,
|
||||
self.circuit_threshold,
|
||||
error or "unknown"
|
||||
)
|
||||
if self.is_circuit_open:
|
||||
logger.error(
|
||||
"Circuit breaker OPEN: %d consecutive failures. "
|
||||
"Halting API calls until manual reset.",
|
||||
self._consecutive_failures
|
||||
)
|
||||
|
||||
def wait_if_needed(self):
|
||||
"""Wait with exponential backoff if there have been failures. Raises CircuitBreakerOpen."""
|
||||
with self._lock:
|
||||
if self.is_circuit_open:
|
||||
raise CircuitBreakerOpen(
|
||||
f"Circuit breaker open after {self._consecutive_failures} consecutive failures. "
|
||||
f"Call reset() to retry."
|
||||
)
|
||||
|
||||
if self._consecutive_failures == 0:
|
||||
return
|
||||
|
||||
delay = min(
|
||||
self.base_delay * (2 ** (self._consecutive_failures - 1)),
|
||||
self.max_delay
|
||||
)
|
||||
# Add jitter (±25%)
|
||||
jitter = delay * random.uniform(-0.25, 0.25)
|
||||
actual_delay = max(0, delay + jitter)
|
||||
|
||||
logger.info("API backoff: waiting %.1fs (attempt %d)",
|
||||
actual_delay, self._consecutive_failures + 1)
|
||||
|
||||
# Sleep outside the lock
|
||||
time.sleep(actual_delay)
|
||||
|
||||
def reset(self):
|
||||
"""Manually reset the circuit breaker."""
|
||||
with self._lock:
|
||||
self._consecutive_failures = 0
|
||||
self._current_delay = 0.0
|
||||
logger.info("Circuit breaker reset")
|
||||
+147
@@ -0,0 +1,147 @@
|
||||
"""Atomic download cache with hash verification for Mosaic."""
|
||||
import os
|
||||
import hashlib
|
||||
import shutil
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class DownloadCache:
|
||||
"""Manages atomic downloads with hash-sharded storage and integrity verification."""
|
||||
|
||||
def __init__(self, base_dir: str = "output/raw"):
|
||||
self.base_dir = Path(base_dir)
|
||||
self.base_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
def get_shard_path(self, content_hash: str, collection: str, ext: str = "") -> Path:
|
||||
"""Get hash-sharded storage path: raw/{collection}/{ab}/{cd}/{hash}.{ext}"""
|
||||
if len(content_hash) < 4:
|
||||
raise ValueError(f"Invalid hash: {content_hash}")
|
||||
ab = content_hash[:2]
|
||||
cd = content_hash[2:4]
|
||||
filename = content_hash + (f".{ext}" if ext else "")
|
||||
return self.base_dir / collection / ab / cd / filename
|
||||
|
||||
def store_atomic(self, tmp_path: str, content_hash: str, collection: str,
|
||||
ext: str = "", expected_size: Optional[int] = None) -> str:
|
||||
"""Atomically move a temp file to its cache location.
|
||||
|
||||
Args:
|
||||
tmp_path: Path to the temporary download file
|
||||
content_hash: SHA256 hash of the content
|
||||
collection: Source collection name
|
||||
ext: File extension
|
||||
expected_size: Expected file size (from Content-Length)
|
||||
|
||||
Returns:
|
||||
Final path of the cached file
|
||||
|
||||
Raises:
|
||||
ValueError: If hash or size verification fails
|
||||
"""
|
||||
tmp = Path(tmp_path)
|
||||
if not tmp.exists():
|
||||
raise FileNotFoundError(f"Temp file not found: {tmp_path}")
|
||||
|
||||
# Verify size if expected
|
||||
if expected_size is not None:
|
||||
actual_size = tmp.stat().st_size
|
||||
if actual_size != expected_size:
|
||||
raise ValueError(
|
||||
f"Size mismatch: expected {expected_size}, got {actual_size}"
|
||||
)
|
||||
|
||||
# Verify hash
|
||||
actual_hash = self.compute_hash(tmp_path)
|
||||
if actual_hash != content_hash:
|
||||
raise ValueError(
|
||||
f"Hash mismatch: expected {content_hash}, got {actual_hash}"
|
||||
)
|
||||
|
||||
# Atomic move to final location
|
||||
final_path = self.get_shard_path(content_hash, collection, ext)
|
||||
final_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
shutil.move(str(tmp), str(final_path))
|
||||
logger.debug("Cached: %s → %s", tmp_path, final_path)
|
||||
return str(final_path)
|
||||
|
||||
def is_cached(self, content_hash: str, collection: str, ext: str = "") -> bool:
|
||||
"""Check if a file exists in cache with valid hash."""
|
||||
path = self.get_shard_path(content_hash, collection, ext)
|
||||
if not path.exists():
|
||||
return False
|
||||
actual_hash = self.compute_hash(str(path))
|
||||
if actual_hash != content_hash:
|
||||
logger.warning("Cache corruption detected: %s (expected %s, got %s)",
|
||||
path, content_hash, actual_hash)
|
||||
return False
|
||||
return True
|
||||
|
||||
def get_cached_path(self, content_hash: str, collection: str, ext: str = "") -> Optional[str]:
|
||||
"""Get path to cached file, or None if not cached/corrupted."""
|
||||
if self.is_cached(content_hash, collection, ext):
|
||||
return str(self.get_shard_path(content_hash, collection, ext))
|
||||
return None
|
||||
|
||||
def verify_file(self, file_path: str, expected_hash: str) -> bool:
|
||||
"""Verify a file's hash matches expected."""
|
||||
actual = self.compute_hash(file_path)
|
||||
return actual == expected_hash
|
||||
|
||||
def verify_all(self, collection: str) -> tuple[int, int, list[str]]:
|
||||
"""Verify all files in a collection.
|
||||
|
||||
Returns:
|
||||
(verified_count, mismatch_count, list_of_mismatched_paths)
|
||||
"""
|
||||
collection_dir = self.base_dir / collection
|
||||
if not collection_dir.exists():
|
||||
return 0, 0, []
|
||||
|
||||
verified = 0
|
||||
mismatched = 0
|
||||
bad_paths = []
|
||||
|
||||
for root, _, files in os.walk(collection_dir):
|
||||
for fname in files:
|
||||
if fname.startswith('.') or fname.endswith('.tmp'):
|
||||
continue
|
||||
fpath = os.path.join(root, fname)
|
||||
# Extract expected hash from filename
|
||||
expected_hash = Path(fname).stem
|
||||
actual_hash = self.compute_hash(fpath)
|
||||
if actual_hash == expected_hash:
|
||||
verified += 1
|
||||
else:
|
||||
mismatched += 1
|
||||
bad_paths.append(fpath)
|
||||
|
||||
return verified, mismatched, bad_paths
|
||||
|
||||
@staticmethod
|
||||
def compute_hash(file_path: str) -> str:
|
||||
"""Compute SHA256 hash of a file."""
|
||||
sha256 = hashlib.sha256()
|
||||
with open(file_path, 'rb') as f:
|
||||
for chunk in iter(lambda: f.read(8192), b''):
|
||||
sha256.update(chunk)
|
||||
return sha256.hexdigest()
|
||||
|
||||
@staticmethod
|
||||
def hash_content(content: bytes) -> str:
|
||||
"""Compute SHA256 hash of bytes content."""
|
||||
return hashlib.sha256(content).hexdigest()
|
||||
|
||||
def check_disk_space(self, required_mb: int = 100) -> bool:
|
||||
"""Check if enough disk space is available."""
|
||||
stat = shutil.disk_usage(str(self.base_dir))
|
||||
free_mb = stat.free // (1024 * 1024)
|
||||
if free_mb < required_mb:
|
||||
logger.warning("Low disk space: %d MB free, %d MB required",
|
||||
free_mb, required_mb)
|
||||
return False
|
||||
return True
|
||||
@@ -0,0 +1,45 @@
|
||||
"""Token-bucket rate limiter for HTTP requests."""
|
||||
import time
|
||||
import random
|
||||
import threading
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HTTPRateLimiter:
|
||||
"""Token-bucket rate limiter with jitter for polite crawling."""
|
||||
|
||||
def __init__(self, rate: float = 1.0, jitter: float = 0.5):
|
||||
"""
|
||||
Args:
|
||||
rate: Maximum requests per second
|
||||
jitter: Random jitter range in seconds (±jitter/2)
|
||||
"""
|
||||
self.rate = rate
|
||||
self.jitter = jitter
|
||||
self.min_interval = 1.0 / rate if rate > 0 else 0
|
||||
self._last_request = 0.0
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def wait(self):
|
||||
"""Block until the next request is allowed."""
|
||||
with self._lock:
|
||||
now = time.monotonic()
|
||||
elapsed = now - self._last_request
|
||||
wait_time = self.min_interval - elapsed
|
||||
|
||||
if self.jitter > 0:
|
||||
wait_time += random.uniform(-self.jitter / 2, self.jitter / 2)
|
||||
|
||||
if wait_time > 0:
|
||||
logger.debug("Rate limiting: waiting %.2fs", wait_time)
|
||||
time.sleep(wait_time)
|
||||
|
||||
self._last_request = time.monotonic()
|
||||
|
||||
def update_rate(self, rate: float):
|
||||
"""Update the rate limit (e.g., from profile config)."""
|
||||
with self._lock:
|
||||
self.rate = rate
|
||||
self.min_interval = 1.0 / rate if rate > 0 else 0
|
||||
@@ -0,0 +1,141 @@
|
||||
"""Input sanitization for Mosaic — filenames, paths, content types."""
|
||||
import os
|
||||
import re
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Allowed characters in filenames
|
||||
FILENAME_SAFE = re.compile(r'[^a-zA-Z0-9._-]')
|
||||
# Maximum filename length
|
||||
MAX_FILENAME_LENGTH = 200
|
||||
# ANSI escape sequence pattern
|
||||
ANSI_ESCAPE = re.compile(r'\x1b\[[0-9;]*[a-zA-Z]|\x1b\].*?\x07')
|
||||
|
||||
# Magic bytes → MIME type mapping
|
||||
MAGIC_BYTES = {
|
||||
b'%PDF': 'application/pdf',
|
||||
b'\x89PNG': 'image/png',
|
||||
b'\xff\xd8\xff': 'image/jpeg',
|
||||
b'GIF87a': 'image/gif',
|
||||
b'GIF89a': 'image/gif',
|
||||
b'PK\x03\x04': 'application/zip',
|
||||
b'<!DOCTYPE': 'text/html',
|
||||
b'<html': 'text/html',
|
||||
b'<HTML': 'text/html',
|
||||
}
|
||||
|
||||
|
||||
def sanitize_filename(filename: str) -> str:
|
||||
"""Sanitize a filename to safe characters only.
|
||||
|
||||
- Strip to basename (no directory traversal)
|
||||
- Replace unsafe characters with underscore
|
||||
- Limit length
|
||||
"""
|
||||
# Strip to basename
|
||||
filename = os.path.basename(filename)
|
||||
# Replace unsafe characters
|
||||
sanitized = FILENAME_SAFE.sub('_', filename)
|
||||
# Collapse multiple underscores
|
||||
sanitized = re.sub(r'_+', '_', sanitized)
|
||||
# Limit length
|
||||
if len(sanitized) > MAX_FILENAME_LENGTH:
|
||||
name, ext = os.path.splitext(sanitized)
|
||||
sanitized = name[:MAX_FILENAME_LENGTH - len(ext)] + ext
|
||||
# Ensure non-empty
|
||||
if not sanitized or sanitized == '_':
|
||||
sanitized = 'unnamed'
|
||||
return sanitized
|
||||
|
||||
|
||||
def sanitize_path(path: str, base_dir: str) -> Optional[str]:
|
||||
"""Validate a path stays within base_dir (path traversal prevention).
|
||||
|
||||
Returns the resolved path if safe, None if path escapes base_dir.
|
||||
"""
|
||||
base = os.path.realpath(base_dir)
|
||||
resolved = os.path.realpath(os.path.join(base, path))
|
||||
if resolved.startswith(base + os.sep) or resolved == base:
|
||||
return resolved
|
||||
logger.warning("Path traversal blocked: %s → %s (base: %s)", path, resolved, base)
|
||||
return None
|
||||
|
||||
|
||||
def validate_content_type(file_path: str, expected_type: Optional[str] = None) -> str:
|
||||
"""Determine content type by magic bytes, not file extension.
|
||||
|
||||
Args:
|
||||
file_path: Path to the file
|
||||
expected_type: Expected MIME type to validate against
|
||||
|
||||
Returns:
|
||||
Detected MIME type string
|
||||
|
||||
Raises:
|
||||
ValueError: If expected_type is provided and doesn't match
|
||||
"""
|
||||
try:
|
||||
import magic
|
||||
detected = magic.from_file(file_path, mime=True)
|
||||
except ImportError:
|
||||
# Fallback to magic bytes check
|
||||
detected = _detect_by_magic_bytes(file_path)
|
||||
|
||||
if expected_type and detected != expected_type:
|
||||
# Allow some flexibility (e.g., text/plain for text/html)
|
||||
if not _types_compatible(detected, expected_type):
|
||||
raise ValueError(
|
||||
f"Content type mismatch: expected {expected_type}, detected {detected}"
|
||||
)
|
||||
return detected
|
||||
|
||||
|
||||
def _detect_by_magic_bytes(file_path: str) -> str:
|
||||
"""Fallback detection using magic bytes."""
|
||||
try:
|
||||
with open(file_path, 'rb') as f:
|
||||
header = f.read(16)
|
||||
except (IOError, OSError):
|
||||
return 'application/octet-stream'
|
||||
|
||||
for magic, mime in MAGIC_BYTES.items():
|
||||
if header.startswith(magic):
|
||||
return mime
|
||||
|
||||
# Check if it looks like text
|
||||
try:
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
f.read(1024)
|
||||
return 'text/plain'
|
||||
except (UnicodeDecodeError, IOError):
|
||||
return 'application/octet-stream'
|
||||
|
||||
|
||||
def _types_compatible(detected: str, expected: str) -> bool:
|
||||
"""Check if two MIME types are compatible."""
|
||||
# Same type
|
||||
if detected == expected:
|
||||
return True
|
||||
# text/plain is compatible with text/*
|
||||
if detected == 'text/plain' and expected.startswith('text/'):
|
||||
return True
|
||||
# text/html variants
|
||||
if detected.startswith('text/html') and expected.startswith('text/html'):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def strip_ansi(text: str) -> str:
|
||||
"""Strip ANSI escape sequences from text."""
|
||||
return ANSI_ESCAPE.sub('', text)
|
||||
|
||||
|
||||
def sanitize_for_display(text: str, max_length: int = 0) -> str:
|
||||
"""Sanitize text for safe Rich display — strip ANSI, optionally truncate."""
|
||||
clean = strip_ansi(text)
|
||||
if max_length > 0 and len(clean) > max_length:
|
||||
clean = clean[:max_length] + "..."
|
||||
return clean
|
||||
@@ -0,0 +1,253 @@
|
||||
"""Text chunking, cleaning, and token estimation for Mosaic."""
|
||||
import re
|
||||
import hashlib
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Rough token estimation: ~4 chars per token for English text
|
||||
CHARS_PER_TOKEN = 4
|
||||
|
||||
# Default chunk sizes
|
||||
DEFAULT_CHUNK_TOKENS = 8000
|
||||
DEFAULT_OVERLAP_TOKENS = 500
|
||||
EMAIL_MAX_TOKENS = 4000
|
||||
PDF_MAX_TOKENS = 8000
|
||||
|
||||
|
||||
def estimate_tokens(text: str) -> int:
|
||||
"""Estimate token count from text length."""
|
||||
return len(text) // CHARS_PER_TOKEN
|
||||
|
||||
|
||||
def clean_text(text: str) -> str:
|
||||
"""Clean extracted text: normalize whitespace, remove control chars."""
|
||||
if not text:
|
||||
return ""
|
||||
# Remove null bytes and other control characters (keep newlines, tabs)
|
||||
text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]', '', text)
|
||||
# Normalize line endings
|
||||
text = text.replace('\r\n', '\n').replace('\r', '\n')
|
||||
# Collapse excessive blank lines (3+ → 2)
|
||||
text = re.sub(r'\n{3,}', '\n\n', text)
|
||||
# Strip trailing whitespace per line
|
||||
text = '\n'.join(line.rstrip() for line in text.split('\n'))
|
||||
return text.strip()
|
||||
|
||||
|
||||
def chunk_text(text: str, doc_type: str = "text",
|
||||
max_tokens: int = DEFAULT_CHUNK_TOKENS,
|
||||
overlap_tokens: int = DEFAULT_OVERLAP_TOKENS,
|
||||
metadata: Optional[dict] = None) -> list[dict]:
|
||||
"""Chunk text using per-document-type strategy.
|
||||
|
||||
Args:
|
||||
text: The text to chunk
|
||||
doc_type: Document type (html, pdf, email, cable, text)
|
||||
max_tokens: Maximum tokens per chunk
|
||||
overlap_tokens: Token overlap between chunks
|
||||
metadata: Optional metadata (e.g., cable headers)
|
||||
|
||||
Returns:
|
||||
List of dicts with 'text', 'chunk_index', 'total_chunks', 'metadata'
|
||||
"""
|
||||
if not text:
|
||||
return []
|
||||
|
||||
if doc_type == 'cable':
|
||||
return _chunk_cable(text, max_tokens, overlap_tokens, metadata)
|
||||
elif doc_type == 'pdf':
|
||||
return _chunk_pdf(text, max_tokens, overlap_tokens)
|
||||
elif doc_type == 'email':
|
||||
return _chunk_email(text, max_tokens, overlap_tokens, metadata)
|
||||
elif doc_type == 'html':
|
||||
return _chunk_html(text, max_tokens, overlap_tokens)
|
||||
else:
|
||||
return _chunk_default(text, max_tokens, overlap_tokens)
|
||||
|
||||
|
||||
def _chunk_cable(text: str, max_tokens: int, overlap_tokens: int,
|
||||
metadata: Optional[dict] = None) -> list[dict]:
|
||||
"""Chunk cable: headers as structured context, body chunked independently."""
|
||||
headers = metadata or {}
|
||||
header_text = ""
|
||||
if headers:
|
||||
header_lines = []
|
||||
for key in ('SUBJECT', 'ORIGIN', 'CLASSIFICATION', 'TAGS', 'DATE'):
|
||||
if key in headers:
|
||||
header_lines.append(f"{key}: {headers[key]}")
|
||||
header_text = '\n'.join(header_lines)
|
||||
|
||||
# If small enough, return as one chunk
|
||||
total_tokens = estimate_tokens(header_text + '\n\n' + text)
|
||||
if total_tokens <= max_tokens:
|
||||
return [{'text': (header_text + '\n\n' + text).strip(),
|
||||
'chunk_index': 0, 'total_chunks': 1, 'metadata': headers}]
|
||||
|
||||
# Chunk body, prepend headers to each chunk
|
||||
body_chunks = _chunk_default(text, max_tokens - estimate_tokens(header_text) - 50,
|
||||
overlap_tokens)
|
||||
for i, chunk in enumerate(body_chunks):
|
||||
chunk['text'] = header_text + '\n\n' + chunk['text']
|
||||
chunk['metadata'] = headers
|
||||
chunk['chunk_index'] = i
|
||||
chunk['total_chunks'] = len(body_chunks)
|
||||
return body_chunks
|
||||
|
||||
|
||||
def _chunk_pdf(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]:
|
||||
"""Chunk PDF by page boundaries (indicated by form feeds or page markers)."""
|
||||
# Split by form feed or common page break markers
|
||||
pages = re.split(r'\f|\n-{3,}\n|\n={3,}\n', text)
|
||||
pages = [p.strip() for p in pages if p.strip()]
|
||||
|
||||
if not pages:
|
||||
return _chunk_default(text, max_tokens, overlap_tokens)
|
||||
|
||||
chunks = []
|
||||
current_text = ""
|
||||
for page in pages:
|
||||
combined_tokens = estimate_tokens(current_text + '\n\n' + page)
|
||||
if combined_tokens > max_tokens and current_text:
|
||||
chunks.append(current_text.strip())
|
||||
# Keep overlap from end of previous chunk
|
||||
overlap_chars = overlap_tokens * CHARS_PER_TOKEN
|
||||
current_text = current_text[-overlap_chars:] + '\n\n' + page
|
||||
else:
|
||||
current_text = (current_text + '\n\n' + page).strip()
|
||||
|
||||
if current_text.strip():
|
||||
chunks.append(current_text.strip())
|
||||
|
||||
return [{'text': c, 'chunk_index': i, 'total_chunks': len(chunks), 'metadata': {}}
|
||||
for i, c in enumerate(chunks)]
|
||||
|
||||
|
||||
def _chunk_email(text: str, max_tokens: int, overlap_tokens: int,
|
||||
metadata: Optional[dict] = None) -> list[dict]:
|
||||
"""Chunk email: headers + body as atomic unit unless body is too large."""
|
||||
headers = metadata or {}
|
||||
total_tokens = estimate_tokens(text)
|
||||
|
||||
if total_tokens <= EMAIL_MAX_TOKENS:
|
||||
return [{'text': text, 'chunk_index': 0, 'total_chunks': 1, 'metadata': headers}]
|
||||
|
||||
# Split headers from body
|
||||
header_section = ""
|
||||
body = text
|
||||
if '\n\n' in text:
|
||||
parts = text.split('\n\n', 1)
|
||||
header_section = parts[0]
|
||||
body = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
# Chunk body, prepend headers
|
||||
body_chunks = _chunk_default(body, max_tokens - estimate_tokens(header_section) - 50,
|
||||
overlap_tokens)
|
||||
for i, chunk in enumerate(body_chunks):
|
||||
chunk['text'] = header_section + '\n\n' + chunk['text']
|
||||
chunk['metadata'] = headers
|
||||
chunk['chunk_index'] = i
|
||||
chunk['total_chunks'] = len(body_chunks)
|
||||
return body_chunks
|
||||
|
||||
|
||||
def _chunk_html(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]:
|
||||
"""Chunk HTML: section-aware chunking preserving heading hierarchy."""
|
||||
# Split on headings
|
||||
sections = re.split(r'(?=\n#{1,6}\s|\n[A-Z][A-Za-z\s]{3,}\n[=-]+\n)', text)
|
||||
sections = [s for s in sections if s.strip()]
|
||||
|
||||
if not sections or len(sections) == 1:
|
||||
return _chunk_default(text, max_tokens, overlap_tokens)
|
||||
|
||||
chunks = []
|
||||
current_text = ""
|
||||
for section in sections:
|
||||
combined_tokens = estimate_tokens(current_text + '\n\n' + section)
|
||||
if combined_tokens > max_tokens and current_text:
|
||||
chunks.append(current_text.strip())
|
||||
overlap_chars = overlap_tokens * CHARS_PER_TOKEN
|
||||
current_text = current_text[-overlap_chars:] + '\n\n' + section
|
||||
else:
|
||||
current_text = (current_text + '\n\n' + section).strip()
|
||||
|
||||
if current_text.strip():
|
||||
chunks.append(current_text.strip())
|
||||
|
||||
return [{'text': c, 'chunk_index': i, 'total_chunks': len(chunks), 'metadata': {}}
|
||||
for i, c in enumerate(chunks)]
|
||||
|
||||
|
||||
def _chunk_default(text: str, max_tokens: int, overlap_tokens: int) -> list[dict]:
|
||||
"""Default chunking: paragraph-aware with token overlap."""
|
||||
total_tokens = estimate_tokens(text)
|
||||
if total_tokens <= max_tokens:
|
||||
return [{'text': text, 'chunk_index': 0, 'total_chunks': 1, 'metadata': {}}]
|
||||
|
||||
max_chars = max_tokens * CHARS_PER_TOKEN
|
||||
overlap_chars = overlap_tokens * CHARS_PER_TOKEN
|
||||
|
||||
# Split into paragraphs
|
||||
paragraphs = text.split('\n\n')
|
||||
chunks = []
|
||||
current = ""
|
||||
|
||||
for para in paragraphs:
|
||||
if len(current) + len(para) + 2 > max_chars and current:
|
||||
chunks.append(current.strip())
|
||||
# Keep overlap
|
||||
current = current[-overlap_chars:] + '\n\n' + para
|
||||
else:
|
||||
current = (current + '\n\n' + para) if current else para
|
||||
|
||||
if current.strip():
|
||||
chunks.append(current.strip())
|
||||
|
||||
# If still too large (single giant paragraph), force-split
|
||||
final_chunks = []
|
||||
for chunk in chunks:
|
||||
if len(chunk) > max_chars * 1.5:
|
||||
# Force split at sentence boundaries or character limit
|
||||
for sub in _force_split(chunk, max_chars, overlap_chars):
|
||||
final_chunks.append(sub)
|
||||
else:
|
||||
final_chunks.append(chunk)
|
||||
|
||||
return [{'text': c, 'chunk_index': i, 'total_chunks': len(final_chunks), 'metadata': {}}
|
||||
for i, c in enumerate(final_chunks)]
|
||||
|
||||
|
||||
def _force_split(text: str, max_chars: int, overlap_chars: int) -> list[str]:
|
||||
"""Force-split text at sentence boundaries or character limits."""
|
||||
sentences = re.split(r'(?<=[.!?])\s+', text)
|
||||
chunks = []
|
||||
current = ""
|
||||
|
||||
for sent in sentences:
|
||||
if len(current) + len(sent) + 1 > max_chars and current:
|
||||
chunks.append(current.strip())
|
||||
current = current[-overlap_chars:] + ' ' + sent
|
||||
else:
|
||||
current = (current + ' ' + sent) if current else sent
|
||||
|
||||
if current.strip():
|
||||
chunks.append(current.strip())
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
def dedup_text(texts: list[str], threshold: float = 0.95) -> list[str]:
|
||||
"""Remove near-duplicate texts using hash comparison.
|
||||
|
||||
For exact duplicates. Near-duplicate detection would require
|
||||
more sophisticated approaches (minhash, simhash, etc.)
|
||||
"""
|
||||
seen = set()
|
||||
unique = []
|
||||
for text in texts:
|
||||
h = hashlib.md5(text.encode()).hexdigest()
|
||||
if h not in seen:
|
||||
seen.add(h)
|
||||
unique.append(text)
|
||||
return unique
|
||||
Reference in New Issue
Block a user