Update app.py

app.py

@@ -1,10 +1,11 @@
 import sys
 import os
 import pandas as pd
+import pdfplumber
 import json
 import gradio as gr
 from typing import List
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, as_completed
 import hashlib
 import shutil
 import re
@@ -15,21 +16,9 @@ import torch
 import gc
 from diskcache import Cache
 import time
-import asyncio
-
-# Try importing pypdfium2 and pytesseract, fall back to pdfplumber
-try:
-    import pypdfium2 as pdfium
-    import pytesseract
-    from PIL import Image
-    HAS_PYPDFIUM2 = True
-except ImportError:
-    HAS_PYPDFIUM2 = False
-    import pdfplumber
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
-logging.getLogger("pdfminer").setLevel(logging.ERROR)
 logger = logging.getLogger(__name__)
 
 # Persistent directory
@@ -67,78 +56,37 @@ def file_hash(path: str) -> str:
     with open(path, "rb") as f:
         return hashlib.md5(f.read()).hexdigest()
 
-async def extract_all_pages_async(file_path: str, progress_callback=None, force_ocr=False) -> str:
+def extract_all_pages(file_path: str, progress_callback=None) -> str:
     try:
-        extracted_text = ""
-        total_pages = 0
-        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")
                 return ""
 
-            def extract_page(i):
-                page = pdf.pages[i]
-                text = ""
-                # Adjust table settings for complex layouts
-                table_settings = {
-                    "vertical_strategy": "lines",
-                    "horizontal_strategy": "lines",
-                    "explicit_vertical_lines": [],
-                    "explicit_horizontal_lines": [],
-                    "snap_tolerance": 5,
-                    "join_tolerance": 5,
-                    "edge_min_length": 3,
-                    "min_words_vertical": 3,
-                    "min_words_horizontal": 1,
-                    "intersection_tolerance": 5,
-                }
-                tables = page.extract_tables(table_settings=table_settings)
-                if tables:
-                    for table in tables:
-                        table_text = "\n".join(
-                            " | ".join(str(cell) if cell is not None else "" for cell in row)
-                            for row in table if any(cell is not None for cell in row)
-                        )
-                        text += table_text + "\n\n"
-                    logger.debug("Page %d extracted %d tables, text length: %d chars", i + 1, len(tables), len(text))
-                else:
-                    text = page.extract_text() or ""
-                    logger.debug("Page %d no tables, raw text length: %d chars", i + 1, len(text))
-
-                # Force OCR if text is short or force_ocr is True
-                if (not text.strip() or len(text) < 100 or force_ocr) and HAS_PYPDFIUM2 and 'pytesseract' in sys.modules:
-                    try:
-                        logger.info("Attempting OCR for page %d", i + 1)
-                        pdfium_pdf = pdfium.PdfDocument(file_path)
-                        page_bitmap = pdfium_pdf[i].render(scale=2).to_pil()
-                        ocr_text = pytesseract.image_to_string(page_bitmap, lang="eng")
-                        logger.debug("Page %d OCR text length: %d chars", i + 1, len(ocr_text))
-                        text = ocr_text if ocr_text.strip() else text
-                        pdfium_pdf.close()
-                    except Exception as e:
-                        logger.error("OCR failed for page %d: %s", i + 1, e)
-                return (i, f"=== Page {i + 1} ===\n{text.strip()}")
-
-            with ThreadPoolExecutor(max_workers=4) as executor:
-                futures = [executor.submit(extract_page, i) for i in range(total_pages)]
-                for future in futures:
-                    page_num, text = future.result()
-                    text_chunks.append((page_num, text))
-                    logger.debug("Page %d extracted: %s...", page_num + 1, text[:50])
-                    if progress_callback:
-                        progress_callback(page_num + 1, total_pages)
-
-        text_chunks.sort(key=lambda x: x[0])
-        extracted_text = "\n\n".join(chunk[1] for chunk in text_chunks if chunk[1].strip())
-        logger.info("Extracted %d pages, total length: %d chars", total_pages, len(extracted_text))
-        # Force OCR retry if text is too short
-        if len(extracted_text) < 1000 and not force_ocr and HAS_PYPDFIUM2 and 'pytesseract' in sys.modules:
-            logger.info("Text too short, forcing OCR for all pages")
-            return await extract_all_pages_async(file_path, progress_callback, force_ocr=True)
-        return extracted_text
+        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 page in pdf.pages[start:end]:
+                    page_num = start + pdf.pages.index(page)
+                    page_text = page.extract_text() or ""
+                    results.append((page_num, f"=== Page {page_num + 1} ===\n{page_text.strip()}"))
+            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():
+                    text_chunks[page_num] = text
+                processed_pages += batch_size
+                if progress_callback:
+                    progress_callback(min(processed_pages, total_pages), total_pages)
+
+        return "\n\n".join(filter(None, text_chunks))
     except Exception as e:
         logger.error("PDF processing error: %s", e)
         return f"PDF processing error: {str(e)}"
