Spaces:
Sleeping
Sleeping
| """Reading one upload for the live demo: the shared upload checks, the pipeline's own | |
| ingestion (layoutlmv3 tables, Qwen2.5 VL figures), the limits, and the per address daily cap. No Gradio here, so | |
| it runs and is tested anywhere; app.py wraps it for the Space.""" | |
| import contextlib | |
| import hashlib | |
| import io | |
| import ipaddress | |
| import os | |
| import shutil | |
| import tempfile | |
| import threading | |
| import time | |
| from typing import Dict, List, Optional, Tuple | |
| from src.ingestion import upload_checks | |
| from src.ingestion.upload_checks import UploadRefused | |
| DAY_S = 24 * 3600 | |
| def limit(name: str, default: int) -> int: | |
| """A Space variable (MAX_MB, MAX_PAGES, MAX_CHUNKS, DAILY_READS_PER_IP), else the default.""" | |
| return int(os.environ.get(name, default)) | |
| def client_ip(forwarded_for: Optional[str], peer: Optional[str]) -> str: | |
| """The visitor's address: the rightmost public hop of X-Forwarded-For (what the first proxy we trust saw; the | |
| leftmost hops are whatever the client claimed), else the socket peer.""" | |
| return client_hop(forwarded_for, peer)[0] | |
| def client_hop(forwarded_for: Optional[str], peer: Optional[str]) -> Tuple[str, str]: | |
| """(address, where it came from, as "hop k of n" or "peer"); the second part is safe to log, the first is not.""" | |
| hops = [h.strip() for h in (forwarded_for or "").split(",") if h.strip()] | |
| for k in range(len(hops) - 1, -1, -1): | |
| try: | |
| if ipaddress.ip_address(hops[k]).is_global: | |
| return hops[k], f"hop {k + 1} of {len(hops)}" | |
| except ValueError: | |
| continue | |
| return peer or "unknown", "peer" | |
| class DailyCap: | |
| """Successful reads per address in a rolling 24 hours, in memory (lost on restart); safe across threads.""" | |
| def __init__(self): | |
| self._reads: Dict[str, Tuple[float, int]] = {} | |
| self._lock = threading.Lock() | |
| def allowed(self, ip: str, now: float = None) -> bool: | |
| now = time.time() if now is None else now | |
| with self._lock: | |
| start, count = self._reads.get(ip, (now, 0)) | |
| return now - start >= DAY_S or count < limit("DAILY_READS_PER_IP", 10) | |
| def record(self, ip: str, now: float = None) -> None: | |
| now = time.time() if now is None else now | |
| with self._lock: | |
| start, count = self._reads.get(ip, (now, 0)) | |
| self._reads[ip] = (now, 1) if now - start >= DAY_S else (start, count + 1) | |
| def prepare(source: str, name: str) -> Tuple[str, Optional[int]]: | |
| """Checks the upload and copies it, under its own name, into a fresh private folder: (folder, PDF page count).""" | |
| upload_checks.check_name(name) | |
| max_mb = limit("MAX_MB", 10) | |
| if os.path.getsize(source) > max_mb * 1024 * 1024: | |
| raise UploadRefused("too_large", 413, f"{name} is larger than {max_mb} MB.") | |
| folder = tempfile.mkdtemp(prefix="vrag-") | |
| path = os.path.join(folder, name) | |
| shutil.copyfile(source, path) | |
| try: | |
| pages = upload_checks.check_contents(path, name, limit("MAX_PAGES", 5)) | |
| except Exception: | |
| shutil.rmtree(folder, ignore_errors=True) | |
| raise | |
| return folder, pages | |
| def log_name(name: str) -> str: | |
| """A short fingerprint of a visitor's file name, so the log can tell uploads apart without naming them.""" | |
| return hashlib.sha256((name or "").encode("utf-8")).hexdigest()[:8] | |
| def read(folder: str) -> List[Dict]: | |
| """The pipeline's own reading and chunking of the one file in `folder`, as the demo's chunk records. The | |
| ingestion's own log lines name the file and quote reader errors, so they are discarded here.""" | |
| from src.ingestion.chunking import DocumentChunker | |
| from src.ingestion.loader import DocumentLoader | |
| with contextlib.redirect_stdout(io.StringIO()): | |
| chunks = DocumentChunker().chunk_documents(DocumentLoader(raw_data_dir=folder).load_documents()) | |
| max_chunks = limit("MAX_CHUNKS", 2000) | |
| if len(chunks) > max_chunks: | |
| raise UploadRefused("too_many_chunks", 413, f"The document makes {len(chunks)} chunks; the limit is {max_chunks}.") | |
| return [{"id": c["id"], "doc_hash": c["metadata"].get("doc_hash"), "chunk_index": c["metadata"].get("chunk_index"), | |
| "page": c["metadata"].get("page"), "kind": c["metadata"].get("kind", "text"), "content": c["content"]} | |
| for c in chunks] | |
| def discard(*paths: Optional[str]) -> None: | |
| """Deletes the private copy and Gradio's cached upload; nothing of the visitor's file is kept.""" | |
| for path in paths: | |
| if not path: | |
| continue | |
| if os.path.isdir(path): | |
| shutil.rmtree(path, ignore_errors=True) | |
| elif os.path.exists(path): | |
| try: | |
| os.remove(path) | |
| except OSError: | |
| pass | |