app.py

import gradio as gr
import os
import sys
from typing import Dict, Any
import json

# Import the ENHANCED clinical diagnosis exploration functions
from clinical_diag_exp_enhanced import run_clinical_diagnosis_exploration_enhanced

# Get API key from environment variable
try:
    API_KEY = os.environ['OPENAI_API_KEY']
    if not API_KEY:
        raise KeyError("OPENAI_API_KEY is empty")
except KeyError:
    print("ERROR: OPENAI_API_KEY environment variable not found or empty.")
    print("Please set your OpenAI API key as an environment variable:")
    print("export OPENAI_API_KEY='your-api-key-here'")
    sys.exit(1)


def run_clinical_exploration_enhanced(diagnosis_text: str, max_results: int, max_claims: int) -> Dict[str, Any]:
    """
    Run the ENHANCED clinical diagnosis exploration pipeline.
    """
    if not diagnosis_text or diagnosis_text.strip() == "":
        return {
            "error": "Please enter a clinical diagnosis or note to analyze."
        }
    
    try:
        # Run the ENHANCED clinical diagnosis exploration
        results = run_clinical_diagnosis_exploration_enhanced(
            diagnosis_text=diagnosis_text.strip(),
            max_results_per_claim=max_results,
            max_claims=max_claims
        )
        
        # Format the results for display
        formatted_results = {
            "diagnosis": results.get("diagnosis", "No diagnosis extracted"),
            "revised_note": results.get("revised_note", "No revised note generated"),
            "total_claims": results.get("claims_researched", 0),
            "total_articles": results.get("total_articles", 0),
            "total_llm_calls": results.get("total_llm_calls", 0),
            "claims_summary": "",
            "references": results.get("references", "No references available."),
            "missing_information": results.get("missing_information", []),
            "enhancement": results.get("enhancement", "Enhanced features"),
        }
        
        # Create a detailed summary of claims with full research summaries and references
        claims = results.get("claims", [])
        claims_summary = "## Claims Analyzed:\n\n"
        for i, claim in enumerate(claims, 1):
            claims_summary += f"### Claim {i}: {claim['claim_text']}\n"
            claims_summary += f"- **Type:** {claim['claim_type']}\n"
            claims_summary += f"- **Evidence Level:** {claim['evidence_level']}\n"
            claims_summary += f"- **Research Status:** {claim['research_status']}\n"
            claims_summary += f"- **Evidence Quality:** {claim.get('evidence_quality', 'unknown')}\n"
            claims_summary += f"- **Articles Found:** {len(claim['abstracts'])}\n"
            claims_summary += f"- **PubMed Query:** `{claim['pubmed_query']}`\n"
            claims_summary += f"- **Query Strategy:** {claim.get('query_strategy', 'Standard query')}\n\n"
            
            # Add the detailed summary with inline citations
            claims_summary += f"**Detailed Summary:** {claim['detailed_summary']}\n\n"
            claims_summary += f"**Brief Summary:** {claim['evidence_summary']}\n\n"
            
            # Add critical assessment
            if claim.get('critical_assessment'):
                claims_summary += f"**Critical Assessment:** {claim['critical_assessment']}\n\n"
            
            # Add supporting/contradicting evidence if available
            if claim.get('supporting_evidence'):
                claims_summary += "**Supporting Evidence:**\n"
                for evidence in claim['supporting_evidence']:
                    claims_summary += f"- {evidence}\n"
                claims_summary += "\n"
            
            if claim.get('contradicting_evidence'):
                claims_summary += "**Contradicting Evidence:**\n"
                for evidence in claim['contradicting_evidence']:
                    claims_summary += f"- {evidence}\n"
                claims_summary += "\n"
            
            # Add methodological issues if available
            if claim.get('methodological_issues'):
                claims_summary += "**Methodological Issues:**\n"
                for issue in claim['methodological_issues']:
                    claims_summary += f"- {issue}\n"
                claims_summary += "\n"
            
            # Add the claim-specific references
            claims_summary += f"**Claim-Specific References:**\n{claim['claim_references']}\n\n"
            
            claims_summary += "---\n\n"
        
        formatted_results["claims_summary"] = claims_summary
        
        return formatted_results
        
    except Exception as e:
        return {
            "error": f"Error running clinical exploration: {str(e)}"
        }

