Update app.py

app.py

@@ -1,53 +1,37 @@
-import sys
 import os
 import pandas as pd
 import pdfplumber
-import json
+import re
 import gradio as gr
-from typing import List
+from typing import List, Dict
 from concurrent.futures import ThreadPoolExecutor, as_completed
 import hashlib
-import shutil
-import re
-import psutil
-import subprocess
 
-# Persistent directory
+# Persistent directories
 persistent_dir = "/data/hf_cache"
 os.makedirs(persistent_dir, exist_ok=True)
-
-model_cache_dir = os.path.join(persistent_dir, "txagent_models")
-tool_cache_dir = os.path.join(persistent_dir, "tool_cache")
 file_cache_dir = os.path.join(persistent_dir, "cache")
 report_dir = os.path.join(persistent_dir, "reports")
-vllm_cache_dir = os.path.join(persistent_dir, "vllm_cache")
-
-for directory in [model_cache_dir, tool_cache_dir, file_cache_dir, report_dir, vllm_cache_dir]:
+for directory in [file_cache_dir, report_dir]:
     os.makedirs(directory, exist_ok=True)
 
-os.environ["HF_HOME"] = model_cache_dir
-os.environ["TRANSFORMERS_CACHE"] = model_cache_dir
-os.environ["VLLM_CACHE_DIR"] = vllm_cache_dir
-os.environ["TOKENIZERS_PARALLELISM"] = "false"
-os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
-
-current_dir = os.path.dirname(os.path.abspath(__file__))
-src_path = os.path.abspath(os.path.join(current_dir, "src"))
-sys.path.insert(0, src_path)
-
-from txagent.txagent import TxAgent
-
-MEDICAL_KEYWORDS = {'diagnosis', 'assessment', 'plan', 'results', 'medications',
-                    'allergies', 'summary', 'impression', 'findings', 'recommendations'}
+# Medical keywords for PDF extraction
+MEDICAL_KEYWORDS = {
+    'diagnosis', 'assessment', 'plan', 'results', 'medications',
+    'allergies', 'summary', 'impression', 'findings', 'recommendations'
+}
 
 def sanitize_utf8(text: str) -> str:
+    """Sanitize text to handle UTF-8 encoding issues."""
     return text.encode("utf-8", "ignore").decode("utf-8")
 
 def file_hash(path: str) -> str:
+    """Generate MD5 hash of a file."""
     with open(path, "rb") as f:
         return hashlib.md5(f.read()).hexdigest()
 
 def extract_priority_pages(file_path: str) -> str:
+    """Extract text from PDF pages, prioritizing those with medical keywords."""
     try:
         text_chunks = []
         with pdfplumber.open(file_path) as pdf:
@@ -59,82 +43,112 @@ def extract_priority_pages(file_path: str) -> str:
     except Exception as e:
         return f"PDF processing error: {str(e)}"
 
-def convert_file_to_json(file_path: str, file_type: str) -> str:
+def convert_file_to_text(file_path: str, file_type: str) -> str:
+    """Convert supported file types to text, caching results."""
     try:
         h = file_hash(file_path)
-        cache_path = os.path.join(file_cache_dir, f"{h}.json")
+        cache_path = os.path.join(file_cache_dir, f"{h}.txt")
         if os.path.exists(cache_path):
             with open(cache_path, "r", encoding="utf-8") as f:
                 return f.read()
 
         if file_type == "pdf":
             text = extract_priority_pages(file_path)
-            result = json.dumps({"filename": os.path.basename(file_path), "content": text, "status": "initial"})
         elif file_type == "csv":
             df = pd.read_csv(file_path, encoding_errors="replace", header=None, dtype=str,
                              skip_blank_lines=False, on_bad_lines="skip")
-            content = df.fillna("").astype(str).values.tolist()
-            result = json.dumps({"filename": os.path.basename(file_path), "rows": content})
+            text = "\n".join(df.fillna("").astype(str).agg(" ".join, axis=1))
         elif file_type in ["xls", "xlsx"]:
             try:
                 df = pd.read_excel(file_path, engine="openpyxl", header=None, dtype=str)
             except Exception:
                 df = pd.read_excel(file_path, engine="xlrd", header=None, dtype=str)
-            content = df.fillna("").astype(str).values.tolist()
-            result = json.dumps({"filename": os.path.basename(file_path), "rows": content})
+            text = "\n".join(df.fillna("").astype(str).agg(" ".join, axis=1))
         else:
-            result = json.dumps({"error": f"Unsupported file type: {file_type}"})
-        with open(cache_path, "w", encoding="utf-8") as f:
-            f.write(result)
-        return result
-    except Exception as e:
-        return json.dumps({"error": f"Error processing {os.path.basename(file_path)}: {str(e)}"})
+            text = f"Unsupported file type: {file_type}"
 
