Update app.py

app.py

@@ -11,15 +11,6 @@ import shutil
 import re
 import psutil
 import subprocess
-import logging
-import torch
-import gc
-from diskcache import Cache
-import time
-
-# Configure logging with a more specific logger name to avoid conflicts
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger("ClinicalOversightApp")
 
 # Persistent directory
 persistent_dir = "/data/hf_cache"
@@ -46,8 +37,8 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
-# Initialize cache with 10GB limit
-cache = Cache(file_cache_dir, size_limit=10 * 1024**3)
+MEDICAL_KEYWORDS = {'diagnosis', 'assessment', 'plan', 'results', 'medications',
+                    'allergies', 'summary', 'impression', 'findings', 'recommendations'}
 
 def sanitize_utf8(text: str) -> str:
     return text.encode("utf-8", "ignore").decode("utf-8")
@@ -56,60 +47,28 @@ def file_hash(path: str) -> str:
     with open(path, "rb") as f:
         return hashlib.md5(f.read()).hexdigest()
 
-def extract_all_pages(file_path: str, progress_callback=None) -> str:
+def extract_priority_pages(file_path: str) -> str:
     try:
+        text_chunks = []
         with pdfplumber.open(file_path) as pdf:
-            total_pages = len(pdf.pages)
-            if total_pages == 0:
-                logger.error("No pages found in PDF: %s", file_path)
-                return ""
-
-        batch_size = 10
-        batches = [(i, min(i + batch_size, total_pages)) for i in range(0, total_pages, batch_size)]
-        text_chunks = [""] * total_pages
-        processed_pages = 0
-
-        def extract_batch(start: int, end: int) -> List[tuple]:
-            results = []
-            with pdfplumber.open(file_path) as pdf:
-                for idx, page in enumerate(pdf.pages[start:end], start=start):
-                    page_text = page.extract_text() or ""
-                    results.append((idx, f"=== Page {idx + 1} ===\n{page_text.strip()}"))
-                    logger.debug("Extracted page %d, text length: %d chars", idx + 1, len(page_text))
-            return results
-
-        with ThreadPoolExecutor(max_workers=6) as executor:
-            futures = [executor.submit(extract_batch, start, end) for start, end in batches]
-            for future in as_completed(futures):
-                for page_num, text in future.result():
-                    if page_num < len(text_chunks):
-                        text_chunks[page_num] = text
-                    else:
-                        logger.warning("Page number %d out of range for text_chunks (size %d)", page_num, len(text_chunks))
-                processed_pages += batch_size
-                if progress_callback:
-                    progress_callback(min(processed_pages, total_pages), total_pages)
-                logger.info("Processed %d/%d pages for %s", min(processed_pages, total_pages), total_pages, file_path)
-
-        extracted_text = "\n\n".join(filter(None, text_chunks))
-        logger.info("Extracted %d pages from %s, total length: %d chars", total_pages, file_path, len(extracted_text))
-        return extracted_text
+            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()}")
+        return "\n\n".join(text_chunks)
     except Exception as e:
-        logger.error("PDF processing error for %s: %s", file_path, e, exc_info=True)
         return f"PDF processing error: {str(e)}"
 
-def convert_file_to_json(file_path: str, file_type: str, progress_callback=None) -> str:
+def convert_file_to_json(file_path: str, file_type: str) -> str:
     try:
-        file_h = file_hash(file_path)
-        cache_key = f"{file_h}_{file_type}"
-        if cache_key in cache:
-            logger.info("Cache hit for %s (key: %s)", file_path, cache_key)
-            return cache[cache_key]
-        else:
-            logger.info("Cache miss for %s (key: %s), performing fresh extraction", file_path, cache_key)
+        h = file_hash(file_path)
+        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)
             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,
@@ -125,99 +84,42 @@ def convert_file_to_json(file_path: str, file_type: str, progress_callback=None)
             result = json.dumps({"filename": os.path.basename(file_path), "rows": content})
         else:
             result = json.dumps({"error": f"Unsupported file type: {file_type}"})
-
-        cache[cache_key] = result
-        logger.info("Cached extraction for %s, size: %d bytes", file_path, len(result))
+        with open(cache_path, "w", encoding="utf-8") as f:
+            f.write(result)
         return result
     except Exception as e:
-        logger.error("Error processing %s: %s", os.path.basename(file_path), e, exc_info=True)
         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()
