demo-backend/ai-service/app/main.py

# --- Library Imports ---
from fastapi import Depends, FastAPI, Query, UploadFile, WebSocket
from fastapi.exceptions import RequestValidationError
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from sqlalchemy.orm import Session
from starlette.exceptions import HTTPException

from .services.background_services.update_form_data_in_db import update_db
from .services.doc_agent import analyse_quote, doc_upload, doc_upload_lo
from .services.helpers.db_helper import get_user_id, underwrite
from .services.quotation_analysis import quotation_upload
from .services.underwrite_agent import get_underwriting_remarks
from .services.websocket_streaming import general_info, loan_officer, start_streaming
from .services.background_services.get_document import fetch_from_blob

# ---
# --- User Imports ---
from .utils.app_logger.exception_handlers import (
    http_exception_handler,
    request_validation_exception_handler,
    unhandled_exception_handler,
)
from .utils.app_logger.middleware import log_request_middleware
from .utils.celery.worker import celery_app
from .utils.database import get_db
from .utils.jwt import UnauthenticatedException, UnauthorizedException, VerifyToken
from .utils.redis.redis import redis



# ---

# initialize Constants
DOC_VERIFY_QUEUE = "doc_verify_queue"


# initialize  app
app = FastAPI()


# allow cors
origins = ["*"]

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

# initialize logger
app.middleware("http")(log_request_middleware)
app.add_exception_handler(RequestValidationError, request_validation_exception_handler)
app.add_exception_handler(HTTPException, http_exception_handler)
app.add_exception_handler(Exception, unhandled_exception_handler)


verify_token = VerifyToken()

# initialize base routes


@app.get("/", summary="Base Route", description="Display a welcome message for the AI service.")
def base_route():
    """
    Display a welcome message for the AI service.
    """
    return {"msg": "Welcome to ai service"}


@app.get(
    "/ai/openapi.json",
    include_in_schema=False,
    summary="Custom OpenAPI JSON",
    description="Return the custom OpenAPI JSON.",
)
def custom_swagger_ui_html():
    """
    Return the custom OpenAPI JSON.
    """
    return app.openapi()


@app.put("/db", summary="Update DB Mock", description="Update the database with mock data.")
def update_db_mock(db: Session = Depends(get_db)):
    """
    Update the database with mock data.

    - **db**: The database session.
    """
    update_db({}, db)


@app.websocket("/ws2")
async def websocket_endpoint(websocket: WebSocket, db: Session = Depends(get_db), redis_client=redis):
    """
    WebSocket endpoint for streaming data.

    - **websocket**: The WebSocket connection.
    - **db**: The database session.
    - **redis_client**: The Redis client.
    """
    await start_streaming(websocket, db, redis_client)


@app.websocket("/ws1")
async def websocket_general_info(websocket: WebSocket):
    """
    WebSocket endpoint for general information.

    - **websocket**: The WebSocket connection.
    """
    await general_info(websocket)


@app.websocket("/ws3")
async def websocket_loan_officer(websocket: WebSocket, db: Session = Depends(get_db)):
    """
    WebSocket endpoint for loan officer information.

    - **websocket**: The WebSocket connection.
    - **db**: The database session.
    """
    await loan_officer(websocket, db)


@app.post("/doc-agent", summary="Document Agent Upload", description="Upload a document for the document agent.")
async def doc_agent(file: UploadFile, user_profile_id: int = Query(None)):
    """
    Upload a document for the document agent.

    - **file**: The document file.
    - **user_profile_id**: The ID of the user profile.
    """
    fs = await doc_upload(file)
    doc_dict = {"file_path": fs, "user_profile_id": user_profile_id}
    import json

    celery_app.send_task(DOC_VERIFY_QUEUE, args=[json.dumps(doc_dict)])
    return fs


@app.post(
    "/doc-agent-lo", summary="Document Agent Loan Officer Upload", description="Upload a document for the loan officer."
)
async def doc_agent_lo(file: UploadFile):
    """
    Upload a document for the loan officer.

    - **file**: The document file.
    """
    fs = await doc_upload_lo(file)
    return fs


