Update app.py

app.py

@@ -1,32 +1,49 @@
+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 shutil
+import re
+import psutil
+import subprocess
 import multiprocessing
 from functools import partial
-import logging
+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
+
 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()
 
@@ -37,12 +54,12 @@ def extract_page_range(file_path: str, start_page: int, end_page: int) -> str:
         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)
+                text_chunks.append(f"=== Page {start_page + pdf.pages.index(page) + 1} ===\n{page_text.strip()}")
+        return "\n\n".join(text_chunks)
     except Exception:
         return ""
 
-def extract_all_pages(file_path: str) -> str:
+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:
@@ -51,8 +68,8 @@ def extract_all_pages(file_path: str) -> str:
         if total_pages == 0:
             return ""
         
-        # Use 4 processes (adjust based on CPU cores)
-        num_processes = min(4, multiprocessing.cpu_count())
+        # Use 6 processes (adjust based on CPU cores)
+        num_processes = min(6, multiprocessing.cpu_count())
         pages_per_process = max(1, total_pages // num_processes)
         
         # Create page ranges for parallel processing
@@ -64,165 +81,263 @@ def extract_all_pages(file_path: str) -> str:
         # Process page ranges in parallel
         with multiprocessing.Pool(processes=num_processes) as pool:
             extract_func = partial(extract_page_range, file_path)
-            results = pool.starmap(extract_func, ranges)
+            results = []
+            for idx, result in enumerate(pool.starmap(extract_func, ranges)):
+                results.append(result)
+                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(filter(None, results))
+    except Exception as e:
+        return f"PDF processing error: {str(e)}"
 
-def convert_file_to_text(file_path: str, file_type: str) -> 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)
+            text = extract_all_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 ""
+            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 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)
+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}")
 
-    for line in raw_response.splitlines():
+def clean_response(text: str) -> str:
+    """Clean TxAgent response to keep only markdown sections with valid findings."""
+    text = sanitize_utf8(text)
+    # Remove tool call artifacts, None, and reasoning
+    text = re.sub(r"\[.*?\]|\bNone\b|To analyze the patient record excerpt.*?medications\.|Since the previous attempts.*?\.|I need to.*?medications\.|Retrieving tools.*?\.", "", text, flags=re.DOTALL)
+    # Remove extra whitespace and non-markdown content
+    text = re.sub(r"\n{3,}", "\n\n", text)
+    text = re.sub(r"[^\n#\-\*\w\s\.\,\:\(\)]+", "", text)  # Keep markdown-relevant characters
+    
+    # Extract markdown sections with valid findings
+    sections = []
+    current_section = None
+    lines = text.splitlines()
+    for line in lines:
         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) -> str:
-    """Analyze medical records and return 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)]
-    
-    # Placeholder for analysis (replace with model or rule-based logic)
-    raw_response_template = """
-    Missed Diagnoses:
-    - Undiagnosed hypertension despite elevated BP readings.
-    - Family history of diabetes not evaluated for prediabetes risk.
-    
-    Medication Conflicts:
-    - SSRIs and NSAIDs detected, increasing GI bleeding risk.
-    
-    Incomplete Assessments:
-    - No cardiac stress test despite chest pain.
-    
-    Urgent Follow-up:
-    - Abnormal ECG requires cardiology referral.
-    """
-    
-    # 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 re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line):
+            current_section = line
+            sections.append([current_section])
+        elif current_section and re.match(r"-\s*.+", line) and not re.match(r"-\s*No issues identified", line):
+            sections[-1].append(line)
     
-    # Format final response
-    response = ["### Clinical Oversight Analysis\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("- None identified.")
-        response.append("")
-    
-    response.append("### Summary")
-    summary = ("The analysis identified potential oversights in diagnosis, medication management, "
-               "assessments, and follow-up needs. Immediate action is recommended.") if has_findings else \
-              "No significant oversights identified. Continue monitoring."
-    response.append(summary)
+    # Combine only non-empty sections
+    cleaned = []
+    for section in sections:
+        if len(section) > 1:  # Section has at least one finding
+            cleaned.append("\n".join(section))
     
-    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 and return results."""
-        history.append({"role": "user", "content": message})
-        history.append({"role": "assistant", "content": "⏳ Extracting text from files..."})
-        yield history, None
-
-        extracted_text = ""
-        file_hash_value = ""
-        if files:
-            with ThreadPoolExecutor(max_workers=4) as executor:
-                futures = [executor.submit(convert_file_to_text, f.name, f.name.split(".")[-1].lower()) 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.pop()  # Remove "Extracting..."
-        history.append({"role": "assistant", "content": "⏳ Analyzing medical records..."})
-        yield history, None
-
-        report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt") if file_hash_value else None
-
-        try:
-            response = analyze_medical_records(extracted_text)
-            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
+    text = "\n\n".join(cleaned).strip()
+    if not text:
+        text = ""  # Return empty string if no valid findings
+    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 and 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).
+Output ONLY the markdown-formatted findings, with bullet points under each heading. Do NOT include reasoning, tool calls, or intermediate steps. If no issues are found in a section, state "No issues identified." Ensure the output is specific to the provided text and avoids generic responses.
+
+Example Output:
+### Missed Diagnoses
+- Elevated BP noted without diagnosis. Missed due to inconsistent visits. Risks: stroke. Recommend: BP monitoring, antihypertensives.
+### Medication Conflicts
+- No issues identified.
+### Incomplete Assessments
+- Chest pain not evaluated. Time constraints likely cause. Risks: cardiac issues. Recommend: ECG, stress test.
+### Urgent Follow-up
+- Abnormal creatinine not addressed. Delayed lab review. Risks: renal failure. Recommend: nephrology referral.
+
+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 and re.search(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", cleaned):
+                                        chunk_response += cleaned + "\n\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 and re.search(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", cleaned):
+                                chunk_response += cleaned + "\n\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
+                    if chunk_response:
+                        combined_response += f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response}\n"
+                    else:
+                        combined_response += f"--- Analysis for Chunk {chunk_idx} ---\nNo oversights identified for this chunk.\n\n"
+
+                # Finalize UI with complete response
+                if combined_response.strip() and not all("No oversights identified" in chunk for chunk in combined_response.split("--- Analysis for Chunk")):
+                    history[-1]["content"] = combined_response.strip()
+                else:
+                    history.append({"role": "assistant", "content": "No oversights identified in the provided records."})
+
+                # 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])
@@ -230,14 +345,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
+    )