-        logger.info("[%s] CPU: %.1f%% | RAM: %dMB / %dMB", tag, cpu, mem.used // (1024**2), mem.total // (1024**2))
+        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(", ")
-            logger.info("[%s] GPU: %sMB / %sMB | Utilization: %s%%", tag, used, total, util)
+            print(f"[{tag}] GPU: {used}MB / {total}MB | Utilization: {util}%")
     except Exception as e:
-        logger.error("[%s] GPU/CPU monitor failed: %s", tag, e)
+        print(f"[{tag}] GPU/CPU monitor failed: {e}")
 
 def clean_response(text: str) -> str:
     text = sanitize_utf8(text)
-    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)
-    text = re.sub(r"\n{3,}", "\n\n", text)
-    text = re.sub(r"[^\n#\-\*\w\s\.\,\:\(\)]+", "", text)
-
-    sections = {}
-    current_section = None
-    lines = text.splitlines()
-    for line in lines:
-        line = line.strip()
-        if not line:
-            continue
-        section_match = re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line)
-        if section_match:
-            current_section = section_match.group(1)
-            if current_section not in sections:
-                sections[current_section] = []
-            continue
-        finding_match = re.match(r"-\s*.+", line)
-        if finding_match and current_section and not re.match(r"-\s*No issues identified", line):
-            sections[current_section].append(line)
-
-    cleaned = []
-    for heading, findings in sections.items():
-        if findings:
-            cleaned.append(f"### {heading}\n" + "\n".join(findings))
-
-    text = "\n\n".join(cleaned).strip()
-    logger.debug("Cleaned response length: %d chars", len(text))
-    return text if text else ""
-
-def summarize_findings(combined_response: str) -> str:
-    if not combined_response or all("No oversights identified" in chunk for chunk in combined_response.split("--- Analysis for Chunk")):
-        logger.info("No clinical oversights identified in analysis")
-        return "### Summary of Clinical Oversights\nNo critical oversights identified in the provided records."
-
-    sections = {}
-    lines = combined_response.splitlines()
-    current_section = None
-    for line in lines:
-        line = line.strip()
-        if not line:
-            continue
-        section_match = re.match(r"###\s*(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", line)
-        if section_match:
-            current_section = section_match.group(1)
-            if current_section not in sections:
-                sections[current_section] = []
-            continue
-        finding_match = re.match(r"-\s*(.+)", line)
-        if finding_match and current_section:
-            sections[current_section].append(finding_match.group(1))
-
-    summary_lines = []
-    for heading, findings in sections.items():
-        if findings:
-            summary = f"- **{heading}**: {'; '.join(findings[:2])}. Risks: {heading.lower()} may lead to adverse outcomes. Recommend: urgent review and specialist referral."
-            summary_lines.append(summary)
-
-    if not summary_lines:
-        logger.info("No clinical oversights identified after summarization")
-        return "### Summary of Clinical Oversights\nNo critical oversights identified."
-
-    summary = "### Summary of Clinical Oversights\n" + "\n".join(summary_lines)
-    logger.info("Summarized findings: %s", summary[:100])
-    return summary
+    # Remove tool calls, JSON data, and repetitive phrases
+    text = re.sub(r"\[TOOL_CALLS\].*", "", text, flags=re.DOTALL)
+    text = re.sub(r"\['get_[^\]]+\']\n?", "", text)  # Remove tool names
+    text = re.sub(r"\{'meta':\s*\{.*?\}\s*,\s*'results':\s*\[.*?\]\}\n?", "", text, flags=re.DOTALL)  # Remove JSON
+    text = re.sub(r"To analyze the medical records for clinical oversights.*?begin by reviewing.*?\n", "", text, flags=re.DOTALL)
+    text = re.sub(r"\n{3,}", "\n\n", text).strip()
+    # Only keep text under analysis headings or relevant content
+    if not re.search(r"(Missed Diagnoses|Medication Conflicts|Incomplete Assessments|Urgent Follow-up)", text):
+        return ""
+    return text
 
 def init_agent():
-    logger.info("Initializing model...")
+    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")
@@ -229,150 +131,148 @@ def init_agent():
         rag_model_name="mims-harvard/ToolRAG-T1-GTE-Qwen2-1.5B",
         tool_files_dict={"new_tool": target_tool_path},
         force_finish=True,
-        enable_checker=False,
-        step_rag_num=4,
+        enable_checker=True,
+        step_rag_num=2,
         seed=100,
         additional_default_tools=[],
     )
     agent.init_model()
     log_system_usage("After Load")
-    logger.info("Agent Ready")
+    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="Detailed Analysis", height=600, type="messages")
-        final_summary = gr.Markdown(label="Summary of Clinical Oversights")
+        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 Full Report")
-        progress_bar = gr.Progress()
-
-        prompt_template = """
-Analyze the patient record excerpt for clinical oversights. Provide a concise, evidence-based summary in markdown with findings grouped under headings (e.g., 'Missed Diagnoses'). For each finding, include clinical context, risks, and recommendations. Output only markdown bullet points under headings. If no issues, state "No issues identified".
-Patient Record Excerpt (Chunk {0} of {1}):
-{chunk}
-"""
 
-        def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
+        def analyze(message: str, history: List[dict], files: List):
             history.append({"role": "user", "content": message})