@@ -147,59 +95,28 @@ def convert_file_to_json(file_path: str, file_type: str, progress_callback=None)
     try:
         file_h = file_hash(file_path)
         cache_key = f"{file_h}_{file_type}"
-        # Bypass cache to force fresh extraction
-        logger.info("Forcing fresh extraction for %s", file_path)
-        # if cache_key in cache:
-        #     logger.info("Using cached extraction for %s", file_path)
-        #     return cache[cache_key]
+        if cache_key in cache:
+            return cache[cache_key]
 
         if file_type == "pdf":
-            text = asyncio.run(extract_all_pages_async(file_path, progress_callback, force_ocr=False))
+            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":
-            try:
-                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})
-                logger.info("CSV processed, rows: %d", len(content))
-            except Exception as e:
-                logger.error("CSV processing failed: %s", e)
-                result = json.dumps({"error": f"CSV processing failed: {str(e)}"})
+            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})
         elif file_type in ["xls", "xlsx"]:
             try:
-                # Try all sheets to maximize data
-                xl = pd.ExcelFile(file_path, engine="openpyxl")
-                content = []
-                for sheet_name in xl.sheet_names:
-                    try:
-                        df = pd.read_excel(file_path, sheet_name=sheet_name, engine="openpyxl", header=None, dtype=str)
-                        sheet_content = df.fillna("").astype(str).values.tolist()
-                        content.extend(sheet_content)
-                        logger.debug("Excel sheet %s processed, rows: %d", sheet_name, len(sheet_content))
-                    except Exception as e:
-                        logger.warning("Excel sheet %s failed: %s", sheet_name, e)
-                if not content:
-                    logger.error("No valid data extracted from Excel")
-                    result = json.dumps({"error": "No valid data extracted from Excel"})
-                else:
-                    result = json.dumps({"filename": os.path.basename(file_path), "rows": content})
-                    logger.info("Excel processed, total rows: %d", len(content))
-            except Exception as e:
-                logger.error("Excel processing failed: %s", e)
-                try:
-                    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})
-                    logger.info("Excel processed with xlrd, rows: %d", len(content))
-                except Exception as e2:
-                    logger.error("Excel processing failed with xlrd: %s", e2)
-                    result = json.dumps({"error": f"Excel processing failed: {str(e)}"})
+                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:
             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))
         return result
     except Exception as e:
         logger.error("Error processing %s: %s", os.path.basename(file_path), e)
@@ -222,63 +139,66 @@ def log_system_usage(tag=""):
 
 def clean_response(text: str) -> str:
     text = sanitize_utf8(text)
-    text = text.replace("[", "").replace("]", "").replace("None", "")
-    text = text.replace("\n\n\n", "\n\n")
+    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
-    seen_lines = set()
-    for line in text.splitlines():
+    lines = text.splitlines()
+    for line in lines:
         line = line.strip()
-        if not line or line in seen_lines:
+        if not line:
             continue
-        seen_lines.add(line)
         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)
-            sections.setdefault(current_section, [])
+            if current_section not in sections:
+                sections[current_section] = []
             continue
-        if current_section and line.startswith("- "):
+        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 = [f"### {heading}\n" + "\n".join(findings) for heading, findings in sections.items() if findings]
-    result = "\n\n".join(cleaned).strip()
-    logger.debug("Cleaned response length: %d chars", len(result))
-    return result or "No oversights identified"
-
-def summarize_findings(all_responses: List[str]) -> str:
-    combined_response = "\n\n".join(all_responses)
-    if not combined_response or all("No oversights identified" in resp.lower() for resp in all_responses):
-        return "### Comprehensive Clinical Oversight Summary\nNo critical oversights were identified across the provided patient records after thorough analysis."
-
-    sections = {
-        "Missed Diagnoses": [],
-        "Medication Conflicts": [],
-        "Incomplete Assessments": [],
-        "Urgent Follow-up": []
-    }
+
+    cleaned = []
+    for heading, findings in sections.items():
+        if findings:
+            cleaned.append(f"### {heading}\n" + "\n".join(findings))
+
+    text = "\n\n".join(cleaned).strip()
+    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")):
+        return "### Summary of Clinical Oversights\nNo critical oversights identified in the provided records."
+
+    sections = {}
+    lines = combined_response.splitlines()
     current_section = None
-    seen_findings = set()
-    for line in combined_response.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
-        if current_section and line.startswith("- ") and line not in seen_findings:
-            sections[current_section].append(line)
-            seen_findings.add(line)
+        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_lines.append(f"### {heading}")
-            for finding in findings:
-                summary_lines.append(f"{finding}\n  - **Risks**: Potential adverse outcomes if not addressed.\n  - **Recommendation**: Immediate clinical review and follow-up.")
+            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:
+        return "### Summary of Clinical Oversights\nNo critical oversights identified."
 
-    result = "### Comprehensive Clinical Oversight Summary\n" + "\n".join(summary_lines) if summary_lines else "### Comprehensive Clinical Oversight Summary\nNo critical oversights identified."
-    logger.debug("Summary length: %d chars", len(result))
-    return result
+    return "### Summary of Clinical Oversights\n" + "\n".join(summary_lines)
 
 def init_agent():
     logger.info("Initializing model...")
