main.py

"""
FastAPI application for RAG-based document Q&A system
"""

# 1st party imports
from contextlib import asynccontextmanager

# 3rd party imports
from fastapi import FastAPI
from sqlalchemy import text as SQLText
from fastapi.middleware.cors import CORSMiddleware

# local imports
from log_config import get_logger
from settings import project_settings
from database.vector_db import QdrantVectorDB
from database.postgres_db import get_sync_session
from routers import chat_routes, doc_routes, auth_routes


# create a logger
logger = get_logger(__name__)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """
    Async context manager for FastAPI application lifespan.

    Handles startup and shutdown events for the application.
    On startup, recreates the vector database collection to ensure a clean state.

    Args:
        app: The FastAPI application instance.

    Yields:
        None: Control is yielded back to the application after startup completes.
    """

    # log the starting
    logger.info("Starting the application.")

    # create the vector db client and recreate the collection
    logger.info("Connecting to vector database.")
    vector_db_client = QdrantVectorDB(
        vector_db_url=project_settings.VECTOR_DB_URL,
        size=project_settings.EMBED_SIZE,
    )

    # initialize the collection
    logger.info("Initializing the vector database collection.")
    vector_db_client.initialize_collection()

    # check connection to the postgres Database
    logger.info("Checking connection to postgres database.")
    with get_sync_session() as session:
        session.execute(SQLText("SELECT 1"))

    # log the completion
    logger.info("Application started successfully.")

    # yield
    yield


# Create FastAPI app
app = FastAPI(
    title="RAG Document Q&A API",
    description="A FastAPI application for document-based question answering using RAG",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc",
    lifespan=lifespan,
)

# Add CORS middleware
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# Health check endpoint
@app.get("/health")
async def health_check():
    """Health check endpoint"""
    return {"status": "healthy", "message": "RAG API is running."}


# Include routers
app.include_router(auth_routes.router, prefix="/auth", tags=["authentication"])
app.include_router(chat_routes.router, prefix="/chats", tags=["chat"])
app.include_router(doc_routes.router, prefix="/docs", tags=["documents"])