Update main.py

main.py

@@ -105,7 +105,7 @@ async def generate_audio(text, filename):
         if os.path.exists(filename): os.remove(filename)
         return None
 
-async def safe_generate(prompt, retries=3):
+async def safe_generate(prompt, retries=5):
     if client is None: raise Exception("AI Client is not initialized.")
     for attempt in range(retries):
         try:
@@ -113,7 +113,7 @@ async def safe_generate(prompt, retries=3):
             return await loop.run_in_executor(None, lambda: client.predict(prompt=prompt, api_name="/generate_questions"))
         except Exception as e:
             if attempt == retries - 1: raise e
-            await asyncio.sleep(2)
+            await asyncio.sleep(5)
 
 def parse_question_output(raw_output: str):
     if not raw_output: return None, None
@@ -156,10 +156,10 @@ async def refill_specific_pool(track_id: int, difficulty: int, count: int, sessi
                         "created_at": firestore.SERVER_TIMESTAMP
                     })
                     success_count += 1
-                    print(f"Successfully added question {success_count}/{count}")
-                    await asyncio.sleep(3)
+                    print(f"Successfully added {track_text} question {success_count}/{count}")
+                    await asyncio.sleep(2)
         except Exception as e:
-            print(f"Refill logic error: {e}")
+            print(f"Refill error: {e}")
             await asyncio.sleep(5)
 
 # =========================================
@@ -167,12 +167,22 @@ async def refill_specific_pool(track_id: int, difficulty: int, count: int, sessi
 # =========================================
 @app.post("/generate-session")
 async def generate_session(request: GenerateSessionRequest, background_tasks: BackgroundTasks):
-    s_type, diff, t_id = request.sessionType, request.difficultyLevel, request.trackName
+    s_type, t_id = request.sessionType, request.trackName
+    
+    # FORCED ADJUSTMENT: Behavioral always uses difficulty 0
+    if s_type == 0:
+        diff = 0
+    else:
+        diff = request.difficultyLevel
+
     query = db.collection("questions_pool").where(filter=FieldFilter("session_type", "==", s_type))
     
     if s_type == 1: # Technical
         if t_id is None: raise HTTPException(status_code=400, detail="trackName is required for technical sessions.")
         query = query.where(filter=FieldFilter("track_id", "==", t_id)).where(filter=FieldFilter("difficulty", "==", diff))
+    else:
+        # For behavioral, we filter specifically by the forced difficulty 0
+        query = query.where(filter=FieldFilter("difficulty", "==", 0))
     
     docs = query.limit(10).get()
     final_questions = []
@@ -189,7 +199,7 @@ async def generate_session(request: GenerateSessionRequest, background_tasks: Ba
         snap = query.count().get()
         current_count = snap[0][0].value
         if current_count < 50:
-            print(f"Stock is low ({current_count}). Starting background refill...")
+            print(f"Stock for {('Behavioral' if s_type==0 else TECH_CATEGORIES[t_id])} is low ({current_count}). Refilling...")
             await refill_specific_pool(t_id if s_type == 1 else -1, diff, 50 - current_count, s_type)
 
     background_tasks.add_task(check_and_refill_background)
@@ -199,6 +209,38 @@ async def generate_session(request: GenerateSessionRequest, background_tasks: Ba
     
     return {"session_id": request.sessionId, "questions": final_questions}
 
+@app.get("/admin/prefill-all")
+async def prefill_all(background_tasks: BackgroundTasks):
+    """Checks and refills ALL tracks and behavioral questions to 50 items each."""
+    async def run_sync():
+        print("Starting Global Smart Prefill...")
+        
+        # 1. Sync Behavioral
+        beh_snap = db.collection("questions_pool").where(filter=FieldFilter("session_type", "==", 0)).count().get()
+        beh_count = beh_snap[0][0].value
+        if beh_count < 50:
+            print(f"Syncing Behavioral: adding {50-beh_count}")
+            await refill_specific_pool(-1, 0, 50 - beh_count, 0)
+        
+        # 2. Sync Technical
+        for t_id, t_name in TECH_CATEGORIES.items():
+            for d_id, d_name in DIFFICULTY_MAP.items():
+                query = db.collection("questions_pool")\
+                          .where(filter=FieldFilter("session_type", "==", 1))\
+                          .where(filter=FieldFilter("track_id", "==", t_id))\
+                          .where(filter=FieldFilter("difficulty", "==", d_id))
+                
+                snap = query.count().get()
+                current = snap[0][0].value
+                if current < 50:
+                    print(f"Syncing {t_name} ({d_name}): adding {50-current}")
+                    await refill_specific_pool(t_id, d_id, 50 - current, 1)
+        
+        print("Global Smart Prefill Completed!")
+
+    background_tasks.add_task(run_sync)
+    return {"message": "Global prefill process started in the background."}
+
 @app.post("/cleanup-audio")
 async def cleanup_audio(request: CleanupRequest, background_tasks: BackgroundTasks):
     def delete_job(urls):
@@ -210,22 +252,6 @@ async def cleanup_audio(request: CleanupRequest, background_tasks: BackgroundTas
     background_tasks.add_task(delete_job, request.audioUrls)
     return {"message": "Cloudinary cleanup process initiated."}
 
-@app.get("/system-cleanup")
-async def system_cleanup(background_tasks: BackgroundTasks):
-    """Manually trigger a cleanup for broken records in Firestore"""
-    def run_cleanup():
-        docs = db.collection("questions_pool").get()
-        count = 0
-        for doc in docs:
-            data = doc.to_dict()
-            if not data.get("audio_url"):
-                db.collection("questions_pool").document(doc.id).delete()
-                count += 1
-        print(f"Manual cleanup finished. Removed {count} broken records.")
-
-    background_tasks.add_task(run_cleanup)
-    return {"message": "System cleanup task started in background."}
-
 @app.get("/health")
 async def health():
     return {"status": "active", "ai_model_connected": client is not None}