Spaces:
Running
Running
Enhance API security and functionality by adding authentication middleware and session management. Updated app.py to include the new auth router and integrated authentication checks for protected endpoints. Modified requirements.txt to include necessary libraries for session handling. Updated .env.example to include authentication credentials. Improved retrieval functions with query expansion for better medical term matching and enriched context in responses.
ddc9c77
| """ | |
| Background initialization system for preloading heavy components during app startup. | |
| This module handles the eager loading of embedding models, retrievers, and chunks | |
| to improve first-question response time. | |
| """ | |
| import threading | |
| import time | |
| from typing import Optional, Callable | |
| from .config import logger | |
| class BackgroundInitializer: | |
| """Manages background initialization of heavy components""" | |
| def __init__(self): | |
| self._initialization_thread: Optional[threading.Thread] = None | |
| self._initialization_complete = threading.Event() | |
| self._initialization_error: Optional[Exception] = None | |
| self._progress_callback: Optional[Callable[[str, int], None]] = None | |
| self._status = "Not started" | |
| def set_progress_callback(self, callback: Callable[[str, int], None]): | |
| """Set a callback function to receive progress updates""" | |
| self._progress_callback = callback | |
| def _update_progress(self, message: str, percentage: int): | |
| """Update progress and call callback if set""" | |
| self._status = message | |
| logger.info(f"🔄 Background Init: {message} ({percentage}%)") | |
| if self._progress_callback: | |
| try: | |
| self._progress_callback(message, percentage) | |
| except Exception as e: | |
| logger.error(f"Progress callback error: {e}") | |
| def _initialize_components(self): | |
| """Initialize all heavy components in background thread""" | |
| try: | |
| self._update_progress("Starting background initialization...", 0) | |
| # Step 1: Load embedding model (this is the heaviest component) | |
| self._update_progress("Loading embedding model...", 10) | |
| from .config import get_embedding_model | |
| embedding_model = get_embedding_model() | |
| self._update_progress("Embedding model loaded successfully", 40) | |
| # Step 2: Initialize retrievers (this will load chunks and create vector store) | |
| self._update_progress("Initializing retrievers and loading chunks...", 50) | |
| from .retrievers import _ensure_initialized | |
| _ensure_initialized() | |
| self._update_progress("Retrievers initialized successfully", 90) | |
| # Step 3: Learn medical terminology from corpus | |
| self._update_progress("Learning medical terminology from corpus...", 92) | |
| try: | |
| from .medical_terminology import learn_from_corpus | |
| from . import utils | |
| # Load chunks to learn from | |
| chunks = utils.load_chunks() | |
| if chunks: | |
| # Convert to format expected by learner | |
| documents = [{'content': chunk.page_content} for chunk in chunks[:1000]] # Limit for performance | |
| learn_from_corpus(documents) | |
| logger.info(f"Learned medical terminology from {len(documents)} documents") | |
| except Exception as e: | |
| logger.warning(f"Could not learn terminology from corpus: {e}") | |
| # Step 4: Warm up LLM (optional, lightweight) | |
| self._update_progress("Warming up LLM...", 97) | |
| from .config import get_llm | |
| llm = get_llm() | |
| self._update_progress("All components initialized successfully", 100) | |
| logger.info("✅ Background initialization completed successfully") | |
| except Exception as e: | |
| self._initialization_error = e | |
| logger.error(f"❌ Background initialization failed: {e}") | |
| self._update_progress(f"Initialization failed: {str(e)}", -1) | |
| finally: | |
| self._initialization_complete.set() | |
| def start_background_initialization(self): | |
| """Start background initialization in a separate thread""" | |
| if self._initialization_thread is not None: | |
| logger.warning("Background initialization already started") | |
| return | |
| logger.info("🚀 Starting background initialization...") | |
| self._initialization_thread = threading.Thread( | |
| target=self._initialize_components, | |
| name="BackgroundInitializer", | |
| daemon=True | |
| ) | |
| self._initialization_thread.start() | |
| def is_complete(self) -> bool: | |
| """Check if initialization is complete""" | |
| return self._initialization_complete.is_set() | |
| def wait_for_completion(self, timeout: Optional[float] = None) -> bool: | |
| """Wait for initialization to complete""" | |
| return self._initialization_complete.wait(timeout) | |
| def get_status(self) -> str: | |
| """Get current initialization status""" | |
| return self._status | |
| def get_error(self) -> Optional[Exception]: | |
| """Get initialization error if any""" | |
| return self._initialization_error | |
| def is_successful(self) -> bool: | |
| """Check if initialization completed successfully""" | |
| return self.is_complete() and self._initialization_error is None | |
| # Global initializer instance | |
| _background_initializer = BackgroundInitializer() | |
| def start_background_initialization(progress_callback: Optional[Callable[[str, int], None]] = None): | |
| """Start background initialization with optional progress callback""" | |
| if progress_callback: | |
| _background_initializer.set_progress_callback(progress_callback) | |
| _background_initializer.start_background_initialization() | |
| def wait_for_initialization(timeout: Optional[float] = None) -> bool: | |
| """Wait for background initialization to complete""" | |
| return _background_initializer.wait_for_completion(timeout) | |
| def is_initialization_complete() -> bool: | |
| """Check if background initialization is complete""" | |
| return _background_initializer.is_complete() | |
| def get_initialization_status() -> str: | |
| """Get current initialization status""" | |
| return _background_initializer.get_status() | |
| def get_initialization_error() -> Optional[Exception]: | |
| """Get initialization error if any""" | |
| return _background_initializer.get_error() | |
| def is_initialization_successful() -> bool: | |
| """Check if initialization completed successfully""" | |
| return _background_initializer.is_successful() | |