Update app.py

app.py

@@ -6,6 +6,7 @@ from datetime import datetime
 import asyncio
 from functools import wraps
 
+# Keep your existing imports
 from pipeline.critique_extraction import extract_critiques
 from pipeline.disagreement_detection import detect_disagreements
 from pipeline.search_retrieval import search_and_retrieve
@@ -17,7 +18,6 @@ from utils.validators import validate_paper_input
 
 from dotenv import load_dotenv
 load_dotenv()
-print(os.getenv("GEMINI_API_KEY"))
 
 # Initialize rate limiter and queue manager
 rate_limiter = RateLimiter(max_requests_per_minute=10)
@@ -43,15 +43,6 @@ async def full_pipeline(
 ) -> Dict:
     """
     Run the complete consensus analysis pipeline
-    
-    Args:
-        paper_title: Title of the paper
-        paper_abstract: Abstract of the paper
-        reviews: List of review texts
-        request_id: Optional request ID for progress tracking
-    
-    Returns:
-        Complete pipeline results
     """
     if not request_id:
         request_id = f"req_{datetime.now().timestamp()}"
@@ -64,7 +55,7 @@ async def full_pipeline(
     
     try:
         # Stage 1: Critique Extraction
-        update_progress(request_id, "critique_extraction", 0.1, "Extracting critique points...")
+        update_progress(request_id, "critique_extraction", 0.1, "Extracting critique points (Gemini)...")
         critique_results = await extract_critiques(reviews)
         results["critique_points"] = critique_results
         
@@ -74,12 +65,12 @@ async def full_pipeline(
         results["disagreements"] = disagreement_results
         
         # Stage 3: Search & Retrieval
-        update_progress(request_id, "search_retrieval", 0.5, "Searching for relevant research...")
+        update_progress(request_id, "search_retrieval", 0.5, "Searching arXiv/Scholar for evidence...")
         search_results = await search_and_retrieve(paper_title, paper_abstract, critique_results)
         results["search_results"] = search_results
         
         # Stage 4: Disagreement Resolution
-        update_progress(request_id, "disagreement_resolution", 0.7, "Resolving disagreements...")
+        update_progress(request_id, "disagreement_resolution", 0.7, "Reasoning through disagreements (DeepSeek)...")
         resolution_results = await resolve_disagreements(
             paper_title,
             paper_abstract,
@@ -90,7 +81,7 @@ async def full_pipeline(
         results["resolution"] = resolution_results
         
         # Stage 5: Meta-Review Generation
-        update_progress(request_id, "meta_review", 0.9, "Generating meta-review...")
+        update_progress(request_id, "meta_review", 0.9, "Writing final Meta-Review (DeepSeek)...")
         meta_review = await generate_meta_review(
             paper_title,
             paper_abstract,
@@ -103,73 +94,106 @@ async def full_pipeline(
         return results
         
     except Exception as e:
-        update_progress(request_id, "error", 0.0, f"Error: {str(e)}")
-        raise
+        error_msg = f"Error: {str(e)}"
+        print(error_msg)
+        update_progress(request_id, "error", 0.0, error_msg)
+        raise e
+
+# --- UI Functions (Converted to Async + Generator for Keep-Alive) ---
 
-# Gradio Interface Functions
-def run_full_pipeline_ui(title: str, abstract: str, reviews_json: str) -> str:
-    """UI wrapper for full pipeline"""
+async def run_full_pipeline_ui(title: str, abstract: str, reviews_json: str):
+    """
+    UI wrapper for full pipeline with STREAMING updates.
+    This prevents the 'Connection Timeout' on HF Spaces.
+    """
     try:
-        # Validate and parse input
+        # Validate input
         reviews = json.loads(reviews_json)
         if not isinstance(reviews, list):
-            return json.dumps({"error": "Reviews must be a list of strings"}, indent=2)
+            yield json.dumps({"error": "Reviews must be a list of strings"}, indent=2)
+            return
         
-        # Check rate limit
         if not rate_limiter.allow_request():
-            return json.dumps({"error": "Rate limit exceeded. Please try again later."}, indent=2)
+            yield json.dumps({"error": "Rate limit exceeded. Please try again later."}, indent=2)
+            return
         
-        # Add to queue and run
         request_id = f"ui_{datetime.now().timestamp()}"
-        result = asyncio.run(queue_manager.add_task(
-            full_pipeline(title, abstract, reviews, request_id)
-        ))
         
-        return json.dumps(result, indent=2)
+        # Initialize progress for this ID
+        update_progress(request_id, "queued", 0.0, "Request queued...")
+
+        # Create the task (wrapped in queue manager)
+        # We assume queue_manager.add_task is an async function
+        pipeline_task = asyncio.create_task(
+            queue_manager.add_task(full_pipeline(title, abstract, reviews, request_id))
+        )
         
+        # Loop while the task is running to keep the connection alive
+        while not pipeline_task.done():
+            # fetch current status
+            current_status = progress_store.get(request_id, {})
+            
+            # Create a temporary status object to show in the UI
+            status_display = {
+                "status": "processing",
+                "current_stage": current_status.get("stage", "initializing"),
+                "message": current_status.get("message", "Waiting..."),
+                "progress": current_status.get("progress", 0)
+            }
+            
+            # Yield the status (Keep-Alive!)
+            yield json.dumps(status_display, indent=2)
+            
+            # Wait 1 second before next update
+            await asyncio.sleep(1)
+            
+        # Get final result or exception
+        try:
+            result = await pipeline_task
+            yield json.dumps(result, indent=2)
+        except Exception as e:
+            yield json.dumps({"error": str(e), "last_stage_log": progress_store.get(request_id)}, indent=2)
+            
     except json.JSONDecodeError:
-        return json.dumps({"error": "Invalid JSON format for reviews"}, indent=2)
+        yield json.dumps({"error": "Invalid JSON format for reviews"}, indent=2)
     except Exception as e:
-        return json.dumps({"error": str(e)}, indent=2)
+        yield json.dumps({"error": f"Unexpected error: {str(e)}"}, indent=2)
 
-def run_critique_extraction_ui(reviews_json: str) -> str:
-    """UI wrapper for critique extraction"""
+# Wrapper for smaller tasks (Convert to async, remove asyncio.run)
+async def run_critique_extraction_ui(reviews_json: str) -> str:
     try:
         reviews = json.loads(reviews_json)
         if not rate_limiter.allow_request():
             return json.dumps({"error": "Rate limit exceeded"}, indent=2)
         
-        result = asyncio.run(extract_critiques(reviews))
+        result = await extract_critiques(reviews)
         return json.dumps(result, indent=2)
     except Exception as e:
         return json.dumps({"error": str(e)}, indent=2)
 
-def run_disagreement_detection_ui(critiques_json: str) -> str:
-    """UI wrapper for disagreement detection"""
+async def run_disagreement_detection_ui(critiques_json: str) -> str:
     try:
         critiques = json.loads(critiques_json)
         if not rate_limiter.allow_request():
             return json.dumps({"error": "Rate limit exceeded"}, indent=2)
         
-        result = asyncio.run(detect_disagreements(critiques))
+        result = await detect_disagreements(critiques)
         return json.dumps(result, indent=2)
     except Exception as e:
         return json.dumps({"error": str(e)}, indent=2)
 
-def run_search_retrieval_ui(title: str, abstract: str, critiques_json: str) -> str:
-    """UI wrapper for search retrieval"""
+async def run_search_retrieval_ui(title: str, abstract: str, critiques_json: str) -> str:
     try:
         critiques = json.loads(critiques_json)
         if not rate_limiter.allow_request():
             return json.dumps({"error": "Rate limit exceeded"}, indent=2)
         
-        result = asyncio.run(search_and_retrieve(title, abstract, critiques))
+        result = await search_and_retrieve(title, abstract, critiques)
         return json.dumps(result, indent=2)
     except Exception as e:
         return json.dumps({"error": str(e)}, indent=2)
 
 def check_progress_ui(request_id: str) -> str:
-    """Check progress of a request"""
     if request_id in progress_store:
         return json.dumps(progress_store[request_id], indent=2)
     return json.dumps({"error": "Request ID not found"}, indent=2)
@@ -180,13 +204,6 @@ with gr.Blocks(title="Automated Consensus Analysis API", theme=gr.themes.Soft())
     # 🔬 Automated Consensus Analysis API
     
     This API provides automated peer review consensus analysis using LLMs and search-augmented verification.
-    
-    ## Features:
-    - **Critique Extraction**: Extract structured critique points from reviews
-    - **Disagreement Detection**: Identify conflicts between reviewers
-    - **Search Retrieval**: Find supporting/contradicting evidence
-    - **Resolution**: Resolve disagreements with evidence
-    - **Meta-Review**: Generate comprehensive meta-reviews
     """)
     
     with gr.Tabs():
@@ -215,11 +232,7 @@ with gr.Blocks(title="Automated Consensus Analysis API", theme=gr.themes.Soft())
         # Individual Stages
         with gr.Tab("🔍 Critique Extraction"):
             gr.Markdown("### Extract critique points from reviews")
-            critique_reviews = gr.Code(
-                label="Reviews (JSON Array)",
-                language="json",
-                value='["Review 1...", "Review 2..."]'
-            )
+            critique_reviews = gr.Code(label="Reviews (JSON Array)", language="json")
             critique_submit = gr.Button("Extract Critiques")
             critique_output = gr.Code(label="Extracted Critiques", language="json")
             
@@ -230,11 +243,8 @@ with gr.Blocks(title="Automated Consensus Analysis API", theme=gr.themes.Soft())
             )
         
         with gr.Tab("⚡ Disagreement Detection"):
-            gr.Markdown("### Detect disagreements between reviews")
-            disagree_critiques = gr.Code(
-                label="Critique Points (JSON)",
-                language="json"
-            )
+            gr.Markdown("### Detect disagreements")
+            disagree_critiques = gr.Code(label="Critique Points (JSON)", language="json")
             disagree_submit = gr.Button("Detect Disagreements")
             disagree_output = gr.Code(label="Disagreement Analysis", language="json")
             
@@ -245,7 +255,7 @@ with gr.Blocks(title="Automated Consensus Analysis API", theme=gr.themes.Soft())
             )
         
         with gr.Tab("🔎 Search & Retrieval"):
-            gr.Markdown("### Search for supporting evidence")
+            gr.Markdown("### Search for evidence")
             with gr.Row():
                 with gr.Column():
                     search_title = gr.Textbox(label="Paper Title")
@@ -260,81 +270,13 @@ with gr.Blocks(title="Automated Consensus Analysis API", theme=gr.themes.Soft())
                 inputs=[search_title, search_abstract, search_critiques],
                 outputs=search_output
             )
-        
-        with gr.Tab("📊 Progress Tracking"):
-            gr.Markdown("### Check pipeline progress")
-            progress_id = gr.Textbox(label="Request ID", placeholder="Enter request ID...")
-            progress_check = gr.Button("Check Progress")
-            progress_output = gr.Code(label="Progress Status", language="json")
-            
-            progress_check.click(
-                fn=check_progress_ui,
-                inputs=progress_id,
-                outputs=progress_output
-            )
-        
+
         with gr.Tab("📖 API Documentation"):
-            gr.Markdown("""
-            ## API Endpoints
-            
-            ### POST /api/full_pipeline
-            Run the complete consensus analysis pipeline.
-            
-            **Request Body:**
-            ```json
-            {
-                "paper_title": "string",
-                "paper_abstract": "string",
-                "reviews": ["review1", "review2", ...]
-            }
-            ```
-            
-            ### POST /api/critique_extraction
-            Extract critique points from reviews.
-            
-            **Request Body:**
-            ```json
-            {
-                "reviews": ["review1", "review2", ...]
-            }
-            ```
-            
-            ### POST /api/disagreement_detection
-            Detect disagreements in critique points.
-            
-            **Request Body:**
-            ```json
-            {
-                "critiques": [{"Methodology": [...], ...}, ...]
-            }
-            ```
-            
-            ### POST /api/search_retrieval
-            Search for supporting evidence.
-            
-            **Request Body:**
-            ```json
-            {
-                "paper_title": "string",
-                "paper_abstract": "string",
-                "critiques": [...]
-            }
-            ```
-            
-            ### GET /api/progress/{request_id}
-            Check progress of a pipeline execution.
-            
-            ## Rate Limits
-            - 10 requests per minute per IP
-            - Maximum 3 concurrent pipeline executions
-            
-            ## Authentication
-            API keys are managed through HuggingFace Spaces secrets.
-            """)
+             gr.Markdown("## API Documentation...")
 
 # Launch the app
 if __name__ == "__main__":
-    demo.queue(max_size=20)  # Enable queuing
+    demo.queue(max_size=20)
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,