Spaces:

rohitdeshmukh318
/

Data_analysis_agent

Sleeping

Auto-run DB migrations on startup to ensure tables exist

1f93fec 5 months ago

5.8 kB

"""
api/main.py
FastAPI application — mounts all routers, configures CORS, rate limiting, lifespan.
"""

import os
from contextlib import asynccontextmanager
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

from fastapi import FastAPI, Request
from fastapi.middleware.cors import CORSMiddleware
from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded

from api.routers import query, upload, history, dashboard, report, schema
from api.routers import metrics as metrics_router
from api.routers import profile as profile_router


# ── Rate limiter ──────────────────────────────────────────────────────────────
limiter = Limiter(key_func=get_remote_address)

import asyncio
import httpx

async def keep_awake():
    """Background task to ping the server and prevent HF Space from sleeping."""
    while True:
        try:
            await asyncio.sleep(600)  # Ping every 10 minutes
            async with httpx.AsyncClient() as client:
                await client.get("http://127.0.0.1:7860/health", timeout=10.0)
                print("[Keep-Awake] Pinged health endpoint to prevent sleep.")
        except asyncio.CancelledError:
            break
        except Exception as e:
            print(f"[Keep-Awake] Ping failed: {e}")

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: pre-warm LLM clients + embedding model
    from llm import get_groq_client, get_embedder
    get_groq_client()
    # Pre-download and load embedding model (first run downloads ~420MB)
    print("[Lifespan] Loading embedding model (first run may download ~420MB)...")
    get_embedder()
    print("[Lifespan] Embedding model loaded successfully.")
    
    # Run automatic database migrations on boot
    print("[Lifespan] Ensuring database tables exist...")
    try:
        from db.pool import get_connection, release_connection
        import os
        conn = get_connection()
        try:
            with conn.cursor() as cur:
                with open(os.path.join(os.path.dirname(os.path.dirname(__file__)), "scripts", "migrate.sql"), "r") as f:
                    cur.execute(f.read())
                conn.commit()
            print("[Lifespan] Database migrations executed successfully.")
        finally:
            release_connection(conn)
    except Exception as e:
        print(f"[Lifespan] Failed to run database migrations: {e}")
    
    # Start the keep-awake background task
    ping_task = asyncio.create_task(keep_awake())
    
    yield
    
    # Shutdown: cancel task and close DB pool
    ping_task.cancel()
    try:
        from db.pool import get_pool
        get_pool().closeall()
    except Exception:
        pass


app = FastAPI(
    title="Cloud Data Analyst Agent",
    version="2.0.0",
    description="AI-powered data analyst with self-correcting LangGraph agent, "
                "multi-turn conversations, anomaly detection, and real-time observability.",
    docs_url="/docs",
    lifespan=lifespan,
)

# Attach rate limiter
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

# CORS — allow the Vite frontend and any Render preview URL
ALLOWED_ORIGINS = os.environ.get(
    "ALLOWED_ORIGINS",
    "http://localhost:5173,http://localhost:3000",
).split(",")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"] if os.environ.get("DEMO_MODE") == "true" else ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Routers
app.include_router(query.router,            prefix="/api/query",     tags=["query"])
app.include_router(upload.router,           prefix="/api/upload",    tags=["upload"])
app.include_router(history.router,          prefix="/api/history",   tags=["history"])
app.include_router(dashboard.router,        prefix="/api/dashboard", tags=["dashboard"])
app.include_router(report.router,           prefix="/api/report",    tags=["report"])
app.include_router(schema.router,           prefix="/api/schema",    tags=["schema"])
app.include_router(metrics_router.router,   prefix="/api/metrics",   tags=["metrics"])
app.include_router(profile_router.router,   prefix="/api/profile",   tags=["profile"])


@app.get("/health")
async def health():
    return {"status": "ok", "version": "2.0.0"}

# ── Serve Frontend ────────────────────────────────────────────────────────────
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse

frontend_dist = os.path.join(os.path.dirname(os.path.dirname(__file__)), "frontend", "dist")

# Mount the assets directory specifically if it exists
assets_dir = os.path.join(frontend_dist, "assets")
if os.path.exists(assets_dir):
    app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")
    