@app.post("/quotation-analysis", summary="Quotation Analysis", description="Analyze a quotation document.")
async def quotation_analysis(file: UploadFile):
    """
    Analyze a quotation document.

    - **file**: The document file.
    """
    return await analyse_quote(file)


@app.get(
    "/pre-underwrite-details",
    summary="Pre-Underwrite Details",
    description="Fetch pre-underwrite details for a specified application.",
)
async def pre_underwrite_details(
    application_id: int,
    user_id: str,
    db: Session = Depends(get_db),
):
    """
    Fetch pre-underwrite details for a specified application.

    - **application_id**: The ID of the application.
    - **user_id**: The ID of the user.
    - **db**: The database session.
    """
    return underwrite(application_id, user_id, db)


@app.get(
    "/underwrite",
    summary="Underwrite Status and Details",
    description="Fetch the status and details of an underwrite application.",
)
async def underwrite_status_and_details(application_id: int, db: Session = Depends(get_db)):
    """
    Fetch the status and details of an underwrite application.

    - **application_id**: The ID of the application.
    - **db**: The database session.
    """
    user_id = get_user_id(application_id, db)
    data = underwrite(application_id, user_id, db)
    underwrite_remarks = get_underwriting_remarks(data)
    return underwrite_remarks


@app.post("/underwrite", summary="Create Underwrite", description="Create a new underwrite application.")
async def underwrite_create(application_id: int):
    """
    Create a new underwrite application.

    - **application_id**: The ID of the application.
    """
    return JSONResponse(status_code=202, content={"code": 202, "Status": "in Progress"})


@app.websocket("/ws1_secure")
async def secure_websocket_general_info(ws: WebSocket):
    try:
        # Step 1: Extract requested subprotocols
        requested_subprotocols = ws.headers.get("sec-websocket-protocol")

        # Step 2: Select the subprotocol you want to use (if multiple are requested)
        if requested_subprotocols:
            subprotocol = requested_subprotocols.split(",")[0].strip()
            token = requested_subprotocols.split(",")[1].strip()
        else:
            subprotocol = None
            token = None

        # Step 3: Verify token
        await verify_token.verify_websocket(ws, token)

        # Continue with general info logic
        await general_info(ws, None, subprotocol)

    except (UnauthenticatedException, UnauthorizedException):
        await ws.close(code=1008)


@app.websocket("/ws2_secure")
async def secure_websocket_endpoint(ws: WebSocket, db: Session = Depends(get_db), redis_client=redis):
    try:
        # Step 1: Extract requested subprotocols
        requested_subprotocols = ws.headers.get("sec-websocket-protocol")

        # Step 2: Select the subprotocol you want to use (if multiple are requested)
        if requested_subprotocols:
            subprotocol = requested_subprotocols.split(",")[0].strip()
            token = requested_subprotocols.split(",")[1].strip()
        else:
            subprotocol = None
            token = None

        # Step 3: Verify token
        await verify_token.verify_websocket(ws, token)

        # Continue with streaming logic
        await start_streaming(ws, db, redis_client, subprotocol)
    except (UnauthenticatedException, UnauthorizedException):
        await ws.close(code=1008)


@app.websocket("/ws3_secure")
async def secure_websocket_loan_officer(ws: WebSocket, db: Session = Depends(get_db)):
    try:
        # Step 1: Extract requested subprotocols
        requested_subprotocols = ws.headers.get("sec-websocket-protocol")

        # Step 2: Select the subprotocol you want to use (if multiple are requested)
        if requested_subprotocols:
            subprotocol = requested_subprotocols.split(",")[0].strip()
            token = requested_subprotocols.split(",")[1].strip()
        else:
            subprotocol = None
            token = None

        # Step 3: Verify token
        await verify_token.verify_websocket(ws, token)
        verify_token.is_loan_officer()

        # Continue with loan officer logic
        await loan_officer(ws, db, subprotocol)

    except (UnauthenticatedException, UnauthorizedException):
        await ws.close(code=1008)

@app.get("/income-analytics")
async def post_income_analytics(user_profile_id: int = Query(None),db: Session = Depends(get_db)):
    # print(" reached here ",(user_profile_id))
    # return {"filtered_income"}
    await fetch_from_blob(user_profile_id,db)