diff --git a/src/aeros_simulation/router.py b/src/aeros_simulation/router.py index 09ddd69..fbdf124 100644 --- a/src/aeros_simulation/router.py +++ b/src/aeros_simulation/router.py @@ -216,8 +216,12 @@ async def get_custom_parameters_controller(db_session: DbSession): "status": "success", "message": "Simulation result retrieved successfully", } - -@router.post("/calculate_eaf_contribution", response_model=StandardResponse[dict]) + + + +airflow_router = APIRouter() + +@airflow_router.post("/calculate_eaf_contribution", response_model=StandardResponse[dict]) async def calculate_contribution( db_session: DbSession, simulation_in: SimulationInput, @@ -239,8 +243,18 @@ async def calculate_contribution( start_index = batch_num * batch_size end_index = start_index + batch_size + + if start_index >= len(eqs): + return { + "data": contribution_results, + "status": "success", + "message": "No more equipment to process", + } + if end_index > len(eqs): end_index = len(eqs) + + eqs = eqs[start_index:end_index] for eq in eqs: diff --git a/src/api.py b/src/api.py index 8f77220..7f689ed 100644 --- a/src/api.py +++ b/src/api.py @@ -5,7 +5,7 @@ from fastapi.responses import JSONResponse from pydantic import BaseModel from src.aeros_project.router import router as aeros_project_router -from src.aeros_simulation.router import router as aeros_simulation_router +from src.aeros_simulation.router import router as aeros_simulation_router, airflow_router from src.auth.service import JWTBearer from src.dashboard_model.router import router as dashboard_model_router from src.aeros_equipment.router import router as aeros_equipment_router @@ -51,4 +51,5 @@ aeros_routes.include_router(aeros_equipment_router, prefix="/equipment") authenticated_api_router.include_router(aeros_routes) +api_router.include_router(airflow_router, prefix="/airflow") api_router.include_router(authenticated_api_router)