def create_interface():
    """
    Create the Clinical Diagnosis Explorer interface for extracting diagnostic claims 
    and grounding them in scientific literature with PubMed search.
    """
    
    with gr.Blocks(title="Clinical Diagnosis Explorer", theme=gr.themes.Soft()) as iface:
        
        gr.Markdown("# Clinical Diagnosis Explorer")
        gr.Markdown("**Extract diagnostic claims and ground them in the scientific literature with PubMed search**")
        gr.Markdown("*Features: Diagnosis-first approach, fuzzy search, LLM analysis, evidence-based clinical notes*")
        
        with gr.Tab("Clinical Analysis"):
            gr.Markdown("## Analyze Clinical Diagnosis")
            gr.Markdown("Enter a clinical diagnosis or clinical note to extract diagnostic claims and validate each against medical literature using PubMed search.")
            
            diagnosis_input = gr.Textbox(
                label="Clinical Diagnosis or Note",
                placeholder="Enter the clinical diagnosis, clinical note, or diagnostic reasoning here...\n\nExample: 28-year-old right-handed female with an 18-month history of progressive word comprehension deficits, surface dyslexia, and difficulty naming common objects while speech remains fluent and grammatically intact. MRI reveals bilateral anterior temporal lobe atrophy. Diagnosis: Semantic Variant Primary Progressive Aphasia (svPPA).",
                lines=12
            )
            
            with gr.Row():
                max_results = gr.Slider(
                    minimum=5,
                    maximum=20,
                    value=12,
                    step=1,
                    label="Max Results per Claim",
                    info="Number of PubMed articles to analyze per claim (increased for better analysis)"
                )
                
                max_claims = gr.Slider(
                    minimum=3,
                    maximum=12,
                    value=8,
                    step=1,
                    label="Max Claims to Extract",
                    info="Maximum number of claims to extract from diagnosis"
                )
            
            analyze_btn = gr.Button("🔍 Extract Claims & Search Literature", variant="primary")

            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Results")
                    results_output = gr.JSON(label="Analysis Results", visible=False)
                    error_output = gr.Textbox(label="Error", visible=False)
                
                with gr.Column():
                    gr.Markdown("### Analysis Summary")
                    summary_output = gr.Markdown(label="Analysis Summary")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Primary Diagnosis")
                    diagnosis_output = gr.Markdown(label="Extracted Diagnosis")
                
                with gr.Column():
                    gr.Markdown("### Missing Information")
                    missing_info_output = gr.Markdown(label="Missing Information")
            
            with gr.Row():
                gr.Markdown("### Revised Clinical Note")
                revised_note_output = gr.Markdown(label="Evidence-Based Clinical Note")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("### Claims Analysis")
                    claims_output = gr.Markdown(label="Detailed Claims Analysis")
                
                with gr.Column():
                    gr.Markdown("### References")
                    references_output = gr.Markdown(label="Research References")
            
            def process_enhanced_analysis(diagnosis, max_res, max_clms):
                results = run_clinical_exploration_enhanced(diagnosis, max_res, max_clms)
                
                if "error" in results:
                    return {
                        results_output: gr.update(value=None, visible=False),
                        error_output: gr.update(value=results["error"], visible=True),
                        summary_output: gr.update(value=""),
                        diagnosis_output: gr.update(value=""),
                        missing_info_output: gr.update(value=""),
                        revised_note_output: gr.update(value=""),
                        claims_output: gr.update(value=""),
                        references_output: gr.update(value=""),
                    }
                else:
                    # Format analysis summary
                    summary_text = f"""
## Analysis Summary

**Total Claims Researched:** {results.get("total_claims", 0)}
**Total Articles Found:** {results.get("total_articles", 0)}
**Total LLM Calls:** {results.get("total_llm_calls", 0)}
**Enhanced Features:** {results.get("enhancement", "Enhanced analysis")}
"""
                    
                    # Format diagnosis
                    diagnosis_text = f"""
## Primary Diagnosis

**Extracted Diagnosis:** {results.get("diagnosis", "No diagnosis extracted")}

This diagnosis was identified as the primary clinical condition and used as the foundation for claim extraction and PubMed query generation.
"""
                    
                    # Format missing information
                    missing_info = results.get("missing_information", [])
                    missing_text = "## Missing Information\n\n"
                    if missing_info:
                        for i, info in enumerate(missing_info, 1):
                            missing_text += f"{i}. ❓ {info}\n"
                    else:
                        missing_text += "No missing information identified."
                    
                    return {
                        results_output: gr.update(value=results, visible=True),
                        error_output: gr.update(value="", visible=False),
                        summary_output: gr.update(value=summary_text),
                        diagnosis_output: gr.update(value=diagnosis_text),
                        missing_info_output: gr.update(value=missing_text),
                        revised_note_output: gr.update(value=results["revised_note"]),
                        claims_output: gr.update(value=results["claims_summary"]),
                        references_output: gr.update(value=results["references"]),
                    }
            
            analyze_btn.click(
                fn=process_enhanced_analysis,
                inputs=[diagnosis_input, max_results, max_claims],
                outputs=[results_output, error_output, summary_output, diagnosis_output, missing_info_output, revised_note_output, claims_output, references_output]
            )
            
    return iface

if __name__ == "__main__":
    iface = create_interface()
    iface.launch()