app.py

import gradio as gr
import json
import os
from typing import Dict, List, Optional
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
from pipeline.disagreement_resolution import resolve_disagreements
from pipeline.meta_review import generate_meta_review
from utils.rate_limiter import RateLimiter
from utils.queue_manager import QueueManager
from utils.validators import validate_paper_input

from dotenv import load_dotenv
load_dotenv()

# Initialize rate limiter and queue manager
rate_limiter = RateLimiter(max_requests_per_minute=10)
queue_manager = QueueManager(max_concurrent=3)

# Progress tracking
progress_store = {}

# Example data for quick testing
EXAMPLE_PAPER_TITLE = "Learning Disentangled Representations for CounterFactual Regression"

EXAMPLE_PAPER_ABSTRACT = """We consider the challenge of estimating treatment effects from observational data; and point out that, in general, only some factors based on the observed covariates X contribute to selection of the treatment T, and only some to determining the outcomes Y. We model this by considering three underlying sources of {X, T, Y} and show that explicitly modeling these sources offers great insight to guide designing models that better handle selection bias. This paper is an attempt to conceptualize this line of thought and provide a path to explore it further.
In this work, we propose an algorithm to (1) identify disentangled representations of the above-mentioned underlying factors from any given observational dataset D and (2) leverage this knowledge to reduce, as well as account for, the negative impact of selection bias on estimating the treatment effects from D. Our empirical results show that the proposed method achieves state-of-the-art performance in both individual and population based evaluation measures."""

EXAMPLE_REVIEWS = [
    """Summary:
   The authors consider the problem of estimating average treatment effects when observed X and treatment T causes Y. Observational data for X,T,Y is available and strong ignorability is assumed. Previous work (Shalit et al 2017) introduced learning a representation that is invariant in distribution across treatment and control groups and using that with treatment to estimate Y. However, authors point out that this representation being forced to be invariant still does not drive the selection bias to zero. A follow up work (Hassanpour and Greiner 2019) - corrects for this by using additional importance weighting that estimates the treatment selection bias given the learnt representation. However, the authors point out even this is not complete in general, as X could be determined by three latent factors, one that is the actual confounder between treatment and outcome and the other that affects only the outcome and the other that affects only the treatment. Therefore, the authors propose to have three representations and enforce independence between representation that solely determines outcome and the treatment and make other appropriate terms depend on the respective latent factors. This gives a modified objective with respect to these two prior works.

The authors implement optimize this joint system on synthetic and real world datasets. They show that they outperform all these previous works because of explicitly accounting for confounder, latent factors that solely control only outcome and treatment assignment respectively. 

Pros:
  This paper directly addresses the problems due to Shalit 2017 that are still left open. The experimental results seems convincing on standard benchmarks. 

I vote for accepting the paper. I don't have many concerns about this paper.

Cons:
  - I have one question for the authors - if T and Y(0),Y(1) are independent given X is assumed, then how are we sure that the composite representations (of the three latent factors) are going to necessarily satisfy ignorability provably ?? I guess this cannot be formally established. It would be great for the authors to comment on this.""",
    
    """The paper proposes a new way of estimating treatment effects from observational data, that decouples (disentangles) the observed covariates X into three sets: covariates that contributed to the selection of the treatment T, covariates that cause the outcome Y and covariates that do both. The authors show that by leveraging this additional structure they can improve upon existing methods in both ITE and ATE

The main contributions of the paper are:
* Highlighting the importance of differentiating between treatment and outcome inducing factors and proposing an algorithm to detect the two
* Creating a joint optimisation model that contains the factual loss, the cross entropy (treatment) loss and the imbalance loss

Overall, I like the paper quite a lot, I find it well-written and clearly motivated with a very nice experimental section that it is designed around understanding the behaviour of the proposed model.

In terms of suggestions, I think it will be very interesting to link the approaches using invariant causal representations with existing work in the Counterfactual Risk Minimization [1] literature and to mutualise the experimental setup. 

[1] Swaminathan, Adith, and Thorsten Joachims. "Counterfactual risk minimization: Learning from logged bandit feedback." International Conference on Machine Learning. 2015.""",
    
    """The paper proposes an algorithm that identifies disentangled representation to find out an individual treatment effect. A very specific model that tries to find out the underlying dynamics of such a problem is proposed and is learned by minimizing a suggested objective that takes the strengths of previous approaches. The method is demonstrated in a synthetic dataset and IHDP dataset and shown to outperform other previous methods by a large margin.

My initial review was negative, but I changed my mind after reading a few papers in this area. It seems that explicit learning of underlying factors that are described in (Hassanpour & Greiner, 2019) is a nice idea and works well. My only concern is that the paper has a lot of overlap with (Hassanpour & Greiner, 2019), even using identical figures. I am not sure whether it is OK."""
]

def update_progress(request_id: str, stage: str, progress: float, message: str):
    """Update progress for a request"""
    progress_store[request_id] = {
        "stage": stage,
        "progress": progress,
        "message": message,
        "timestamp": datetime.now().isoformat()
    }

