chore(api): add auto-cleanup job for old vector collections (#99)

backend/app/database.py

@@ -44,14 +44,29 @@ def _migrate_schema():
     for non-destructive changes such as new nullable columns.
     """
     inspector = inspect(engine)
-    existing_columns = {c["name"] for c in inspector.get_columns("users")}
-
-    migrations = [
+    # Migrate users
+    existing_users_columns = {c["name"] for c in inspector.get_columns("users")}
+    users_migrations = [
         ("users", "hf_token", "ALTER TABLE users ADD COLUMN hf_token VARCHAR(255)"),
     ]
+    for table, column, ddl in users_migrations:
+        if column not in existing_users_columns:
+            try:
+                with engine.begin() as conn:
+                    conn.execute(text(ddl))
+                logger.info("Migration: added column %s.%s", table, column)
+            except Exception:
+                logger.warning(
+                    "Migration skipped (may already exist): %s.%s", table, column
+                )
 
-    for table, column, ddl in migrations:
-        if column not in existing_columns:
+    # Migrate documents
+    existing_docs_columns = {c["name"] for c in inspector.get_columns("documents")}
+    docs_migrations = [
+        ("documents", "last_accessed_at", "ALTER TABLE documents ADD COLUMN last_accessed_at TIMESTAMP"),
+    ]
+    for table, column, ddl in docs_migrations:
+        if column not in existing_docs_columns:
             try:
                 with engine.begin() as conn:
                     conn.execute(text(ddl))

backend/app/main.py

@@ -31,6 +31,63 @@ logger = logging.getLogger(__name__)
 settings = get_settings()
 
 
+async def document_cleanup_job():
+    """Background loop to periodically purge documents not accessed in 30 days."""
+    import asyncio
+    from datetime import datetime, timedelta, timezone
+    logger.info("Starting document cleanup background job loop")
+    while True:
+        try:
+            from app.database import SessionLocal
+            from app.models import Document
+            from app.rag.vectorstore import delete_document_chunks
+            from sqlalchemy import or_
+            
+            db = SessionLocal()
+            try:
+                cutoff = datetime.now(timezone.utc) - timedelta(days=30)
+                expired_docs = db.query(Document).filter(
+                    or_(
+                        Document.last_accessed_at < cutoff,
+                        Document.last_accessed_at.is_(None) & (Document.uploaded_at < cutoff)
+                    )
+                ).all()
+                
+                for doc in expired_docs:
+                    logger.info(f"Auto-cleanup: Purging document {doc.id} ('{doc.original_name}') due to inactivity since {doc.last_accessed_at or doc.uploaded_at}")
+                    
+                    # Delete physical file
+                    filepath = os.path.join(settings.UPLOAD_DIR, doc.user_id, doc.filename)
+                    if os.path.exists(filepath):
+                        try:
+                            os.remove(filepath)
+                        except Exception as e:
+                            logger.warning(f"Auto-cleanup: Failed to delete physical file {filepath}: {e}")
+                    
+                    # Delete vectors
+                    try:
+                        delete_document_chunks(document_id=doc.id, user_id=doc.user_id)
+                    except Exception as e:
+                        logger.warning(f"Auto-cleanup: Error deleting vectors for document {doc.id}: {e}")
+                    
+                    # Delete database record
+                    db.delete(doc)
+                
+                db.commit()
+                if expired_docs:
+                    logger.info(f"Auto-cleanup: Purged {len(expired_docs)} documents.")
+            except Exception as exc:
+                logger.error(f"Auto-cleanup job encountered error: {exc}", exc_info=True)
+            finally:
+                db.close()
+                
+        except Exception as e:
+            logger.error(f"Error in document cleanup background loop: {e}", exc_info=True)
+            
+        # Run every 24 hours (86400 seconds)
+        await asyncio.sleep(86400)
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Application startup/shutdown lifecycle."""
@@ -53,10 +110,21 @@ async def lifespan(app: FastAPI):
     except Exception as e:
         logger.warning(f"Failed to pre-load embedding model: {e}")
 
+    # Start background cleanup task
+    import asyncio
+    cleanup_task = asyncio.create_task(document_cleanup_job())
+
     yield
 
     # ── Shutdown ─────────────────────────────────────
     logger.info("Shutting down")
+    cleanup_task.cancel()
+    try:
+        await cleanup_task
+    except asyncio.CancelledError:
+        pass
+    except Exception as e:
+        logger.warning(f"Error cancelling cleanup task: {e}")
 
 
 # ── Create App ───────────────────────────────────────

backend/app/models.py

@@ -90,6 +90,7 @@ class Document(Base):
     status = Column(String(20), default="pending")         # pending | processing | ready | failed
     error_message = Column(Text, nullable=True)
     uploaded_at = Column(DateTime, default=lambda: datetime.now(timezone.utc))
+    last_accessed_at = Column(DateTime, default=lambda: datetime.now(timezone.utc), nullable=True)
     summary = Column(Text, nullable=True)  # Optional summary of the document's content
 
     # Relationships

backend/app/routes/chat.py

@@ -146,6 +146,10 @@ def ask_question(
                     status_code=400,
                     detail=f"Document is still {doc.status}. Please wait for processing to complete.",
                 )
+            
+            # Update last_accessed_at timestamp
+            doc.last_accessed_at = datetime.now(timezone.utc)
+            db.commit()
 
         result = generate_answer(
             question=payload.question,
@@ -224,6 +228,10 @@ def ask_question_stream(
                 status_code=400,
                 detail=f"Document is still {doc.status}. Please wait for processing to complete.",
             )
+        
+        # Update last_accessed_at timestamp
+        doc.last_accessed_at = datetime.now(timezone.utc)
+        db.commit()
 
     started_at = time.perf_counter()
 

backend/scripts/document_cleanup.py

@@ -0,0 +1,137 @@
+"""
+Cleanup script for ChromaDB vectors and records of documents inactive for 30 days.
+
+By default, a document is considered inactive if it has not been accessed (last_accessed_at)
+or uploaded (if last_accessed_at is missing) for 30 days.
+
+Run manually:
+    python backend/scripts/document_cleanup.py
+
+Environment:
+    DOCUMENT_CLEANUP_INACTIVE_DAYS=30
+    DOCUMENT_CLEANUP_DRY_RUN=true
+"""
+from __future__ import annotations
+
+import logging
+import os
+import sys
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+
+from sqlalchemy import or_, inspect, text
+
+# Allow running this file directly from the repository root.
+BACKEND_DIR = Path(__file__).resolve().parents[1]
+if str(BACKEND_DIR) not in sys.path:
+    sys.path.insert(0, str(BACKEND_DIR))
+
+from app.database import SessionLocal  # noqa: E402
+from app.models import Document  # noqa: E402
+from app.rag.vectorstore import delete_document_chunks  # noqa: E402
+from app.config import get_settings  # noqa: E402
+
+logger = logging.getLogger("document_cleanup")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
+
+settings = get_settings()
+
+def _env_bool(name: str, default: bool = False) -> bool:
+    value = os.getenv(name)
+    if value is None:
+        return default
+    return value.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def ensure_last_accessed_at_column() -> None:
+    """Ensure last_accessed_at column exists in database (handles SQLite local installs)."""
+    db = SessionLocal()
+    try:
+        bind = db.get_bind()
+        inspector = inspect(bind)
+        columns = {column["name"] for column in inspector.get_columns("documents")}
+        if "last_accessed_at" not in columns:
+            logger.info("Adding missing documents.last_accessed_at column")
+            db.execute(text("ALTER TABLE documents ADD COLUMN last_accessed_at TIMESTAMP"))
+            db.commit()
+    finally:
+        db.close()
+
+
+def cleanup_inactive_documents(
+    inactive_days: int | None = None,
+    dry_run: bool | None = None,
+) -> dict[str, int]:
+    """Delete database records, physical files, and Chroma collections for inactive documents."""
+    ensure_last_accessed_at_column()
+
+    days = inactive_days or int(os.getenv("DOCUMENT_CLEANUP_INACTIVE_DAYS", "30"))
+    is_dry_run = _env_bool("DOCUMENT_CLEANUP_DRY_RUN", False) if dry_run is None else dry_run
+    cutoff = datetime.now(timezone.utc) - timedelta(days=days)
+
+    stats = {
+        "scanned": 0,
+        "eligible": 0,
+        "deleted": 0,
+        "failed": 0,
+    }
+
+    db = SessionLocal()
+    try:
+        # Select documents whose last_accessed_at or uploaded_at is before cutoff
+        docs = db.query(Document).filter(
+            or_(
+                Document.last_accessed_at < cutoff,
+                Document.last_accessed_at.is_(None) & (Document.uploaded_at < cutoff)
+            )
+        ).all()
+
+        for doc in docs:
+            stats["scanned"] += 1
+            last_activity = doc.last_accessed_at or doc.uploaded_at
+            
+            stats["eligible"] += 1
+            logger.info(
+                "Document %s ('%s') inactive since %s; purging dry_run=%s",
+                doc.id,
+                doc.original_name,
+                last_activity,
+                is_dry_run,
+            )
+
+            if is_dry_run:
+                continue
+
+            try:
+                # 1. Delete file from disk
+                filepath = os.path.join(settings.UPLOAD_DIR, doc.user_id, doc.filename)
+                if os.path.exists(filepath):
+                    os.remove(filepath)
+                    logger.info("Deleted physical file: %s", filepath)
+
+                # 2. Delete vectors from ChromaDB
+                delete_document_chunks(document_id=doc.id, user_id=doc.user_id)
+
+                # 3. Delete from SQL database
+                db.delete(doc)
+                stats["deleted"] += 1
+            except Exception as exc:
+                stats["failed"] += 1
+                logger.warning(
+                    "Failed purging document %s: %s",
+                    doc.id,
+                    exc,
+                    exc_info=True,
+                )
+
+        if not is_dry_run:
+            db.commit()
+
+        logger.info("Document cleanup complete: %s", stats)
+        return stats
+    finally:
+        db.close()
+
+
+if __name__ == "__main__":
+    cleanup_inactive_documents()