From 7d1210c8892b9b321dd6daeeaa9d2b58c5345725 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Wed, 27 Nov 2024 15:57:37 +0100 Subject: [PATCH] feat: Add custom exception handling to dataset router Added custom exceptions for dataset router Feature COG-502 --- cognee/api/v1/datasets/routers/get_datasets_router.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/cognee/api/v1/datasets/routers/get_datasets_router.py b/cognee/api/v1/datasets/routers/get_datasets_router.py index a3a30278b..1e62f7f49 100644 --- a/cognee/api/v1/datasets/routers/get_datasets_router.py +++ b/cognee/api/v1/datasets/routers/get_datasets_router.py @@ -156,18 +156,13 @@ def get_datasets_router() -> APIRouter: dataset_data = await get_dataset_data(dataset.id) if dataset_data is None: - raise HTTPException(status_code=404, detail=f"No data found in dataset ({dataset_id}).") + raise EntityNotFoundError(message=f"No data found in dataset ({dataset_id}).") matching_data = [data for data in dataset_data if str(data.id) == data_id] # Check if matching_data contains an element if len(matching_data) == 0: - return JSONResponse( - status_code=404, - content={ - "detail": f"Data ({data_id}) not found in dataset ({dataset_id})." - } - ) + raise EntityNotFoundError(message= f"Data ({data_id}) not found in dataset ({dataset_id}).") data = matching_data[0]