@app.get("/{full_path:path}")
async def serve_frontend(full_path: str):
    # Prevent API routes from being intercepted if they fall through
    if full_path.startswith("api/"):
        from fastapi import HTTPException
        raise HTTPException(status_code=404, detail="API route not found")
        
    file_path = os.path.join(frontend_dist, full_path)
    if full_path and os.path.isfile(file_path):
        return FileResponse(file_path)
    # Fallback to index.html for React Router
    index_file = os.path.join(frontend_dist, "index.html")
    if os.path.isfile(index_file):
        return FileResponse(index_file)
    return {"error": "Frontend build not found"}

	"""
	api/main.py
	FastAPI application — mounts all routers, configures CORS, rate limiting, lifespan.
	"""

	import os
	from contextlib import asynccontextmanager
	from dotenv import load_dotenv

	# Load environment variables from .env file
	load_dotenv()

	from fastapi import FastAPI, Request
	from fastapi.middleware.cors import CORSMiddleware
	from slowapi import Limiter, _rate_limit_exceeded_handler
	from slowapi.util import get_remote_address
	from slowapi.errors import RateLimitExceeded

	from api.routers import query, upload, history, dashboard, report, schema
	from api.routers import metrics as metrics_router
	from api.routers import profile as profile_router


	# ── Rate limiter ──────────────────────────────────────────────────────────────
	limiter = Limiter(key_func=get_remote_address)

	import asyncio
	import httpx

	async def keep_awake():
	"""Background task to ping the server and prevent HF Space from sleeping."""
	while True:
	try:
	await asyncio.sleep(600) # Ping every 10 minutes
	async with httpx.AsyncClient() as client:
	await client.get("http://127.0.0.1:7860/health", timeout=10.0)
	print("[Keep-Awake] Pinged health endpoint to prevent sleep.")
	except asyncio.CancelledError:
	break
	except Exception as e:
	print(f"[Keep-Awake] Ping failed: {e}")

	@asynccontextmanager
	async def lifespan(app: FastAPI):
	# Startup: pre-warm LLM clients + embedding model
	from llm import get_groq_client, get_embedder
	get_groq_client()
	# Pre-download and load embedding model (first run downloads ~420MB)
	print("[Lifespan] Loading embedding model (first run may download ~420MB)...")
	get_embedder()
	print("[Lifespan] Embedding model loaded successfully.")

	# Run automatic database migrations on boot
	print("[Lifespan] Ensuring database tables exist...")
	try:
	from db.pool import get_connection, release_connection
	import os
	conn = get_connection()
	try:
	with conn.cursor() as cur:
	with open(os.path.join(os.path.dirname(os.path.dirname(__file__)), "scripts", "migrate.sql"), "r") as f:
	cur.execute(f.read())
	conn.commit()
	print("[Lifespan] Database migrations executed successfully.")
	finally:
	release_connection(conn)
	except Exception as e:
	print(f"[Lifespan] Failed to run database migrations: {e}")

	# Start the keep-awake background task
	ping_task = asyncio.create_task(keep_awake())

	yield

	# Shutdown: cancel task and close DB pool
	ping_task.cancel()
	try:
	from db.pool import get_pool
	get_pool().closeall()
	except Exception:
	pass


	app = FastAPI(
	title="Cloud Data Analyst Agent",
	version="2.0.0",
	description="AI-powered data analyst with self-correcting LangGraph agent, "
	"multi-turn conversations, anomaly detection, and real-time observability.",
	docs_url="/docs",
	lifespan=lifespan,
	)

	# Attach rate limiter
	app.state.limiter = limiter
	app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

	# CORS — allow the Vite frontend and any Render preview URL
	ALLOWED_ORIGINS = os.environ.get(
	"ALLOWED_ORIGINS",
	"http://localhost:5173,http://localhost:3000",
	).split(",")

	app.add_middleware(
	CORSMiddleware,
	allow_origins=["*"] if os.environ.get("DEMO_MODE") == "true" else ALLOWED_ORIGINS,
	allow_credentials=True,
	allow_methods=["*"],
	allow_headers=["*"],
	)

	# Routers
	app.include_router(query.router, prefix="/api/query", tags=["query"])
	app.include_router(upload.router, prefix="/api/upload", tags=["upload"])
	app.include_router(history.router, prefix="/api/history", tags=["history"])
	app.include_router(dashboard.router, prefix="/api/dashboard", tags=["dashboard"])
	app.include_router(report.router, prefix="/api/report", tags=["report"])
	app.include_router(schema.router, prefix="/api/schema", tags=["schema"])
	app.include_router(metrics_router.router, prefix="/api/metrics", tags=["metrics"])
	app.include_router(profile_router.router, prefix="/api/profile", tags=["profile"])


	@app.get("/health")
	async def health():
	return {"status": "ok", "version": "2.0.0"}

	# ── Serve Frontend ────────────────────────────────────────────────────────────
	from fastapi.staticfiles import StaticFiles
	from fastapi.responses import FileResponse

	frontend_dist = os.path.join(os.path.dirname(os.path.dirname(__file__)), "frontend", "dist")

	# Mount the assets directory specifically if it exists
	assets_dir = os.path.join(frontend_dist, "assets")
	if os.path.exists(assets_dir):
	app.mount("/assets", StaticFiles(directory=assets_dir), name="assets")

	@app.get("/{full_path:path}")
	async def serve_frontend(full_path: str):
	# Prevent API routes from being intercepted if they fall through
	if full_path.startswith("api/"):
	from fastapi import HTTPException
	raise HTTPException(status_code=404, detail="API route not found")

	file_path = os.path.join(frontend_dist, full_path)
	if full_path and os.path.isfile(file_path):
	return FileResponse(file_path)
	# Fallback to index.html for React Router
	index_file = os.path.join(frontend_dist, "index.html")
	if os.path.isfile(index_file):
	return FileResponse(index_file)
	return {"error": "Frontend build not found"}