Update app.py

app.py

@@ -4,6 +4,8 @@ import os
 import gradio as gr
 from dotenv import load_dotenv
 
+from huggingface_hub import HfApi
+
 from llm_interface import ERROR_503_DICT  # Import error dict
 from llm_interface import parse_qwen_response, query_qwen_endpoint
 
@@ -30,13 +32,14 @@ load_dotenv()
 
 # --- Constants ---
 HF_TOKEN = os.getenv("HF_TOKEN")
+ENDPOINT_NAME = "qwen2-5-coder-32b-instruct-pmf"
 DATASET_ID = "yjernite/spaces-privacy-reports"
 CACHE_INFO_MSG = "\n\n*(Report retrieved from cache)*"
 DEFAULT_SELECTION = "HuggingFaceTB/SmolVLM2"
 
 TRUNCATION_WARNING = """**⚠️ Warning:** The input data (code and/or prior analysis) was too long for the AI model's context limit and had to be truncated. The analysis below may be incomplete or based on partial information.\n\n---\n\n"""
 
-ERROR_503_USER_MESSAGE = """**503 Service Unavailable**: It appears that the analysis model endpoint is currently down or starting up. 
+ERROR_503_USER_MESSAGE = """**Service Unavailable**: It appears that the analysis model endpoint is currently down or starting up. 
 
 You have a few options:
 
@@ -237,9 +240,53 @@ def _run_live_analysis(space_id: str, progress=gr.Progress(track_tqdm=True)):
             gr.update(visible=True, open=False),
         )
 
-    # --- Step 2: Fetch Code Files (if not cached) ---
-    progress(2 / steps, desc="Step 2/8: Fetching code files...")
-    logging.info("Step 2/8: Fetching code files...")
+    # --- Step 2: Check Endpoint Status --- 
+    progress(2 / steps, desc="Step 2/8: Checking endpoint status...")
+    logging.info("Step 2/8: Checking endpoint status...")
+    yield (
+        gr.update(value="Checking whether model endpoint is active...", visible=True),
+        gr.update(value="", visible=True),
+        gr.update(visible=True, open=True),
+        gr.update(visible=True, open=False),
+    )
+
+    endpoint_ready = False
+    if HF_TOKEN and HF_INFERENCE_ENDPOINT_URL:
+        try:
+            api = HfApi(token=HF_TOKEN)
+            endpoint = api.get_inference_endpoint(name=ENDPOINT_NAME)
+            status = endpoint.status
+
+            logging.info(f"Endpoint '{ENDPOINT_NAME}' status: {status}")
+
+            if status == 'running':
+                endpoint_ready = True
+            else:
+                logging.warning(f"Endpoint '{ENDPOINT_NAME}' is not ready (Status: {status}).")
+                if status == 'scaledToZero':
+                    logging.info(f"Endpoint '{ENDPOINT_NAME}' is scaled to zero. Attempting to resume...")
+                    endpoint.resume()
+                msg_503 = f"The status of the Qwen2.5-Coder-32B-Instruct endpoint powering the analysis is currently: {status}\n\n" + ERROR_503_USER_MESSAGE
+                yield (
+                    gr.update(value=msg_503, visible=True),
+                    gr.update(value="", visible=False),
+                    gr.update(visible=True, open=True),
+                    gr.update(visible=False)
+                )
+                return # Stop analysis, user needs to retry
+        except Exception as e:
+            logging.error(f"Error checking endpoint status for {ENDPOINT_NAME}: {e}")
+            yield (
+                gr.update(value=f"Error checking analysis endpoint status: {e}", visible=True),
+                gr.update(value="", visible=False),
+                gr.update(visible=True, open=True),
+                gr.update(visible=False)
+            )
+            return # Stop analysis
+    
+    # --- Step 3: Fetch Code Files (if not cached) ---
+    progress(3 / steps, desc="Step 3/8: Fetching code files...")
+    logging.info("Step 3/8: Fetching code files...")
     code_files = get_space_code_files(space_id)
     if not code_files:
         error_msg = f"Could not retrieve code files for '{space_id}'. Check ID and ensure it's a public Space."
@@ -252,11 +299,11 @@ def _run_live_analysis(space_id: str, progress=gr.Progress(track_tqdm=True)):
         )
         return  # End generation on error
 
-    # --- Step 3: Generate DETAILED Privacy Report (LLM Call 1) ---
+    # --- Step 4: Generate DETAILED Privacy Report (LLM Call 1) ---
     progress(
-        3 / steps, desc="Step 3/8: Generating detailed privacy report (AI Call 1)..."
+        4 / steps, desc="Step 4/8: Generating detailed privacy report (AI Call 1)..."
     )
-    logging.info("Step 3/8: Generating detailed privacy analysis report...")
+    logging.info("Step 4/8: Generating detailed privacy analysis report...")
     yield (
         gr.update(value="Generating detailed privacy report...", visible=True),
         gr.update(value="Generating detailed privacy report via AI...", visible=True),
@@ -307,9 +354,6 @@ def _run_live_analysis(space_id: str, progress=gr.Progress(track_tqdm=True)):
         gr.update(visible=True, open=True),
     )
 
-    # --- Step 4: Extract Model IDs ---
-    progress(4 / steps, desc="Step 4/8: Extracting model IDs...")
-    logging.info("Step 4/8: Extracting potential model IDs...")
 
     # --- Step 5: Fetch Model Descriptions ---
     progress(5 / steps, desc="Step 5/8: Fetching model descriptions...")