from fastapi import FastAPI, Request from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse, RedirectResponse from contextlib import asynccontextmanager from app.routers import shield, auth from app.models.deberta_model import classifier from app.database.db import init_db import os @asynccontextmanager async def lifespan(app: FastAPI): # Startup: Initialize Database & Load classifier model print("Initializing Database...") try: init_db() print("Database initialized successfully!") except Exception as e: print(f"Error initializing DB: {e}") print("Loading DeBERTa-v3 model. This might take a few moments on the first run...") try: classifier.initialize() print("DeBERTa-v3 model loaded successfully!") except Exception as e: print(f"Error pre-loading model: {e}") print("Initializing SemanticGuard embedding model...") try: from app.services.semantic_guard import semantic_guard semantic_guard.initialize() print("SemanticGuard embedding model loaded successfully!") except Exception as e: print(f"Error pre-loading SemanticGuard: {e}") yield # Shutdown: Clean up if needed pass app = FastAPI( title="GuardrailShield AI", description="Defensive Guardrail System sitting in front of LLMs", version="1.0.0", lifespan=lifespan ) # Include routers app.include_router(shield.router) app.include_router(auth.router) # Mount Static Files static_dir = os.path.join(os.path.dirname(__file__), "static") if os.path.exists(static_dir): app.mount("/static", StaticFiles(directory=static_dir), name="static") @app.get("/") def get_chat_client(): """Serves the public user chat interface.""" return FileResponse(os.path.join(static_dir, "chat.html")) @app.get("/login") def get_login_page(): """Serves the admin login page.""" return FileResponse(os.path.join(static_dir, "login.html")) @app.get("/dashboard") def get_dashboard_ui(request: Request): """Serves the security admin dashboard.""" return FileResponse(os.path.join(static_dir, "index.html")) @app.get("/health", tags=["Health"]) def health_check(): return {"status": "healthy"}