@@ -294,9 +214,7 @@ def init_agent():
         tool_files_dict={"new_tool": target_tool_path},
         force_finish=True,
         enable_checker=False,
-        enable_rag=False,
-        init_rag_num=0,
-        step_rag_num=0,
+        step_rag_num=4,
         seed=100,
         additional_default_tools=[],
     )
@@ -308,8 +226,8 @@ def init_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", visible=False)
-        final_summary = gr.Markdown(label="Comprehensive Clinical Oversight Summary")
+        chatbot = gr.Chatbot(label="Detailed Analysis", height=600, type="messages")
+        final_summary = gr.Markdown(label="Summary of Clinical Oversights")
         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")
@@ -317,14 +235,12 @@ def create_ui(agent):
         progress_bar = gr.Progress()
 
         prompt_template = """
-Analyze the patient record excerpt for clinical oversights. Provide a detailed, evidence-based summary in markdown with findings grouped under headings: Missed Diagnoses, Medication Conflicts, Incomplete Assessments, Urgent Follow-up. For each finding, include clinical context, risks, and recommendations. Output only markdown bullet points under headings. If no issues, state "No oversights identified" once.
-
-Patient Record Excerpt:
+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}
 """
 
-        async def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
-            start_time = time.time()
+        def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
             history.append({"role": "user", "content": message})
             yield history, None, ""
 
@@ -335,75 +251,73 @@ Patient Record Excerpt:
                     progress(current / total, desc=f"Extracting text... Page {current}/{total}")
                     return history, None, ""
 
-                futures = [convert_file_to_json(f.name, f.name.split(".")[-1].lower(), update_extraction_progress) for f in files]
-                results = [sanitize_utf8(future) for future in futures]
-                extracted = "\n".join([json.loads(r).get("content", "") for r in results if "content" in json.loads(r)])
-                file_hash_value = file_hash(files[0].name) if files else ""
+                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.append({"role": "assistant", "content": "✅ Text extraction complete."})
                 yield history, None, ""
-                logger.info("Extracted text length: %d chars", len(extracted))
-
-            chunk_size = 3000
-            chunks = [extracted[i:i + chunk_size] for i in range(0, max(len(extracted), 1), chunk_size)] or [""]
-            logger.info("Created %d chunks", len(chunks))
-            for i, chunk in enumerate(chunks):
-                logger.debug("Chunk %d content: %s...", i + 1, chunk[:100])
-            all_responses = []
+
+            chunk_size = 6000
+            chunks = [extracted[i:i + chunk_size] for i in range(0, len(extracted), chunk_size)]
+            combined_response = ""
             batch_size = 2
 
             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(chunk=chunk[:2000]) for chunk in batch_chunks]
+                    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)}")
 
-                    async def process_chunk(prompt):
-                        chunk_response = ""
-                        raw_outputs = []
-                        for chunk_output in agent.run_gradio_chat(
-                            message=prompt, history=[], temperature=0.2, max_new_tokens=512, max_token=2048, 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:
-                                        raw_outputs.append(m.content)
-                                        cleaned = clean_response(m.content)
+                    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(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():
-                                raw_outputs.append(chunk_output)
-                                cleaned = clean_response(chunk_output)
-                                chunk_response += cleaned + "\n\n"
-                        logger.debug("Raw outputs for chunk: %s", raw_outputs[:100])
-                        logger.debug("Chunk response length: %d chars", len(chunk_response))
-                        return chunk_response
-
-                    futures = [process_chunk(prompt) for prompt in batch_prompts]
-                    batch_responses = await asyncio.gather(*futures)
-                    all_responses.extend([resp.strip() for resp in batch_responses if resp.strip()])
-                    torch.cuda.empty_cache()
-                    gc.collect()
-
-                summary = summarize_findings(all_responses)
-                history.append({"role": "assistant", "content": "Analysis complete. See summary below."})
-                yield history, None, summary
+                            batch_responses.append(chunk_response)
+                            torch.cuda.empty_cache()
+                            gc.collect()
+
+                    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")):
+                    history[-1]["content"] = combined_response.strip()
+                else:
+                    history.append({"role": "assistant", "content": "No oversights identified in the provided records."})
 
+                summary = summarize_findings(combined_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(summary)
+                        f.write(combined_response + "\n\n" + summary)
                 yield history, report_path if report_path and os.path.exists(report_path) else None, summary
-                logger.info("Analysis took %.2f seconds", time.time() - start_time)
 
             except Exception as e:
                 logger.error("Analysis error: %s", e)
                 history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
-                yield history, None, f"### Comprehensive Clinical Oversight Summary\nError occurred during analysis: {str(e)}"
-                logger.info("Analysis took %.2f seconds", time.time() - start_time)
+                yield history, None, f"### Summary of Clinical Oversights\nError occurred during analysis: {str(e)}"
 
         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])