-def log_system_usage(tag=""):
-    try:
-        cpu = psutil.cpu_percent(interval=1)
-        mem = psutil.virtual_memory()
-        print(f"[{tag}] CPU: {cpu}% | RAM: {mem.used // (1024**2)}MB / {mem.total // (1024**2)}MB")
-        result = subprocess.run(
-            ["nvidia-smi", "--query-gpu=memory.used,memory.total,utilization.gpu", "--format=csv,nounits,noheader"],
-            capture_output=True, text=True
-        )
-        if result.returncode == 0:
-            used, total, util = result.stdout.strip().split(", ")
-            print(f"[{tag}] GPU: {used}MB / {total}MB | Utilization: {util}%")
+        with open(cache_path, "w", encoding="utf-8") as f:
+            f.write(text)
+        return text
     except Exception as e:
-        print(f"[{tag}] GPU/CPU monitor failed: {e}")
-
-def clean_response(text: str) -> str:
-    text = sanitize_utf8(text)
-    text = re.sub(r"\[TOOL_CALLS\].*", "", text, flags=re.DOTALL)
-    text = re.sub(r"\n{3,}", "\n\n", text).strip()
-    return text
-
-def init_agent():
-    print("🔁 Initializing model...")
-    log_system_usage("Before Load")
-    default_tool_path = os.path.abspath("data/new_tool.json")
-    target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
-    if not os.path.exists(target_tool_path):
-        shutil.copy(default_tool_path, target_tool_path)
-
-    agent = TxAgent(
-        model_name="mims-harvard/TxAgent-T1-Llama-3.1-8B",
-        rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
-        tool_files_dict={"new_tool": target_tool_path},
-        force_finish=True,
-        enable_checker=True,
-        step_rag_num=4,
-        seed=100,
-        additional_default_tools=[],
-    )
-    agent.init_model()
-    log_system_usage("After Load")
-    print("✅ Agent Ready")
-    return agent
-
-def create_ui(agent):
+        return f"Error processing {os.path.basename(file_path)}: {str(e)}"
+
+def parse_analysis_response(raw_response: str) -> Dict[str, List[str]]:
+    """Parse raw analysis response into structured sections."""
+    sections = {
+        "Missed Diagnoses": [],
+        "Medication Conflicts": [],
+        "Incomplete Assessments": [],
+        "Urgent Follow-up": []
+    }
+    current_section = None
+    lines = raw_response.split("\n")
+
+    for line in lines:
+        line = line.strip()
+        if not line:
+            continue
+        if line.startswith("Missed Diagnoses"):
+            current_section = "Missed Diagnoses"
+        elif line.startswith("Medication Conflicts"):
+            current_section = "Medication Conflicts"
+        elif line.startswith("Incomplete Assessments"):
+            current_section = "Incomplete Assessments"
+        elif line.startswith("Urgent Follow-up"):
+            current_section = "Urgent Follow-up"
+        elif current_section and line.startswith("-"):
+            sections[current_section].append(line)
+    
+    return sections
+
+def analyze_medical_records(extracted_text: str) -> str:
+    """Analyze medical records for clinical oversights and return structured response."""
+    # Placeholder for dynamic analysis (replace with actual model or rule-based logic)
+    # Example response to demonstrate flexibility with varying content
+    raw_response = """
+    Missed Diagnoses:
+    - Undiagnosed hypertension despite elevated BP readings.
+    - Family history of diabetes not evaluated for prediabetes risk.
+    
+    Medication Conflicts:
+    - Concurrent use of SSRIs and NSAIDs detected, increasing risk of gastrointestinal bleeding.
+    - Beta-blocker prescribed without assessing asthma history, risking bronchospasm.
+    
+    Incomplete Assessments:
+    - No cardiac stress test despite reported chest pain.
+    - Social history lacks documentation of substance use or living conditions.
+    
+    Urgent Follow-up:
+    - Abnormal ECG results require immediate cardiology referral.
+    - Elevated liver enzymes not addressed, needing hepatology consultation.
+    """
+    
+    # Parse the raw response into sections
+    parsed = parse_analysis_response(raw_response)
+    
+    # Format the response
+    response = ["### Clinical Oversight Analysis\n"]
+    has_findings = False
+    for section, items in parsed.items():
+        response.append(f"#### {section}")
+        if items:
+            response.extend(items)
+            has_findings = True
+        else:
+            response.append("- None identified.")
+        response.append("")  # Add newline for readability
+    
+    response.append("### Summary")
+    if has_findings:
+        summary = "The analysis identified potential oversights in diagnosis, medication management, assessments, and follow-up needs. Immediate action is recommended to address critical findings and ensure comprehensive patient care."
+    else:
+        summary = "No significant clinical oversights were identified in the provided records. Continue monitoring and ensure complete documentation."
+    response.append(summary)
+    
+    return "\n".join(response)
+
+def create_ui():
+    """Create Gradio UI for clinical oversight analysis."""
     with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Markdown("<h1 style='text-align: center;'>🩺 Clinical Oversight Assistant</h1>")
         chatbot = gr.Chatbot(label="Analysis", height=600, type="messages")
@@ -144,106 +158,32 @@ def create_ui(agent):
         download_output = gr.File(label="Download Full Report")
 
         def analyze(message: str, history: List[dict], files: List):
+            """Handle analysis of medical records and update UI."""
             history.append({"role": "user", "content": message})
             history.append({"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."})
             yield history, None
 
