Update app.py

app.py

@@ -1,261 +1,275 @@
+import sys
 import os
 import pandas as pd
 import pdfplumber
-import re
+import json
 import gradio as gr
-from typing import List, Dict
-from concurrent.futures import ThreadPoolExecutor
+from typing import List
+from concurrent.futures import ThreadPoolExecutor, as_completed
 import hashlib
-import multiprocessing
-from functools import partial
-import logging
+import shutil
+import re
+import psutil
+import subprocess
 import time
 
-# Suppress pdfplumber CropBox warnings
-logging.getLogger("pdfplumber").setLevel(logging.ERROR)
-
-# Persistent directories
+# Persistent directory
 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")
-for directory in [file_cache_dir, report_dir]:
+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]:
     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'}
+
 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_page_range(file_path: str, start_page: int, end_page: int) -> str:
-    """Extract text from a range of PDF pages."""
+def extract_priority_pages(file_path: str, progress_callback=None) -> str:
     try:
         text_chunks = []
-        with pdfplumber.open(file_path) as pdf:
-            for page in pdf.pages[start_page:end_page]:
-                page_text = page.extract_text() or ""
-                text_chunks.append(page_text.strip())
-        return "\n".join(text_chunks)
-    except Exception:
-        return ""
-
-def extract_all_pages(file_path: str, progress_callback=None) -> str:
-    """Extract text from all pages of a PDF using parallel processing."""
-    try:
         with pdfplumber.open(file_path) as pdf:
             total_pages = len(pdf.pages)
-        
-        if total_pages == 0:
-            return ""
-        
-        # Use 4 processes (adjust based on CPU cores)
-        num_processes = min(4, multiprocessing.cpu_count())
-        pages_per_process = max(1, total_pages // num_processes)
-        
-        # Create page ranges for parallel processing
-        ranges = [(i * pages_per_process, min((i + 1) * pages_per_process, total_pages))
-                  for i in range(num_processes)]
-        if ranges[-1][1] != total_pages:
-            ranges[-1] = (ranges[-1][0], total_pages)
-        
-        # Process page ranges in parallel
-        with multiprocessing.Pool(processes=num_processes) as pool:
-            extract_func = partial(extract_page_range, file_path)
-            results = []
-            for idx, result in enumerate(pool.starmap(extract_func, ranges)):
-                results.append(result)
+            processed_pages = 0
+            for i, page in enumerate(pdf.pages):
+                page_text = page.extract_text() or ""
+                if i < 3 or any(re.search(rf'\b{kw}\b', page_text.lower()) for kw in MEDICAL_KEYWORDS):
+                    text_chunks.append(f"=== Page {i+1} ===\n{page_text.strip()}")
+                processed_pages += 1
                 if progress_callback:
-                    processed_pages = min((idx + 1) * pages_per_process, total_pages)
                     progress_callback(processed_pages, total_pages)
-        
-        return "\n".join(filter(None, results))
-    except Exception:
-        return ""
+        return "\n\n".join(text_chunks)
+    except Exception as e:
+        return f"PDF processing error: {str(e)}"
 
-def convert_file_to_text(file_path: str, file_type: str, progress_callback=None) -> str:
-    """Convert supported file types to text, caching results."""
+def convert_file_to_json(file_path: str, file_type: str, progress_callback=None) -> str:
     try:
         h = file_hash(file_path)
-        cache_path = os.path.join(file_cache_dir, f"{h}.txt")
+        cache_path = os.path.join(file_cache_dir, f"{h}.json")
         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_all_pages(file_path, progress_callback)
+            text = extract_priority_pages(file_path, progress_callback)
+            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=True, on_bad_lines="skip")
-            text = " ".join(df.fillna("").astype(str).agg(" ".join, axis=1))
+                             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})
         elif file_type in ["xls", "xlsx"]:
-            df = pd.read_excel(file_path, engine="openpyxl", header=None, dtype=str)
-            text = " ".join(df.fillna("").astype(str).agg(" ".join, axis=1))
+            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})
         else:
-            text = ""
-
-        if text:
-            # Compress text by removing redundant whitespace
-            text = re.sub(r'\s+', ' ', text).strip()
-            with open(cache_path, "w", encoding="utf-8") as f:
-                f.write(text)
-        return text
-    except Exception:
-        return ""
-
-def parse_analysis_response(raw_response: str) -> Dict[str, List[str]]:
-    """Parse raw analysis response into structured sections using regex."""
-    sections = {
-        "Missed Diagnoses": [],
-        "Medication Conflicts": [],
-        "Incomplete Assessments": [],
-        "Urgent Follow-up": []
-    }
-    current_section = None
-    section_pattern = re.compile(r"^(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up):$", re.MULTILINE)
-    item_pattern = re.compile(r"^- .+$", re.MULTILINE)
-
-    for line in raw_response.splitlines():
-        line = line.strip()
-        if not line:
-            continue
-        if section_pattern.match(line):
-            current_section = line[:-1]
-        elif current_section and item_pattern.match(line):
-            sections[current_section].append(line)
-    
-    return sections
-
-def analyze_medical_records(extracted_text: str, progress_callback=None) -> str:
-    """Analyze medical records and return generalized structured response."""
-    # Split text into chunks to handle large inputs
-    chunk_size = 10000
-    chunks = [extracted_text[i:i + chunk_size] for i in range(0, len(extracted_text), chunk_size)]
-    
-    # Generalized analysis template (replace with model or rule-based logic)
-    raw_response_template = """
-    Missed Diagnoses:
-    - Chronic conditions potentially missed due to inconsistent monitoring of vital signs or symptoms. This may occur when patient visits are infrequent or records lack longitudinal tracking, leading to undetected trends. Undiagnosed conditions can progress, increasing risks of complications like organ damage. Recommended action: Implement regular screening protocols and trend analysis for key indicators (e.g., blood pressure, glucose levels).
-    - Risk factors for hereditary or lifestyle-related diseases not screened despite documented family history or patient demographics. Screening oversights often stem from time constraints or lack of standardized protocols. Delayed diagnosis may lead to preventable disease progression. Recommended action: Establish routine risk assessments based on family history and clinical guidelines.
-    
-    Medication Conflicts:
-    - Potential interactions from polypharmacy or untracked over-the-counter medications. Conflicts may arise when multiple prescribers are involved or patients self-medicate, increasing risks of adverse events like bleeding or toxicity. Recommended action: Conduct comprehensive medication reconciliation at each visit and educate patients on reporting all medications.
-    
-    Incomplete Assessments:
-    - Symptoms reported but not fully evaluated due to incomplete documentation or failure to follow clinical guidelines. This can occur in busy clinical settings where time limits prioritize acute issues over thorough investigation. Unaddressed symptoms may mask serious conditions, delaying treatment. Recommended action: Standardize symptom evaluation protocols and ensure adequate time for comprehensive assessments.
-    
-    Urgent Follow-up:
-    - Critical findings requiring specialist referral or additional testing delayed due to communication gaps or scheduling issues. Delays often result from fragmented care coordination or underestimation of findings' severity. Untreated critical issues can lead to rapid deterioration. Recommended action: Establish clear referral pathways and prioritize urgent findings with defined timelines.
-    """
-    
-    # Aggregate findings across chunks
-    all_sections = {
-        "Missed Diagnoses": set(),
-        "Medication Conflicts": set(),
-        "Incomplete Assessments": set(),
-        "Urgent Follow-up": set()
-    }
-    
-    for chunk_idx, chunk in enumerate(chunks, 1):
-        # Simulate analysis per chunk (replace with real logic)
-        raw_response = raw_response_template
-        parsed = parse_analysis_response(raw_response)
-        for section, items in parsed.items():
-            all_sections[section].update(items)
-        if progress_callback:
-            progress_callback(chunk_idx, len(chunks))
-    
-    # Format generalized response
-    response = ["### Clinical Oversight Analysis\n"]
-    response.append("This analysis reviews patient records to identify common reasons for potential oversights that could impact clinical outcomes. Findings highlight systemic or procedural gaps, associated risks, and actionable recommendations applicable across various patient records.\n")
-    has_findings = False
-    for section, items in all_sections.items():
-        response.append(f"#### {section}")
-        if items:
-            response.extend(sorted(items))
-            has_findings = True
-        else:
-            response.append("- No issues identified in this category.")
-        response.append("")
-    
-    response.append("### Summary")
-    if has_findings:
-        summary = ("The analysis identified common procedural and systemic gaps that may lead to oversights in diagnosis, medication management, assessments, and follow-up care. These gaps, such as inconsistent monitoring, incomplete documentation, or communication delays, pose risks of disease progression, adverse events, or delayed treatment. Recommended actions include standardizing screening and assessment protocols, improving medication reconciliation, and establishing clear referral pathways. Implementing these measures can enhance patient safety and care quality across diverse clinical scenarios.")
-    else:
-        summary = ("No significant oversights were identified in the provided records. Current practices appear aligned with general clinical standards. To maintain care quality, continue regular monitoring, ensure comprehensive documentation, and adhere to guideline-based screening and follow-up protocols.")
-    response.append(summary)
-    
-    return "\n".join(response)
-
-def create_ui():
-    """Create Gradio UI for clinical oversight analysis."""
-    def analyze(message: str, history: List[dict], files: List):
-        """Handle analysis with animated progress updates."""
-        history.append({"role": "user", "content": message})
-        yield history, None
-
-        extracted_text = ""
-        file_hash_value = ""
-        if files:
-            # Progress callback for extraction
-            total_pages = 0
-            processed_pages = 0
-            def update_extraction_progress(current, total):
-                nonlocal processed_pages, total_pages
-                processed_pages = current
-                total_pages = total
-                animation = ["🌀", "🔄", "⚙️", "🔃"][(int(time.time() * 2) % 4)]
-                history[-1] = {"role": "assistant", "content": f"Extracting text... {animation} Page {processed_pages}/{total_pages}"}
-                return history, None
-
-            with ThreadPoolExecutor(max_workers=4) as executor:
-                futures = [executor.submit(convert_file_to_text, f.name, f.name.split(".")[-1].lower(), update_extraction_progress) for f in files]
-                results = [f.result() for f in futures]
-                extracted_text = "\n".join(sanitize_utf8(r) for r in results if r)
-                file_hash_value = file_hash(files[0].name) if files else ""
-
-        history.append({"role": "assistant", "content": "✅ Text extraction complete."})
-        yield history, None
-
-        report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
-
-        try:
-            # Progress callback for analysis
-            total_chunks = 0
-            processed_chunks = 0
-            def update_analysis_progress(current, total):
-                nonlocal processed_chunks, total_chunks
-                processed_chunks = current
-                total_chunks = total
-                animation = ["🔍", "📊", "🧠", "🔎"][(int(time.time() * 2) % 4)]
-                history[-1] = {"role": "assistant", "content": f"Analyzing records... {animation} Chunk {processed_chunks}/{total_chunks}"}
-                return history, None
-
-            history.append({"role": "assistant", "content": "Analyzing records... 🔍"})
-            yield history, None
-            response = analyze_medical_records(extracted_text, update_analysis_progress)
-            
-            history.pop()  # Remove "Analyzing..."
-            history.append({"role": "assistant", "content": response})
-            if report_path:
-                with open(report_path, "w", encoding="utf-8") as f:
-                    f.write(response)
-            yield history, report_path if report_path and os.path.exists(report_path) else None
-        except Exception as e:
-            history.pop()  # Remove "Analyzing..."
-            history.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
-            yield history, None
+            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)}"})
+
+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}%")
+    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):
     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")
         file_upload = gr.File(file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple")
         msg_input = gr.Textbox(placeholder="Ask about potential oversights...", show_label=False)
         send_btn = gr.Button("Analyze", variant="primary")
-        download_output = gr.File(label="Download Report")
+        download_output = gr.File(label="Download Full Report")
+
+        def analyze(message: str, history: List[dict], files: List):
+            history.append({"role": "user", "content": message})
+            history.append({"role": "assistant", "content": "⏳ Extracting text from files..."})
+            yield history, None
+
+            extracted = ""
+            file_hash_value = ""
+            if files:
+                # Progress callback for extraction
+                total_pages = 0
+                processed_pages = 0
+                def update_extraction_progress(current, total):
+                    nonlocal processed_pages, total_pages
+                    processed_pages = current
+                    total_pages = total
+                    animation = ["🌀", "🔄", "⚙️", "🔃"][(int(time.time() * 2) % 4)]
+                    history[-1] = {"role": "assistant", "content": f"Extracting text... {animation} Page {processed_pages}/{total_pages}"}
+                    return history, None
+
+                with ThreadPoolExecutor(max_workers=6) as executor:
+                    futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower(), update_extraction_progress) for f in files]
+                    results = [sanitize_utf8(f.result()) for f in as_completed(futures)]
+                    extracted = "\n".join(results)
+                    file_hash_value = file_hash(files[0].name) if files else ""
+
+            history.pop()  # Remove extraction message
+            history.append({"role": "assistant", "content": "✅ Text extraction complete."})
+            yield history, None
+
+            # 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 = """
+You are a medical analysis assistant. Analyze the following patient record excerpt for clinical oversights. Provide a concise, evidence-based summary in markdown format under these headings: Missed Diagnoses, Medication Conflicts, Incomplete Assessments, and Urgent Follow-up. For each finding, include:
+- Clinical context (why the issue was missed or relevant details from the record).
+- Potential risks if unaddressed (e.g., disease progression, adverse events).
+- Actionable recommendations (e.g., tests, referrals, medication adjustments).
+If no issues are found in a section, state "No issues identified." Ensure the output is specific to the provided text, formatted as markdown with bullet points under each heading, and avoids generic or static responses.
+
+Patient Record Excerpt (Chunk {0} of {1}):
+{chunk}
+
+### Missed Diagnoses
+- ...
+
+### Medication Conflicts
+- ...
+
+### Incomplete Assessments
+- ...
+
+### Urgent Follow-up
+- ...
+"""
+
+            try:
+                # Process each chunk and stream results in real-time
+                for chunk_idx, chunk in enumerate(chunks, 1):
+                    # Update UI with chunk progress
+                    animation = ["🔍", "📊", "🧠", "🔎"][(int(time.time() * 2) % 4)]
+                    history.append({"role": "assistant", "content": f"Analyzing records... {animation} Chunk {chunk_idx}/{len(chunks)}"})
+                    yield history, None
+
+                    prompt = prompt_template.format(chunk_idx, len(chunks), chunk=chunk[:4000])  # Truncate to avoid token limits
+                    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("Analyzing"):
+                                            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("Analyzing"):
+                                    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."})
+
+                # 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)
+                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
 
         send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])
         msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])
@@ -263,14 +277,12 @@ def create_ui():
 
 if __name__ == "__main__":
     print("🚀 Launching app...")
-    try:
-        demo = create_ui()
-        demo.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)}")
+    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
+    )