ai-service/app/main.py

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

# --- User Imports ---
from .utils.app_logger.exception_handlers import (
    request_validation_exception_handler,
    http_exception_handler,
    unhandled_exception_handler,
)
from .utils.app_logger.middleware import log_request_middleware
from .services.websocket_streaming import start_streaming, general_info
from .services.doc_agent import doc_upload
from .services.underwrite_agent import get_underwriting_remarks
from .utils.redis.redis import redis
from .services.background_services.update_form_data_in_db import update_db
from .utils.database import get_db
from .utils.celery.worker import celery_app
from .services.helpers.db_helper import underwrite, get_user_id
# ---

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)


# initialize base routes

@app.get("/")
def base_route():
    return {"msg": "Welcome to ai service"}


@app.get("/ai/openapi.json", include_in_schema=False)
def custom_swagger_ui_html():
    return app.openapi()


@app.put("/db")
def update_db_mock(db: Session = Depends(get_db)):
    update_db({}, db)


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket, db: Session = Depends(get_db), redis_client=redis):
    await start_streaming(websocket, db, redis_client)


@app.websocket("/ws1")
async def websocket_general_info(websocket: WebSocket):
    await general_info(websocket)


@app.post("/doc-agent")
async def doc_agent(file: UploadFile, user_profile_id: int = Query(None)):
    fs = await doc_upload(file)
    doc_dict = {
        'file_path': fs['file_path'],
        'user_profile_id': user_profile_id
    }
    import json
    celery_app.send_task(DOC_VERIFY_QUEUE, args=[json.dumps(doc_dict)])
    return fs

@app.get("/pre-underwrite-details")
async def pre_underwrite_details(application_id: int,user_id: str,db: Session = Depends(get_db),
):
    return underwrite(application_id, user_id, db)

@app.get('/underwrite')
async def underwrite_status_and_details(applicationId: int, db: Session = Depends(get_db)):
    user_id = get_user_id(applicationId,db)
    data = underwrite(applicationId, user_id, db)
    underwrite_remarks = get_underwriting_remarks(data)
    return underwrite_remarks


@app.post('/underwrite')
async def underwrite_create(applicationId: int):
    return JSONResponse(status_code=202, content={"code": 202, "Status": "in Progress"})