"""Reference document persistence, delete, and re-ingest.""" from __future__ import annotations import logging import shutil import threading import time from datetime import UTC, datetime from pathlib import Path from backend.config import settings from backend.core import ingest from backend.core.rag_store import TIER_REFERENCE, get_rag_store from backend.core.report_session import ( UploadedDocument, delete_document, get_document, list_documents, save_document, ) from backend.core.style_profile import invalidate_style_profile from backend.utils import tenant_store logger = logging.getLogger(__name__) _reingest_lock = threading.Lock() _reingest_running: set[str] = set() def is_reingest_running(tenant_id: str) -> bool: with _reingest_lock: return tenant_id in _reingest_running def recover_stale_processing_documents(tenant_id: str) -> int: """Reset orphaned ``processing`` rows when no re-ingest worker is active.""" if is_reingest_running(tenant_id): return 0 recovered = 0 for _doc_id, doc in list_documents(tenant_id).items(): if doc.status == "processing": doc.status = "complete" doc.error = None save_document(tenant_id, doc) recovered += 1 if recovered: logger.info( "Recovered %d stale processing document(s) for tenant=%s", recovered, tenant_id, ) return recovered def recover_all_tenants_stale_processing() -> int: """On startup, clear processing flags left by a killed background worker.""" tenants_root = settings.data_dir_path / "tenants" if not tenants_root.is_dir(): return 0 total = 0 for tenant_dir in sorted(tenants_root.iterdir()): if not tenant_dir.is_dir(): continue if not (tenant_dir / "compat_documents.json").is_file(): continue total += recover_stale_processing_documents(tenant_dir.name) return total def reingest_progress(tenant_id: str) -> dict: docs = list_documents(tenant_id) counts = {"complete": 0, "processing": 0, "failed": 0, "pending": 0} for doc in docs.values(): key = doc.status if doc.status in counts else "pending" counts[key] = counts.get(key, 0) + 1 return { "total": len(docs), "running": is_reingest_running(tenant_id), **counts, } def document_created_at_iso(doc: UploadedDocument) -> str: """Serialize ``created_at`` for API responses (ISO-8601 UTC).""" ts = doc.created_at if isinstance(ts, (int, float)) and ts > 0: return datetime.fromtimestamp(ts, tz=UTC).isoformat() if isinstance(ts, str) and ts.strip(): return ts.strip() return datetime.now(tz=UTC).isoformat() def persist_reference_file( tenant_id: str, document_id: str, source_path: Path, *, original_filename: str, ) -> Path: """Copy an uploaded file into tenant storage for later re-ingest.""" suffix = source_path.suffix.lower() or Path(original_filename).suffix.lower() dest = tenant_store.reference_upload_path(tenant_id, document_id, suffix) shutil.copy2(source_path, dest) return dest def ingest_and_register( tenant_id: str, source_path: Path, *, original_filename: str, document_id: str | None = None, ) -> UploadedDocument: """Ingest a reference file and record it in the document library.""" from backend.core.report_session import new_document_id doc_id = document_id or new_document_id() stored = persist_reference_file( tenant_id, doc_id, source_path, original_filename=original_filename ) chunks = ingest.ingest_reference(tenant_id, stored) doc = UploadedDocument( document_id=doc_id, filename=original_filename or stored.name, status="complete", ingested_chunks=chunks, storage_path=str(stored), file_size=stored.stat().st_size if stored.is_file() else 0, created_at=time.time(), ) save_document(tenant_id, doc) invalidate_style_profile(tenant_id) return doc def remove_reference_document(tenant_id: str, document_id: str) -> int: """Remove chunks, stored file, and library record. Returns chunks removed.""" doc = get_document(tenant_id, document_id) if doc is None: raise KeyError("Document not found") removed = get_rag_store().remove_document( tenant_id, TIER_REFERENCE, source_filename=doc.filename, doc_id=f"reference:{doc.filename}", ) if doc.storage_path: path = Path(doc.storage_path) if path.is_file(): path.unlink(missing_ok=True) delete_document(tenant_id, document_id) invalidate_style_profile(tenant_id) return removed def reingest_reference_document(tenant_id: str, document_id: str) -> UploadedDocument: """Re-process one stored reference file through the current pipeline.""" doc = get_document(tenant_id, document_id) if doc is None: raise KeyError("Document not found") path = Path(doc.storage_path) if doc.storage_path else None if path is None or not path.is_file(): raise FileNotFoundError("Source file is no longer on disk; cannot re-ingest.") doc.status = "processing" doc.error = None save_document(tenant_id, doc) get_rag_store().remove_document( tenant_id, TIER_REFERENCE, source_filename=doc.filename, doc_id=f"reference:{doc.filename}", ) chunks = ingest.ingest_reference(tenant_id, path) doc.status = "complete" doc.error = None doc.ingested_chunks = chunks doc.file_size = path.stat().st_size save_document(tenant_id, doc) invalidate_style_profile(tenant_id) return doc def reingest_all_documents( tenant_id: str, *, skip_document_ids: set[str] | None = None, ) -> dict: skip = skip_document_ids or set() docs = list_documents(tenant_id) queued: list[str] = [] skipped_missing = 0 skipped_active = 0 for doc_id, doc in docs.items(): if doc_id in skip: skipped_active += 1 continue try: logger.info("Re-ingesting %s for tenant=%s", doc.filename, tenant_id) updated = reingest_reference_document(tenant_id, doc_id) queued.append(doc_id) logger.info( "Re-ingested %s (%d chunks)", updated.filename, updated.ingested_chunks, ) except FileNotFoundError: skipped_missing += 1 except Exception as exc: # noqa: BLE001 doc.status = "failed" doc.error = str(exc) save_document(tenant_id, doc) return { "queued": len(queued), "document_ids": queued, "skipped_active": skipped_active, "skipped_missing_file": skipped_missing, "detail": ( f"Re-ingested {len(queued)} document(s); " f"skipped {skipped_active} blocked and {skipped_missing} missing-file." ), } def schedule_reingest_all_documents( tenant_id: str, *, skip_document_ids: set[str] | None = None, ) -> dict: """Queue a full-library re-ingest on a background thread (non-blocking HTTP). Only the document currently being embedded is marked ``processing`` so a server restart cannot strand the whole library in that state. """ skip = skip_document_ids or set() docs = list_documents(tenant_id) to_queue = [doc_id for doc_id in docs if doc_id not in skip] with _reingest_lock: if tenant_id in _reingest_running: progress = reingest_progress(tenant_id) return { "queued": 0, "document_ids": [], "skipped_active": len(skip), "skipped_missing_file": 0, "reingest_running": True, "progress": progress, "detail": ( f"Re-ingest already running " f"({progress['processing']} processing, " f"{progress['complete']} ready)." ), } _reingest_running.add(tenant_id) # Clear orphaned processing flags from a prior killed worker. recover_stale_processing_documents(tenant_id) def _worker() -> None: try: logger.info( "Background re-ingest started for tenant=%s (%d documents)", tenant_id, len(to_queue), ) reingest_all_documents(tenant_id, skip_document_ids=skip) except Exception: # noqa: BLE001 logger.exception("Background re-ingest failed for tenant=%s", tenant_id) recover_stale_processing_documents(tenant_id) finally: with _reingest_lock: _reingest_running.discard(tenant_id) logger.info("Background re-ingest finished for tenant=%s", tenant_id) threading.Thread( target=_worker, name=f"reingest-{tenant_id}", daemon=True, ).start() progress = reingest_progress(tenant_id) return { "queued": len(to_queue), "document_ids": to_queue, "skipped_active": len(skip), "skipped_missing_file": 0, "reingest_running": True, "progress": progress, "detail": ( f"Re-ingest started in the background for {len(to_queue)} document(s). " "Status updates every few seconds as each file completes." ), } def schedule_reingest_reference_document(tenant_id: str, document_id: str) -> dict: """Queue a single-document re-ingest on a background thread.""" doc = get_document(tenant_id, document_id) if doc is None: raise KeyError("Document not found") path = Path(doc.storage_path) if doc.storage_path else None if path is None or not path.is_file(): raise FileNotFoundError("Source file is no longer on disk; cannot re-ingest.") doc.status = "processing" doc.error = None save_document(tenant_id, doc) def _worker() -> None: try: reingest_reference_document(tenant_id, document_id) except Exception as exc: # noqa: BLE001 failed = get_document(tenant_id, document_id) if failed is not None: failed.status = "failed" failed.error = str(exc) save_document(tenant_id, failed) threading.Thread( target=_worker, name=f"reingest-{tenant_id}-{document_id[:8]}", daemon=True, ).start() return { "queued": 1, "document_ids": [document_id], "detail": f"Re-ingest started in the background for {doc.filename}.", }