async def full_pipeline(
    paper_title: str,
    paper_abstract: str,
    reviews: List[str],
    request_id: Optional[str] = None
) -> Dict:
    """
    Run the complete consensus analysis pipeline
    """
    if not request_id:
        request_id = f"req_{datetime.now().timestamp()}"
    
    results = {
        "request_id": request_id,
        "paper_title": paper_title,
        "paper_abstract": paper_abstract
    }
    
    try:
        # Stage 1: Critique Extraction
        update_progress(request_id, "critique_extraction", 0.1, "Extracting critique points (Gemini)...")
        critique_results = await extract_critiques(reviews)
        results["critique_points"] = critique_results
        
        # Stage 2: Disagreement Detection
        update_progress(request_id, "disagreement_detection", 0.3, "Detecting disagreements...")
        disagreement_results = await detect_disagreements(critique_results)
        results["disagreements"] = disagreement_results
        
        # Stage 3: Search & Retrieval
        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, "Reasoning through disagreements (DeepSeek)...")
        resolution_results = await resolve_disagreements(
            paper_title,
            paper_abstract,
            disagreement_results,
            critique_results,
            search_results
        )
        results["resolution"] = resolution_results
        
        # Stage 5: Meta-Review Generation
        update_progress(request_id, "meta_review", 0.9, "Writing final Meta-Review (DeepSeek)...")
        meta_review = await generate_meta_review(
            paper_title,
            paper_abstract,
            resolution_results,
            search_results
        )
        results["meta_review"] = meta_review
        
        update_progress(request_id, "complete", 1.0, "Pipeline complete!")
        return results
        
    except Exception as e:
        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) ---

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 input
        reviews = json.loads(reviews_json)
        if not isinstance(reviews, list):
            yield json.dumps({"error": "Reviews must be a list of strings"}, indent=2)
            return
        
        if not rate_limiter.allow_request():
            yield json.dumps({"error": "Rate limit exceeded. Please try again later."}, indent=2)
            return
        
        request_id = f"ui_{datetime.now().timestamp()}"
        
        # 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:
        yield json.dumps({"error": "Invalid JSON format for reviews"}, indent=2)
    except Exception as e:
        yield json.dumps({"error": f"Unexpected error: {str(e)}"}, indent=2)

# 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 = await extract_critiques(reviews)
        return json.dumps(result, indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)}, indent=2)

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 = await detect_disagreements(critiques)
        return json.dumps(result, indent=2)
    except Exception as e:
        return json.dumps({"error": str(e)}, indent=2)

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 = 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:
    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)

def load_example():
    """Load example paper and reviews into the form"""
    return (
        EXAMPLE_PAPER_TITLE,
        EXAMPLE_PAPER_ABSTRACT,
        json.dumps(EXAMPLE_REVIEWS, indent=2)
    )

def load_example_critiques():
    """Load example reviews for critique extraction"""
    return json.dumps(EXAMPLE_REVIEWS, indent=2)

