webapp/tipitaka-api/app/main.py

import logging
import os
import asyncio
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from contextlib import asynccontextmanager

from app.config import get_settings
from app.database.sqlite_db import get_db
from app.routers import pages, search, ai, health, reference

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Set PyTorch threads to 1 to prevent pegging CPU and causing HF container crashes
    try:
        import torch
        torch.set_num_threads(1)
        torch.set_num_interop_threads(1)
        logger.info("PyTorch thread limits set to 1 to optimize system responsiveness.")
    except Exception as e:
        logger.warning(f"Failed to set PyTorch thread limits: {e}")

    # Startup: Load DB to memory for performance
    settings = get_settings()
    db = get_db()
    
    logger.info(f"Using database at: {db.db_path}")
    
    # Run heavy DB loading in a background thread to keep startup responsive
    import anyio
    
    async def init_db():
        try:
            logger.info("Initializing services in background...")
            
            # 1. Load SQLite
            await anyio.to_thread.run_sync(db.load_to_memory)
            logger.info("Database loaded to memory.")
            await anyio.to_thread.run_sync(db.ensure_search_log_table)
            logger.info("Search log table ready.")
            
            # 2. Load RAG (This will trigger snapshot extraction if needed)
            from app.routers.ai import get_llm_service
            logger.info("Pre-loading RAG Service (Qdrant + Embedding Model)...")
            await anyio.to_thread.run_sync(get_llm_service)
            logger.info("RAG Service initialized.")

            # 3. Load Qwen Query Transform Model (if runtime is transformers)
            qwen_runtime = os.getenv("QWEN_RUNTIME", "transformers")
            if qwen_runtime == "transformers":
                from app.services.query_transform_service import QueryTransformService
                model_id = os.getenv("QWEN_MODEL_ID", "Qwen/Qwen2.5-1.5B-Instruct")
                logger.info(f"Pre-loading Qwen Model ({model_id}) for Query Transformation...")
                await anyio.to_thread.run_sync(QueryTransformService.preload, model_id)
                logger.info("Qwen Model preloaded.")
            
        except asyncio.CancelledError:
            logger.info("Background initialization cancelled (system shutting down).")
        except Exception as e:
            logger.error(f"Background initialization failed: {e}")

    # Start the background initialization task
    task = asyncio.create_task(init_db())
    
    yield
    # Shutdown logic (if any)
    logger.info("Shutting down...")

app = FastAPI(
    title="Tipitaka API",
    description="API for accessing MCU Thai Tipitaka with AI Search capabilities",
    version="1.0.0",
    lifespan=lifespan
)

settings = get_settings()

# CORS Middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=[settings.CORS_ORIGINS],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Include API Routers (must be before static mount)
app.include_router(pages.router, prefix="/api")
app.include_router(search.router, prefix="/api")
app.include_router(ai.router, prefix="/api")
app.include_router(health.router, prefix="/api")
app.include_router(reference.router, prefix="/api")

@app.get("/health")
async def health():
    return {"status": "healthy"}

# ── Serve built frontend in production (Docker/HF Space) ──
if settings.SERVE_STATIC:
    static_dir = settings.STATIC_DIR or os.path.join(
        os.path.dirname(__file__), "..", "..", "tipitaka-web", "dist"
    )
    if os.path.isdir(static_dir):
        from fastapi.responses import FileResponse, HTMLResponse
        from fastapi.staticfiles import StaticFiles
        from pathlib import Path
        from starlette.requests import Request

        static_path = Path(static_dir)

        # Mount /assets so Vite-built absolute paths (/assets/...) resolve correctly.
        # We intentionally do NOT mount at "/" to avoid StaticFiles swallowing SPA routes.
        assets_dir = static_path / "assets"
        if assets_dir.is_dir():
            logger.info(f"Mounting /assets from: {assets_dir}")
            # Vite builds with hashed filenames (/assets/index-XXXX.js), safe to cache forever
            class _HashStaticFiles(StaticFiles):
                async def get_response(self, path: str, scope):
                    resp = await super().get_response(path, scope)
                    resp.headers["Cache-Control"] = "public, max-age=31536000, immutable"
                    return resp
            app.mount("/assets", _HashStaticFiles(directory=str(assets_dir)), name="assets")

        # Serve individual root-level static files (favicon, icons, manifest, robots.txt, etc.)
        # without needing a full root mount.
        # We explicitly list common root-level file patterns that Vite copies from public/.
        for _static_pattern in [
            "/favicon.ico",
            "/favicon.svg",
            "/manifest.json",
            "/robots.txt",
            "/icon-192.png",
            "/icon-512.png",
            "/apple-touch-icon.png",
            "/apple-touch-icon-precomposed.png",
            "/apple-touch-icon-120x120.png",
            "/apple-touch-icon-120x120-precomposed.png",
            "/icons.svg",
        ]:
            @app.get(_static_pattern, include_in_schema=False)
            async def serve_root_static(request: Request, _p=_static_pattern):
                file = static_path / _p.lstrip("/")
                if file.exists():
                    return FileResponse(
                        file,
                        headers={"Cache-Control": "public, max-age=0, must-revalidate"},
                    )
                from fastapi import HTTPException
                raise HTTPException(status_code=404)

        # ── SPA catch-all: serve index.html for all non-API, non-asset routes ──
        # This is reliable because it's an explicit FastAPI route, not middleware.
        # StaticFiles routes above take priority (registered first); this catches the rest.
        logger.info(f"Registering SPA catch-all, index.html: {static_path / 'index.html'}")

        @app.get("/{full_path:path}", include_in_schema=False)
        async def serve_spa(full_path: str):
            index_file = static_path / "index.html"
            return HTMLResponse(
                content=index_file.read_text(encoding="utf-8"),
                headers={"Cache-Control": "no-cache, no-store, must-revalidate"},
            )

    else:
        logger.warning(f"SERVE_STATIC=true but directory not found: {static_dir}")
else:
    # Only register root route when NOT serving static frontend
    @app.get("/")
    async def root():
        return {
            "message": "Tipitaka API is running",
            "docs": "/docs",
            "version": "1.0.0"
        }