-            extracted = ""
+            extracted_text = ""
             file_hash_value = ""
             if files:
                 with ThreadPoolExecutor(max_workers=6) as executor:
-                    futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) for f in files]
-                    results = [sanitize_utf8(f.result()) for f in as_completed(futures)]
-                    extracted = "\n".join(results)
+                    futures = [executor.submit(convert_file_to_text, f.name, f.name.split(".")[-1].lower()) for f in files]
+                    extracted_text = "\n".join(sanitize_utf8(f.result()) for f in as_completed(futures))
                     file_hash_value = file_hash(files[0].name) if files else ""
 
-            # Split extracted text into chunks of ~6,000 characters
-            chunk_size = 6000
-            chunks = [extracted[i:i + chunk_size] for i in range(0, len(extracted), chunk_size)]
-            combined_response = ""
-
-            prompt_template = f"""
-Analyze the medical records for clinical oversights. Provide a concise, evidence-based summary under these headings:
-1. **Missed Diagnoses**:
-   - Identify inconsistencies in history, symptoms, or tests.
-   - Consider psychiatric, neurological, infectious, autoimmune, genetic conditions, family history, trauma, and developmental factors.
-2. **Medication Conflicts**:
-   - Check for contraindications, interactions, or unjustified off-label use.
-   - Assess if medications worsen diagnoses or cause adverse effects.
-3. **Incomplete Assessments**:
-   - Note missing or superficial cognitive, psychiatric, social, or family assessments.
-   - Highlight gaps in medical history, substance use, or lab/imaging documentation.
-4. **Urgent Follow-up**:
-   - Flag abnormal lab results, imaging, behaviors, or legal history needing immediate reassessment or referral.
-Medical Records (Chunk {0} of {1}):
-{{chunk}}
-Begin analysis:
-"""
-
+            # Analyze extracted text
+            history.pop()  # Remove "Analyzing..." message
             try:
-                if history and history[-1]["content"].startswith("⏳"):
-                    history.pop()
-
-                # Process each chunk and stream results in real-time
-                for chunk_idx, chunk in enumerate(chunks, 1):
-                    # Update UI with progress
-                    history.append({"role": "assistant", "content": f"🔄 Processing Chunk {chunk_idx} of {len(chunks)}..."})
-                    yield history, None
-
-                    prompt = prompt_template.format(chunk_idx, len(chunks), chunk=chunk)
-                    chunk_response = ""
-                    for chunk_output in agent.run_gradio_chat(
-                        message=prompt,
-                        history=[],
-                        temperature=0.2,
-                        max_new_tokens=1024,
-                        max_token=4096,
-                        call_agent=False,
-                        conversation=[],
-                    ):
-                        if chunk_output is None:
-                            continue
-                        if isinstance(chunk_output, list):
-                            for m in chunk_output:
-                                if hasattr(m, 'content') and m.content:
-                                    cleaned = clean_response(m.content)
-                                    if cleaned:
-                                        chunk_response += cleaned + "\n"
-                                        # Update UI with partial response
-                                        if history[-1]["content"].startswith("🔄"):
-                                            history[-1] = {"role": "assistant", "content": f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"}
-                                        else:
-                                            history[-1]["content"] = f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"
-                                        yield history, None
-                        elif isinstance(chunk_output, str) and chunk_output.strip():
-                            cleaned = clean_response(chunk_output)
-                            if cleaned:
-                                chunk_response += cleaned + "\n"
-                                # Update UI with partial response
-                                if history[-1]["content"].startswith("🔄"):
-                                    history[-1] = {"role": "assistant", "content": f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"}
-                                else:
-                                    history[-1]["content"] = f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response.strip()}"
-                                yield history, None
-
-                    # Append completed chunk response to combined response
-                    combined_response += f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response}\n"
-
-                # Finalize UI with complete response
-                if combined_response:
-                    history[-1]["content"] = combined_response.strip()
-                else:
-                    history.append({"role": "assistant", "content": "No oversights identified."})
+                response = analyze_medical_records(extracted_text)
+                history.append({"role": "assistant", "content": response})
 
                 # Generate report file
                 report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
                 if report_path:
                     with open(report_path, "w", encoding="utf-8") as f:
-                        f.write(combined_response)
+                        f.write(response)
                 yield history, report_path if report_path and os.path.exists(report_path) else None
-
             except Exception as e:
-                print("🚨 ERROR:", e)
                 history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
                 yield history, None
 
@@ -253,12 +193,14 @@ Begin analysis:
 
 if __name__ == "__main__":
     print("🚀 Launching app...")
-    agent = init_agent()
-    demo = create_ui(agent)
-    demo.queue(api_open=False).launch(
-        server_name="0.0.0.0",
-        server_port=7860,
-        show_error=True,
-        allowed_paths=[report_dir],
-        share=False
-    )
+    try:
+        demo = create_ui()
+        demo.queue(api_open=False).launch(
+            server_name="0.0.0.0",
+            server_port=7860,
+            show_error=True,
+            allowed_paths=[report_dir],
+            share=False
+        )
+    except Exception as e:
+        print(f"Failed to launch app: {str(e)}")