# Build Gradio Interface
with gr.Blocks(title="Automated Consensus Analysis API", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🔬 Automated Consensus Analysis API
    
    This API provides automated peer review consensus analysis using LLMs and search-augmented verification.
    
    **Pipeline stages:**
    1. 🔍 **Critique Extraction** - Extract structured critique points from reviews (Gemini)
    2. ⚡ **Disagreement Detection** - Identify conflicts between reviewers
    3. 🔎 **Search & Retrieval** - Find supporting evidence from academic sources
    4. 🧠 **Disagreement Resolution** - AI-powered resolution with reasoning (DeepSeek-R1)
    5. 📝 **Meta-Review Generation** - Comprehensive synthesis of all analyses
    """)
    
    with gr.Tabs():
        # Full Pipeline Tab
        with gr.Tab("📋 Full Pipeline"):
            gr.Markdown("### Run the complete analysis pipeline")
            with gr.Row():
                with gr.Column():
                    full_title = gr.Textbox(
                        label="Paper Title", 
                        placeholder="Enter paper title...",
                        lines=2
                    )
                    full_abstract = gr.Textbox(
                        label="Paper Abstract", 
                        lines=8, 
                        placeholder="Enter paper abstract..."
                    )
                    full_reviews = gr.Code(
                        label="Reviews (JSON Array)",
                        language="json",
                        value='["Review 1 text...", "Review 2 text..."]',
                        lines=15
                    )
                    
                    with gr.Row():
                        load_example_btn = gr.Button("📥 Load Example", variant="secondary")
                        full_submit = gr.Button("🚀 Run Full Pipeline", variant="primary")
                    
                    gr.Markdown("""
                    💡 **Tip:** Click "Load Example" to populate the form with a sample ICLR 2020 paper 
                    on Counterfactual Regression with 3 peer reviews.
                    """)
                    
                with gr.Column():
                    full_output = gr.Code(label="Results", language="json", lines=30)
            
            load_example_btn.click(
                fn=load_example,
                inputs=[],
                outputs=[full_title, full_abstract, full_reviews]
            )
            
            full_submit.click(
                fn=run_full_pipeline_ui,
                inputs=[full_title, full_abstract, full_reviews],
                outputs=full_output
            )
        
        # Individual Stages
        with gr.Tab("🔍 Critique Extraction"):
            gr.Markdown("### Extract critique points from reviews")
            gr.Markdown("Extract structured critique points categorized by: Methodology, Experiments, Clarity, Significance, Novelty")
            
            critique_reviews = gr.Code(
                label="Reviews (JSON Array)", 
                language="json",
                lines=15
            )
            
            with gr.Row():
                load_critique_example_btn = gr.Button("📥 Load Example", variant="secondary")
                critique_submit = gr.Button("Extract Critiques", variant="primary")
            
            critique_output = gr.Code(label="Extracted Critiques", language="json", lines=20)
            
            load_critique_example_btn.click(
                fn=load_example_critiques,
                inputs=[],
                outputs=critique_reviews
            )
            
            critique_submit.click(
                fn=run_critique_extraction_ui,
                inputs=critique_reviews,
                outputs=critique_output
            )
        
        with gr.Tab("⚡ Disagreement Detection"):
            gr.Markdown("### Detect disagreements between reviewers")
            gr.Markdown("Compares critique points from multiple reviews and identifies conflicts with disagreement scores (0-1).")
            
            disagree_critiques = gr.Code(
                label="Critique Points (JSON) - Output from Critique Extraction", 
                language="json",
                lines=15
            )
            disagree_submit = gr.Button("Detect Disagreements", variant="primary")
            disagree_output = gr.Code(label="Disagreement Analysis", language="json", lines=20)
            
            disagree_submit.click(
                fn=run_disagreement_detection_ui,
                inputs=disagree_critiques,
                outputs=disagree_output
            )
        
        with gr.Tab("🔎 Search & Retrieval"):
            gr.Markdown("### Search for supporting evidence")
            gr.Markdown("Searches Semantic Scholar, arXiv, and Tavily for state-of-the-art research and evidence to validate critiques.")
            
            with gr.Row():
                with gr.Column():
                    search_title = gr.Textbox(label="Paper Title", lines=2)
                    search_abstract = gr.Textbox(label="Paper Abstract", lines=5)
                    search_critiques = gr.Code(
                        label="Critiques (JSON) - Output from Critique Extraction", 
                        language="json",
                        lines=10
                    )
                    search_submit = gr.Button("Search Evidence", variant="primary")
                with gr.Column():
                    search_output = gr.Code(label="Search Results", language="json", lines=25)
            
            search_submit.click(
                fn=run_search_retrieval_ui,
                inputs=[search_title, search_abstract, search_critiques],
                outputs=search_output
            )

        with gr.Tab("📖 API Documentation"):
            gr.Markdown("""
            ## API Endpoints
            
            ### Full Pipeline
            **Endpoint**: `/api/full_pipeline`  
            **Method**: POST  
            
            ```json
            {
              "paper_title": "Your Paper Title",
              "paper_abstract": "Your paper abstract...",
              "reviews": ["Review 1 text...", "Review 2 text..."]
            }
            ```
            
            ### Individual Stages
            
            | Endpoint | Description |
            |----------|-------------|
            | `/api/critique_extraction` | Extract critique points from reviews |
            | `/api/disagreement_detection` | Detect disagreements between critiques |
            | `/api/search_retrieval` | Search for supporting evidence |
            
            ### Rate Limits
            - **10 requests per minute** per client
            - **Maximum 3 concurrent** pipeline executions
            
            ### Example Python Usage
            ```python
            import requests
            
            response = requests.post(
                "https://your-space.hf.space/api/full_pipeline",
                json={
                    "paper_title": "Novel Approach to X",
                    "paper_abstract": "We propose...",
                    "reviews": ["Review 1...", "Review 2..."]
                }
            )
            result = response.json()
            print(result["meta_review"])
            ```
            """)
        
        with gr.Tab("ℹ️ About"):
            gr.Markdown("""
            ## About This Tool
            
            This API provides automated peer review consensus analysis using state-of-the-art LLMs 
            and search-augmented verification.
            
            ### Models Used
            - **Gemini 2.5 Flash Lite** - Critique extraction and disagreement detection
            - **DeepSeek-R1** - Disagreement resolution and meta-review generation (with reasoning)
            
            ### Search Sources
            - Semantic Scholar
            - arXiv
            - Tavily Search
            - Google Scholar (via SerpAPI)
            
            ### Example Paper
            The example paper is from **ICLR 2020**: *"Learning Disentangled Representations for 
            CounterFactual Regression"* which addresses estimating treatment effects from observational 
            data using disentangled representations.
            
            ### License
            MIT License - See repository for details.
            """)

# Launch the app
if __name__ == "__main__":
    demo.queue(max_size=20)
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )