Spaces:
Sleeping
Sleeping
File size: 3,854 Bytes
1c223a0 eea0a57 1c223a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | from __future__ import annotations
import re
from pathlib import Path
from uuid import uuid4
from rag_chunk_visualizer.core.config import settings
from rag_chunk_visualizer.core.logging import get_logger
from rag_chunk_visualizer.models.document import RawDocument
logger = get_logger(__name__)
ALLOWED_EXTENSIONS = {"txt", "md"}
class DocumentValidationError(ValueError):
pass
def sanitize_filename(filename: str) -> str:
sanitized = re.sub(r"[^A-Za-z0-9._-]+", "_", filename).strip("._")
return sanitized or "document.txt"
def get_extension(filename: str) -> str:
return Path(filename).suffix.lower().lstrip(".")
def validate_uploaded_file(filename: str, size_bytes: int) -> None:
extension = get_extension(filename)
if extension not in ALLOWED_EXTENSIONS:
raise DocumentValidationError(
f"Unsupported file type for '{filename}'. Allowed types: .txt, .md"
)
max_bytes = settings.max_upload_mb * 1024 * 1024
if size_bytes > max_bytes:
raise DocumentValidationError(
f"File '{filename}' is too large ({size_bytes} bytes). "
f"Maximum allowed size is {settings.max_upload_mb} MB."
)
def decode_text_bytes(raw_bytes: bytes, filename: str) -> str:
encodings_to_try = ("utf-8-sig", "utf-8")
for encoding in encodings_to_try:
try:
return raw_bytes.decode(encoding)
except UnicodeDecodeError:
continue
raise DocumentValidationError(
f"Could not decode '{filename}' as UTF-8 text. "
"Please upload a UTF-8 encoded .txt or .md file."
)
def normalize_text(text: str) -> str:
text = text.replace("\r\n", "\n").replace("\r", "\n")
text = re.sub(r"\n{3,}", "\n\n", text)
return text.strip()
def save_uploaded_file(raw_bytes: bytes, doc_id: str, filename: str) -> Path:
safe_name = sanitize_filename(filename)
# Recreate the target dir defensively: on ephemeral hosts (e.g. HF Spaces)
# the data directory can be wiped between requests.
settings.documents_dir.mkdir(parents=True, exist_ok=True)
output_path = settings.documents_dir / f"{doc_id}_{safe_name}"
output_path.write_bytes(raw_bytes)
return output_path
def build_raw_document(filename: str, raw_bytes: bytes) -> RawDocument:
validate_uploaded_file(filename, len(raw_bytes))
decoded_text = decode_text_bytes(raw_bytes, filename)
normalized_text = normalize_text(decoded_text)
if not normalized_text:
raise DocumentValidationError(f"File '{filename}' is empty after text extraction.")
doc_id = uuid4().hex[:8]
saved_path = save_uploaded_file(raw_bytes, doc_id, filename)
document = RawDocument(
doc_id=doc_id,
filename=filename,
extension=get_extension(filename),
source_path=str(saved_path),
char_count=len(normalized_text),
word_count=len(normalized_text.split()),
text=normalized_text,
)
logger.info(
"Built raw document | doc_id=%s | filename=%s | chars=%s | words=%s",
document.doc_id,
document.filename,
document.char_count,
document.word_count,
)
return document
def process_uploaded_files(uploaded_files: list) -> tuple[list[dict], list[str]]:
processed_documents: list[dict] = []
errors: list[str] = []
for uploaded_file in uploaded_files:
filename = uploaded_file.name
raw_bytes = uploaded_file.getvalue()
try:
document = build_raw_document(filename=filename, raw_bytes=raw_bytes)
processed_documents.append(document.to_dict())
except DocumentValidationError as exc:
logger.warning("Document validation failed | filename=%s | error=%s", filename, exc)
errors.append(str(exc))
return processed_documents, errors
|