Update main.py

main.py

@@ -1,12 +1,21 @@
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Request
 from pydantic import BaseModel
-from datasets import load_dataset
+from datasets import load_dataset, Dataset
 
 app = FastAPI()
 
+# Custom exception handler for better error visibility
+@app.middleware("http")
+async def add_process_time_header(request: Request, call_next):
+    try:
+        response = await call_next(request)
+        return response
+    except Exception as exc:
+        return JSONResponse(status_code=500, content={"message": str(exc)})
+
 # Load the datasets from the Hugging Face Hub
-customers_dataset = load_dataset("dwb2023/blackbird-customers", split="train")
-orders_dataset = load_dataset("dwb2023/blackbird-orders", split="train")
+customers_dataset = load_dataset("dwb2023/customers", split="train")
+orders_dataset = load_dataset("dwb2023/orders", split="train")
 
 class GetUserInput(BaseModel):
     key: str
@@ -55,7 +64,7 @@ def update_user(input: UpdateUserInput):
 
     if user_found:
         updated_dataset = Dataset.from_list(customers_data)
-        updated_dataset.push_to_hub("dwb2023/blackbird-customers", split="train")
+        updated_dataset.push_to_hub("dwb2023/customers", split="train")
         customers_dataset = updated_dataset  # Update the global dataset
         return {"message": "User information updated successfully"}
     else:
@@ -91,7 +100,7 @@ def cancel_order(input: CancelOrderInput):
 
     if order_found:
         updated_orders_dataset = Dataset.from_list(orders_data)
-        updated_orders_dataset.push_to_hub("dwb2023/blackbird-orders", split="train")
+        updated_orders_dataset.push_to_hub("dwb2023/orders", split="train")
         orders_dataset = updated_orders_dataset  # Update the global dataset
         return {"status": "Cancelled"}
     else: