Update app.py

app.py

@@ -117,17 +117,20 @@ def health():
     mongo.admin.command("ping")
     return HealthResponse(status="ok", db=db.name, model="Qwen/Qwen-Image-Edit")
 
-
 @app.post("/generate")
 async def generate(
     prompt: str = Form(...),
     image1: UploadFile = File(...),
     image2: Optional[UploadFile] = File(None),
+    user_id: Optional[str] = Form(None),
+    category_id: Optional[str] = Form(None),
     user=Depends(verify_firebase_token)
 ):
     start_time = time.time()
 
-    # PREPARE IMAGE(S)
+    # -------------------------
+    # 1. VALIDATE & READ IMAGES
+    # -------------------------
     try:
         img1_bytes = await image1.read()
         pil_img1 = prepare_image(img1_bytes)
@@ -138,7 +141,7 @@ async def generate(
             metadata={"role": "input"}
         )
     except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Failed to read first image: {e}")
+        raise HTTPException(400, f"Failed to read first image: {e}")
 
     img2_bytes = None
     input2_id = None
@@ -155,9 +158,11 @@ async def generate(
                 metadata={"role": "input"}
             )
         except Exception as e:
-            raise HTTPException(status_code=400, detail=f"Failed to read second image: {e}")
+            raise HTTPException(400, f"Failed to read second image: {e}")
 
-    # COMBINE IF NEEDED
+    # -------------------------
+    # 2. COMBINE IF NEEDED
+    # -------------------------
     if pil_img2:
         total_width = pil_img1.width + pil_img2.width
         max_height = max(pil_img1.height, pil_img2.height)
@@ -167,7 +172,65 @@ async def generate(
     else:
         combined_img = pil_img1
 
-    # INFERENCE
+    # -------------------------
+    # 3. OPTIONAL CATEGORY CLICK LOGIC
+    # -------------------------
+    if user_id and category_id:
+        try:
+            admin_client = MongoClient(os.getenv("ADMIN_MONGODB_URI"))
+            admin_db = admin_client["adminPanel"]
+
+            categories_col = admin_db.categories
+            media_clicks_col = admin_db.media_clicks
+
+            # Validate user_oid & category_oid
+            user_oid = ObjectId(user_id)
+            category_oid = ObjectId(category_id)
+
+            # Check category exists
+            category_doc = categories_col.find_one({"_id": category_oid})
+            if not category_doc:
+                raise HTTPException(400, f"Invalid category_id: {category_id}")
+
+            now = datetime.utcnow()
+
+            # Try updating existing category click
+            update_res = media_clicks_col.update_one(
+                {"userId": user_oid, "categories.categoryId": category_oid},
+                {
+                    "$set": {
+                        "updatedAt": now,
+                        "categories.$.lastClickedAt": now
+                    },
+                    "$inc": {"categories.$.click_count": 1}
+                }
+            )
+
+            # If user or category does not exist → create/push new
+            if update_res.matched_count == 0:
+                media_clicks_col.update_one(
+                    {"userId": user_oid},
+                    {
+                        "$setOnInsert": {"createdAt": now},
+                        "$set": {"updatedAt": now},
+                        "$push": {
+                            "categories": {
+                                "categoryId": category_oid,
+                                "click_count": 1,
+                                "lastClickedAt": now
+                            }
+                        }
+                    },
+                    upsert=True
+                )
+
+        except Exception as e:
+            print("CATEGORY_LOG_ERROR:", e)
+            # but do NOT stop the image generation
+
+    # -------------------------
+    # 4. HF INFERENCE
+    # -------------------------
     try:
         pil_output = hf_client.image_to_image(
             image=combined_img,
@@ -175,7 +238,6 @@ async def generate(
             model="Qwen/Qwen-Image-Edit"
         )
     except Exception as e:
-        # LOG FAILURE
         response_time_ms = round((time.time() - start_time) * 1000)
         logs_collection.insert_one({
             "timestamp": datetime.utcnow(),
@@ -187,9 +249,11 @@ async def generate(
             "error": str(e),
             "response_time_ms": response_time_ms
         })
-        raise HTTPException(status_code=500, detail=f"Inference failed: {e}")
+        raise HTTPException(500, f"Inference failed: {e}")
 
-    # SAVE OUTPUT IMAGE
+    # -------------------------
+    # 5. SAVE OUTPUT IMAGE
+    # -------------------------
     out_buf = io.BytesIO()
     pil_output.save(out_buf, format="PNG")
     out_bytes = out_buf.getvalue()
@@ -207,19 +271,20 @@ async def generate(
         }
     )
 
-    # FINAL RESPONSE TIME
     response_time_ms = round((time.time() - start_time) * 1000)
 
-    # LOG SUCCESS - THIS WAS MISSING THE STATUS AND RESPONSE_TIME_MS FIELDS!
+    # -------------------------
+    # 6. LOG SUCCESS
+    # -------------------------
     logs_collection.insert_one({
         "timestamp": datetime.utcnow(),
-        "status": "success",  # ✅ FIXED: Added status field
+        "status": "success",
         "input1_id": str(input1_id),
         "input2_id": str(input2_id) if input2_id else None,
         "output_id": str(out_id),
         "prompt": prompt,
-        "user_email": user.get("email"),  # ✅ FIXED: Changed from "user" to "user_email"
-        "response_time_ms": response_time_ms  # ✅ FIXED: Added response_time_ms field
+        "user_email": user.get("email"),
+        "response_time_ms": response_time_ms
     })
 
     return JSONResponse({
@@ -228,6 +293,116 @@ async def generate(
         "response_time_ms": response_time_ms
     })
 
+# @app.post("/generate")
+# async def generate(
+#     prompt: str = Form(...),
+#     image1: UploadFile = File(...),
+#     image2: Optional[UploadFile] = File(None),
+#     user=Depends(verify_firebase_token)
+# ):
+#     start_time = time.time()
+
+#     # PREPARE IMAGE(S)
+#     try:
+#         img1_bytes = await image1.read()
+#         pil_img1 = prepare_image(img1_bytes)
+#         input1_id = fs.put(
+#             img1_bytes,
+#             filename=image1.filename,
+#             contentType=image1.content_type,
+#             metadata={"role": "input"}
+#         )
+#     except Exception as e:
+#         raise HTTPException(status_code=400, detail=f"Failed to read first image: {e}")
+
+#     img2_bytes = None
+#     input2_id = None
+#     pil_img2 = None
+
+#     if image2:
+#         try:
+#             img2_bytes = await image2.read()
+#             pil_img2 = prepare_image(img2_bytes)
+#             input2_id = fs.put(
+#                 img2_bytes,
+#                 filename=image2.filename,
+#                 contentType=image2.content_type,
+#                 metadata={"role": "input"}
+#             )
+#         except Exception as e:
+#             raise HTTPException(status_code=400, detail=f"Failed to read second image: {e}")
+
+#     # COMBINE IF NEEDED
+#     if pil_img2:
+#         total_width = pil_img1.width + pil_img2.width
+#         max_height = max(pil_img1.height, pil_img2.height)
+#         combined_img = Image.new("RGB", (total_width, max_height))
+#         combined_img.paste(pil_img1, (0, 0))
+#         combined_img.paste(pil_img2, (pil_img1.width, 0))
+#     else:
+#         combined_img = pil_img1
+
+#     # INFERENCE
+#     try:
+#         pil_output = hf_client.image_to_image(
+#             image=combined_img,
+#             prompt=prompt,
+#             model="Qwen/Qwen-Image-Edit"
+#         )
+#     except Exception as e:
+#         # LOG FAILURE
+#         response_time_ms = round((time.time() - start_time) * 1000)
+#         logs_collection.insert_one({
+#             "timestamp": datetime.utcnow(),
+#             "status": "failure",
+#             "input1_id": str(input1_id),
+#             "input2_id": str(input2_id) if input2_id else None,
+#             "prompt": prompt,
+#             "user_email": user.get("email"),
+#             "error": str(e),
+#             "response_time_ms": response_time_ms
+#         })
+#         raise HTTPException(status_code=500, detail=f"Inference failed: {e}")
+
+#     # SAVE OUTPUT IMAGE
+#     out_buf = io.BytesIO()
+#     pil_output.save(out_buf, format="PNG")
+#     out_bytes = out_buf.getvalue()
+
+#     out_id = fs.put(
+#         out_bytes,
+#         filename=f"result_{input1_id}.png",
+#         contentType="image/png",
+#         metadata={
+#             "role": "output",
+#             "prompt": prompt,
+#             "input1_id": str(input1_id),
+#             "input2_id": str(input2_id) if input2_id else None,
+#             "user_email": user.get("email"),
+#         }
+#     )
+
+#     # FINAL RESPONSE TIME
+#     response_time_ms = round((time.time() - start_time) * 1000)
+
+#     # LOG SUCCESS - THIS WAS MISSING THE STATUS AND RESPONSE_TIME_MS FIELDS!
+#     logs_collection.insert_one({
+#         "timestamp": datetime.utcnow(),
+#         "status": "success",  # ✅ FIXED: Added status field
+#         "input1_id": str(input1_id),
+#         "input2_id": str(input2_id) if input2_id else None,
+#         "output_id": str(out_id),
+#         "prompt": prompt,
+#         "user_email": user.get("email"),  # ✅ FIXED: Changed from "user" to "user_email"
+#         "response_time_ms": response_time_ms  # ✅ FIXED: Added response_time_ms field
+#     })
+
+#     return JSONResponse({
+#         "output_image_id": str(out_id),
+#         "user": user.get("email"),
+#         "response_time_ms": response_time_ms
+#     })
+
 
 
 @app.get("/image/{image_id}")