Blame for reader.py

"""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