Backend with FastAPI¶

"""FastAPI application entry point."""

import contextlib

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

from backend.api.routes import router as api_router
from backend.settings import get_settings

APP_VERSION = "0.1.0"


@contextlib.asynccontextmanager
async def lifespan(app: FastAPI):
    """Startup and shutdown logic.

    Use this to initialize shared resources (DB connections, SDK clients)
    and clean them up on shutdown.
    """
    # Startup: nothing to initialize yet
    yield
    # Shutdown: nothing to clean up yet


def create_app() -> FastAPI:
    settings = get_settings()

    app = FastAPI(
        title=settings.app_name,
        version=APP_VERSION,
        docs_url="/docs",
        redoc_url="/redoc",
        lifespan=lifespan,
    )

    app.add_middleware(
        CORSMiddleware,
        allow_origins=[settings.allow_cors_origin],
        allow_credentials=True,
        allow_methods=["*"],
        allow_headers=["*"],
    )

    app.include_router(api_router, prefix="/api/v1")

    return app


app = create_app()

"""Configuration for the backend service."""

from functools import lru_cache

from pydantic import Field
from pydantic_settings import BaseSettings, SettingsConfigDict


class Settings(BaseSettings):
    model_config = SettingsConfigDict(env_file=".env", case_sensitive=False)

    app_name: str = Field(
        default="My Gaia App",
        validation_alias="APP_NAME",
    )

    gaia_base_url: str = Field(
        default="https://helios.cohesity.com/v2/mcm/gaia",
        validation_alias="GAIA_BASE_URL",
    )
    gaia_api_key: str = Field(
        default="",
        validation_alias="GAIA_API_KEY",
    )
    gaia_verify_ssl: bool = Field(
        default=True,
        validation_alias="GAIA_VERIFY_SSL",
    )
    request_timeout_seconds: float = Field(
        default=30.0,
        validation_alias="REQUEST_TIMEOUT_SECONDS",
    )

    allow_cors_origin: str = Field(
        default="http://localhost:5173",
        validation_alias="ALLOW_CORS_ORIGIN",
    )

    session_ttl_seconds: int = Field(
        default=3600,
        validation_alias="SESSION_TTL_SECONDS",
    )

    log_level: str = Field(default="INFO", validation_alias="LOG_LEVEL")


@lru_cache()
def get_settings() -> Settings:
    return Settings()

"""FastAPI route handlers."""

from fastapi import APIRouter, Depends, HTTPException, status
from fastapi.responses import StreamingResponse

from backend.api.dependencies import get_gaia_client, get_session_api_key
from backend.models.api_models import (
    AskRequest,
    AskResponse,
    LoginRequest,
    LoginResponse,
    DatasetListResponse,
)
from backend.storage.session_store import session_store

from gaia_sdk import GaiaClient

router = APIRouter()


# ── Authentication ────────────────────────────────────────────────

@router.post("/login", response_model=LoginResponse, tags=["Auth"])
async def login(request: LoginRequest):
    """Validate the API key against Gaia and create a session."""
    async with GaiaClient(api_key=request.api_key) as client:
        try:
            await client.list_datasets()
        except Exception:
            raise HTTPException(
                status_code=status.HTTP_401_UNAUTHORIZED,
                detail="Invalid API key — could not authenticate with Gaia.",
            )

    session_id = session_store.create(api_key=request.api_key)
    return LoginResponse(session_id=session_id)


@router.post("/logout", status_code=status.HTTP_204_NO_CONTENT, tags=["Auth"])
async def logout(api_key: str = Depends(get_session_api_key)):
    """Destroy the current session."""
    # get_session_api_key already validated the session;
    # we just need to remove it
    session_store.delete_by_api_key(api_key)


# ── Datasets ──────────────────────────────────────────────────────

@router.get("/datasets", response_model=DatasetListResponse, tags=["Datasets"])
async def list_datasets(client: GaiaClient = Depends(get_gaia_client)):
    """List available Gaia datasets."""
    datasets = await client.list_datasets()
    return DatasetListResponse(datasets=datasets)


# ── Ask ───────────────────────────────────────────────────────────

@router.post("/ask", response_model=AskResponse, tags=["Ask"])
async def ask(
    request: AskRequest,
    client: GaiaClient = Depends(get_gaia_client),
):
    """Send a RAG query to Gaia."""
    result = await client.ask(
        dataset_names=request.dataset_names,
        query=request.query,
        conversation_id=request.conversation_id,
    )
    return AskResponse(
        response=result.response_string,
        query_uid=result.query_uid,
        conversation_id=result.conversation_id,
        documents=result.documents,
    )


# ── Streaming ─────────────────────────────────────────────────────

@router.post("/ask/stream", tags=["Ask"])
async def ask_stream(
    request: AskRequest,
    client: GaiaClient = Depends(get_gaia_client),
):
    """Stream a RAG query response as Server-Sent Events."""

    async def event_generator():
        async for chunk in client.ask_stream_iter(
            dataset_names=request.dataset_names,
            query=request.query,
            conversation_id=request.conversation_id,
        ):
            yield f"event: {chunk.event}\ndata: {chunk.data}\n\n"

    return StreamingResponse(
        event_generator(),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "X-Accel-Buffering": "no",
        },
    )

"""FastAPI dependency providers."""

from fastapi import Depends, Header, HTTPException, status

from backend.settings import get_settings
from backend.storage.session_store import session_store

from gaia_sdk import GaiaClient


def get_session_id(
    x_session_id: str | None = Header(default=None, alias="X-Session-ID"),
) -> str:
    """Extract and validate the session ID from the request header."""
    if not x_session_id:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Missing X-Session-ID header.",
        )
    return x_session_id


def get_session_api_key(
    session_id: str = Depends(get_session_id),
) -> str:
    """Look up the API key for the given session ID."""
    api_key = session_store.get_api_key(session_id)
    if not api_key:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid or expired session.",
        )
    return api_key


async def get_gaia_client(
    api_key: str = Depends(get_session_api_key),
) -> GaiaClient:
    """Construct a GaiaClient bound to the caller's API key.

    The caller is responsible for using the client within
    an async context manager in the route handler, or this
    dependency can yield the client directly.
    """
    settings = get_settings()
    client = GaiaClient(
        api_key=api_key,
        base_url=settings.gaia_base_url,
        timeout=int(settings.request_timeout_seconds),
        verify_ssl=settings.gaia_verify_ssl,
    )
    async with client as c:
        yield c

"""API request and response models."""

from pydantic import BaseModel, Field


class LoginRequest(BaseModel):
    api_key: str = Field(..., description="Cohesity Gaia API key.")


class LoginResponse(BaseModel):
    session_id: str


class AskRequest(BaseModel):
    dataset_names: list[str] = Field(..., alias="datasetNames")
    query: str = Field(..., alias="queryString")
    conversation_id: str | None = Field(None, alias="conversationId")

    model_config = {"populate_by_name": True}


class AskResponse(BaseModel):
    response: str | None
    query_uid: str | None
    conversation_id: str | None
    documents: list | None = None


class DatasetListResponse(BaseModel):
    datasets: list