helpers/lifespan_helper.py

import asyncio
import logging

import psutil

from classes.ocr_reader import OCRReader
from classes.pii_filter import PIIFilter
from classes.session_conversation_store import SessionConversationStore
from classes.session_document_store import SessionDocumentStore
from classes.session_tracker import SessionTracker
from constants import MAX_RAM_USAGE_PERCENT


logger = logging.getLogger("uvicorn")


def run_cleanup(
    session_tracker: SessionTracker,
    session_document_store: SessionDocumentStore,
    session_conversation_store: SessionConversationStore,
):
    logger.info("Running cleanup")
    deleted_session_ids = session_tracker.delete_inactive_sessions()
    if len(deleted_session_ids) > 0:
        logger.info(f"{len(deleted_session_ids)} inactive sessions will be deleted.")
    for session_id in deleted_session_ids:
        session_document_store.delete_session_documents(session_id)
        session_conversation_store.delete_session_conversations(session_id)

    while psutil.virtual_memory().percent > MAX_RAM_USAGE_PERCENT:
        oldest_session_id = session_tracker.delete_oldest_session()
        logger.info(f"Deleting {oldest_session_id} session because of high RAM usage")
        if oldest_session_id is None:
            break
        session_document_store.delete_session_documents(oldest_session_id)
        session_conversation_store.delete_session_conversations(oldest_session_id)


async def cleanup_loop(
    session_tracker: SessionTracker,
    session_document_store: SessionDocumentStore,
    session_conversation_store: SessionConversationStore,
):
    """Run the 4-hour cleanup check every 10 minutes."""
    while True:
        await asyncio.sleep(600)  # Wait 10 minutes
        run_cleanup(session_tracker, session_document_store, session_conversation_store)


def load_heavy_models():
    OCRReader()
    PIIFilter()