|
|
import time |
|
|
from datetime import datetime |
|
|
from contextlib import asynccontextmanager |
|
|
from fastapi import FastAPI, Request |
|
|
from fastapi.middleware.cors import CORSMiddleware |
|
|
from app.routers import execution |
|
|
from app.core.config import get_config |
|
|
|
|
|
|
|
|
|
|
|
@asynccontextmanager |
|
|
async def lifespan(app: FastAPI): |
|
|
""" |
|
|
Execute setup logic before the API starts accepting requests. |
|
|
""" |
|
|
config = get_config() |
|
|
print(f"π [Startup] RiverGen AI Engine ({config.MODEL_NAME}) is warming up...") |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
yield |
|
|
|
|
|
print("π [Shutdown] Cleaning up resources...") |
|
|
|
|
|
|
|
|
app = FastAPI( |
|
|
title="RiverGen AI Engine API", |
|
|
description="Enterprise orchestration API for executing queries across SQL, NoSQL, and Streaming sources.", |
|
|
version="1.0.0", |
|
|
lifespan=lifespan, |
|
|
docs_url="/docs", |
|
|
redoc_url="/redoc" |
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
origins = [ |
|
|
"http://localhost:3000", |
|
|
"https://app.rivergen.ai", |
|
|
"https://staging.rivergen.ai" |
|
|
] |
|
|
|
|
|
app.add_middleware( |
|
|
CORSMiddleware, |
|
|
allow_origins=origins, |
|
|
allow_credentials=True, |
|
|
allow_methods=["GET", "POST", "OPTIONS"], |
|
|
allow_headers=["*"], |
|
|
) |
|
|
|
|
|
|
|
|
|
|
|
@app.middleware("http") |
|
|
async def add_process_time_header(request: Request, call_next): |
|
|
start_time = time.time() |
|
|
response = await call_next(request) |
|
|
process_time = time.time() - start_time |
|
|
response.headers["X-Process-Time"] = str(process_time) |
|
|
return response |
|
|
|
|
|
|
|
|
app.include_router(execution.router, prefix="/api/v1") |
|
|
|
|
|
|
|
|
|
|
|
@app.get("/health", tags=["Monitoring"]) |
|
|
def health_check(): |
|
|
""" |
|
|
Dynamic health check for load balancers. |
|
|
""" |
|
|
return { |
|
|
"status": "healthy", |
|
|
"timestamp": datetime.now().isoformat(), |
|
|
"engine": "RiverGen-v1", |
|
|
"uptime_check": True |
|
|
} |
|
|
|
|
|
@app.get("/", tags=["General"]) |
|
|
def read_root(): |
|
|
return { |
|
|
"message": "RiverGen AI Engine is running.", |
|
|
"docs": "/docs", |
|
|
"health": "/health" |
|
|
} |
|
|
|
|
|
if __name__ == "__main__": |
|
|
import uvicorn |
|
|
|
|
|
uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=True) |