-            yield history, None, ""
-            logger.info("Starting analysis for message: %s", message[:100])
-
-            if not files:
-                logger.error("No files uploaded for analysis")
-                history.append({"role": "assistant", "content": "❌ Please upload a file to analyze."})
-                yield history, None, "### Summary of Clinical Oversights\nNo file uploaded for analysis."
-                return
+            history.append({"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."})
+            yield history, None
 
             extracted = ""
             file_hash_value = ""
-            logger.info("Processing %d uploaded files", len(files))
-            for f in files:
-                logger.info("Processing file: %s", f.name)
-
-            def update_extraction_progress(current, total):
-                progress(current / total, desc=f"Extracting text... Page {current}/{total}")
-                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 ""
-            logger.info("Extraction complete for %d files", len(files))
-            history.append({"role": "assistant", "content": "✅ Text extraction complete."})
-            yield history, None, ""
-
-            logger.info("Extracted text length: %d chars", len(extracted))
-            if len(extracted.strip()) == 0:
-                logger.error("Extracted text is empty")
-                history.append({"role": "assistant", "content": "❌ Extracted text is empty. Please ensure the file contains readable content."})
-                yield history, None, "### Summary of Clinical Oversights\nExtracted text is empty."
-                return
-
-            chunk_size = 6000
+            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)
+                    file_hash_value = file_hash(files[0].name) if files else ""
+
+            # Split extracted text into chunks of ~4,000 characters
+            chunk_size = 4000
             chunks = [extracted[i:i + chunk_size] for i in range(0, len(extracted), chunk_size)]
-            if not chunks:
-                chunks = [""]
-            logger.info("Created %d chunks", len(chunks))
             combined_response = ""
-            batch_size = 2
+
+            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:
+"""
 
             try:
-                for batch_idx in range(0, len(chunks), batch_size):
-                    batch_chunks = chunks[batch_idx:batch_idx + batch_size]
-                    batch_prompts = [prompt_template.format(i + 1, len(chunks), chunk=chunk[:4000]) for i, chunk in enumerate(batch_chunks)]
-                    batch_responses = []
-
-                    progress((batch_idx + 1) / len(chunks), desc=f"Analyzing chunks {batch_idx + 1}-{min(batch_idx + batch_size, len(chunks))}/{len(chunks)}")
-
-                    with ThreadPoolExecutor(max_workers=len(batch_chunks)) as executor:
-                        futures = [executor.submit(agent.run_gradio_chat, prompt, [], 0.2, 512, 2048, False, []) for prompt in batch_prompts]
-                        for future in as_completed(futures):
-                            chunk_response = ""
-                            for chunk_output in future.result():
-                                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*\w+", cleaned):
-                                                chunk_response += cleaned + "\n\n"
-                                elif isinstance(chunk_output, str) and chunk_output.strip():
-                                    cleaned = clean_response(chunk_output)
-                                    if cleaned and re.search(r"###\s*\w+", cleaned):
-                                        chunk_response += cleaned + "\n\n"
-                            batch_responses.append(chunk_response)
-                            torch.cuda.empty_cache()
-                            gc.collect()
-                            logger.debug("Processed chunk response length: %d chars", len(chunk_response))
-
-                    for chunk_idx, chunk_response in enumerate(batch_responses, batch_idx + 1):
-                        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"
-                        history[-1] = {"role": "assistant", "content": combined_response.strip()}
-                        yield history, None, ""
-
-                if combined_response.strip() and not all("No oversights identified" in chunk for chunk in combined_response.split("--- Analysis for Chunk")):
+                if history and history[-1]["content"].startswith("⏳"):
+                    history.pop()
+
+                # Process each chunk and stream cleaned results
+                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"
+                                        # Stream partial response to UI
+                                        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"
+                                # Stream partial response to UI
+                                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
+                    if chunk_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 in the provided records."})
+                    history.append({"role": "assistant", "content": "No oversights identified."})
 
-                summary = summarize_findings(combined_response)
+                # Generate report file with cleaned response
                 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 + "\n\n" + summary)
-                logger.info("Analysis complete, report saved at: %s", report_path if report_path else "None")
-                yield history, report_path if report_path and os.path.exists(report_path) else None, summary
+                        f.write(combined_response)
+                yield history, report_path if report_path and os.path.exists(report_path) else None
 
             except Exception as e:
-                logger.error("Analysis error: %s", e, exc_info=True)
+                print("🚨 ERROR:", e)
                 history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
-                yield history, None, f"### Summary of Clinical Oversights\nError occurred during analysis: {str(e)}"
+                yield history, None
 
-        send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary])
-        msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output, final_summary])
+        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])
     return demo
 
 if __name__ == "__main__":
-    try:
-        logger.info("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
-        )
-    finally:
-        if torch.distributed.is_initialized():
-            torch.distributed.destroy_process_group()
+    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
+    )