api/main.py

"""
ArcisVLM API — FastAPI application entry point.

Start with:
    uvicorn api.main:app --host 0.0.0.0 --port 8000 --reload

For production:
    uvicorn api.main:app --host 0.0.0.0 --port 8000 --workers 1

Note: --workers 1 because the model is loaded in-process. For multi-worker,
use a separate model server with gRPC.
"""

from __future__ import annotations

import os
import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI, WebSocket
from fastapi.middleware.cors import CORSMiddleware

from api.models import HealthResponse
from api.deps import get_model_manager, configure_model_manager

logger = logging.getLogger(__name__)

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s: %(message)s",
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Initialize model and agents on startup, clean up on shutdown."""
    # Allow override via env vars
    config_path = os.environ.get("ARCISVLM_CONFIG", "configs/scale_1.3b.yaml")
    checkpoint_path = os.environ.get("ARCISVLM_CHECKPOINT", None)
    device = os.environ.get("ARCISVLM_DEVICE", "cpu")

    manager = configure_model_manager(
        config_path=config_path,
        checkpoint_path=checkpoint_path,
        device=device,
    )

    # Don't block startup if model fails to load — allow health checks to report status
    try:
        manager.initialize()
    except Exception as e:
        logger.error("Model initialization failed: %s", e)

    logger.info("ArcisVLM API started (model_loaded=%s)", manager.model is not None)
    yield
    manager.shutdown()
    logger.info("ArcisVLM API shutdown")


app = FastAPI(
    title="ArcisVLM API",
    description="Agentic Vision-Language Model for real-time video analytics. "
                "VL-JEPA + MoE architecture with 8 expert agents for VQA, detection, "
                "alerting, counting, OCR, and reasoning across 1000+ IP cameras.",
    version="1.0.0",
    lifespan=lifespan,
)

# CORS — allow dashboard and external clients
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Tighten in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ---------------------------------------------------------------------------
# Health check (top-level, no prefix)
# ---------------------------------------------------------------------------
@app.get("/health", response_model=HealthResponse)
async def health():
    """Health check — returns model and agent status."""
    manager = get_model_manager()
    return HealthResponse(
        status="ok" if manager.is_initialized else "starting",
        model_loaded=manager.model is not None,
        agents_ready=manager.orchestrator is not None,
    )


# ---------------------------------------------------------------------------
# Include route modules
# ---------------------------------------------------------------------------
from api.routes import inference, agents, alerts, streams, metrics  # noqa: E402

app.include_router(inference.router, prefix="/api/v1")
app.include_router(agents.router, prefix="/api/v1")
app.include_router(alerts.router, prefix="/api/v1")
app.include_router(streams.router, prefix="/api/v1")
app.include_router(metrics.router, prefix="/api/v1")


# ---------------------------------------------------------------------------
# WebSocket endpoints
# ---------------------------------------------------------------------------
from api.ws import feed_websocket, chat_websocket  # noqa: E402


@app.websocket("/ws/feed/{camera_id}")
async def ws_feed(ws: WebSocket, camera_id: str):
    """Live detection results stream for a camera."""
    await feed_websocket(ws, camera_id)


@app.websocket("/ws/chat")
async def ws_chat(ws: WebSocket):
    """Interactive VQA chat."""
    await chat_websocket(ws)