Update app.py

app.py

@@ -44,10 +44,14 @@ os.environ.update({
     "TOKENIZERS_PARALLELISM": "false",
     "CUDA_LAUNCH_BLOCKING": "1"
 })
+
+# Add src path for txagent
 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
+
 # ==================== UTILITY FUNCTIONS ====================
 def sanitize_text(text: str) -> str:
     """Clean and sanitize text input"""
@@ -79,8 +83,12 @@ def log_system_resources(tag: str = "") -> None:
 # ==================== FILE PROCESSING ====================
 class FileProcessor:
     @staticmethod
-    def extract_pdf_text(file_path: str) -> str:
-        """Extract text from PDF with parallel processing"""
+    def extract_pdf_text(file_path: str, cache: Cache) -> str:
+        """Extract text from PDF with caching"""
+        cache_key = f"pdf_{get_file_hash(file_path)}"
+        if cache_key in cache:
+            return cache[cache_key]
+        
         try:
             with pdfplumber.open(file_path) as pdf:
                 total_pages = len(pdf.pages)
@@ -100,31 +108,43 @@ class FileProcessor:
             batches = [(i, min(i+batch_size, total_pages)) for i in range(0, total_pages, batch_size)]
             text_chunks = [""] * total_pages
 
-            with ThreadPoolExecutor(max_workers=6) as executor:
+            with ThreadPoolExecutor(max_workers=2) as executor:
                 futures = [executor.submit(process_page_range, 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
 
-            return "\n\n".join(filter(None, text_chunks))
+            result = "\n\n".join(filter(None, text_chunks))
+            cache[cache_key] = result
+            return result
         except Exception as e:
             logger.error(f"PDF processing error: {e}")
             return f"PDF processing error: {str(e)}"
 
     @staticmethod
-    def excel_to_data(file_path: str) -> List[Dict]:
-        """Convert Excel file to structured data"""
+    def excel_to_data(file_path: str, cache: Cache) -> List[Dict]:
+        """Convert Excel file to structured data with caching"""
+        cache_key = f"excel_{get_file_hash(file_path)}"
+        if cache_key in cache:
+            return cache[cache_key]
+        
         try:
             df = pd.read_excel(file_path, engine='openpyxl', header=None, dtype=str)
             content = df.where(pd.notnull(df), "").astype(str).values.tolist()
-            return [{"filename": os.path.basename(file_path), "rows": content, "type": "excel"}]
+            result = [{"filename": os.path.basename(file_path), "rows": content, "type": "excel"}]
+            cache[cache_key] = result
+            return result
         except Exception as e:
             logger.error(f"Excel processing error: {e}")
             return [{"error": f"Excel processing error: {str(e)}"}]
 
     @staticmethod
-    def csv_to_data(file_path: str) -> List[Dict]:
-        """Convert CSV file to structured data"""
+    def csv_to_data(file_path: str, cache: Cache) -> List[Dict]:
+        """Convert CSV file to structured data with caching"""
+        cache_key = f"csv_{get_file_hash(file_path)}"
+        if cache_key in cache:
+            return cache[cache_key]
+        
         try:
             chunks = []
             for chunk in pd.read_csv(
@@ -135,13 +155,15 @@ class FileProcessor:
             
             df = pd.concat(chunks) if chunks else pd.DataFrame()
             content = df.where(pd.notnull(df), "").astype(str).values.tolist()
-            return [{"filename": os.path.basename(file_path), "rows": content, "type": "csv"}]
+            result = [{"filename": os.path.basename(file_path), "rows": content, "type": "csv"}]
+            cache[cache_key] = result
+            return result
         except Exception as e:
             logger.error(f"CSV processing error: {e}")
             return [{"error": f"CSV processing error: {str(e)}"}]
 
     @classmethod
-    def process_file(cls, file_path: str, file_type: str) -> List[Dict]:
+    def process_file(cls, file_path: str, file_type: str, cache: Cache) -> List[Dict]:
         """Route file processing based on type"""
         processors = {
             "pdf": cls.extract_pdf_text,
@@ -154,7 +176,7 @@ class FileProcessor:
             return [{"error": f"Unsupported file type: {file_type}"}]
         
         try:
-            result = processors[file_type](file_path)
+            result = processors[file_type](file_path, cache)
             if file_type == "pdf":
                 return [{
                     "filename": os.path.basename(file_path),
@@ -173,7 +195,7 @@ class TextProcessor:
         self.tokenizer = AutoTokenizer.from_pretrained("mims-harvard/TxAgent-T1-Llama-3.1-8B")
         self.cache = Cache(DIRECTORIES["cache"], size_limit=10*1024**3)
 
-    def chunk_text(self, text: str, max_tokens: int = 1800) -> List[str]:
+    def chunk_text(self, text: str, max_tokens: int = 1200) -> List[str]:
         """Split text into token-limited chunks"""
         tokens = self.tokenizer.encode(text)
         return [
@@ -184,11 +206,7 @@ class TextProcessor:
     def clean_response(self, text: str) -> str:
         """Clean and format model response"""
         text = sanitize_text(text)
-        text = re.sub(
-            r"\[.*?\]|\bNone\b|To analyze the patient record excerpt.*?medications\."
-            r"|Since the previous attempts.*?\.|I need to.*?medications\."
-            r"|Retrieving tools.*?\.", "", text, flags=re.DOTALL
-        )
+        text = re.sub(r"\[.*?\]|\bNone\b", "", text)
         
         diagnoses = []
         in_diagnoses = False
@@ -236,7 +254,7 @@ class TextProcessor:
                     if diagnosis and not re.match(r"No issues identified", diagnosis, re.IGNORECASE):
                         diagnoses.append(diagnosis)
         
-        unique_diagnoses = list(dict.fromkeys(diagnoses))  # Remove duplicates
+        unique_diagnoses = list(dict.fromkeys(diagnoses))
         
         if not unique_diagnoses:
             return "No missed diagnoses were identified in the provided records."
@@ -302,50 +320,43 @@ class ClinicalOversightApp:
                     full_response += cleaned + " "
                     yield {"role": "assistant", "content": full_response}
 
-    def analyze(self, message: str, history: List[dict], files: List) -> Generator[Dict[str, Any], None, None]:
+    def analyze(self, message: str, history: List[dict], files: List) -> Generator[tuple, None, None]:
         """Main analysis pipeline with proper output formatting"""
-        # Initialize all output components
-        outputs = {
-            "chatbot": history.copy(),
-            "download_output": None,
-            "final_summary": "",
-            "progress_text": {"value": "Starting analysis...", "visible": True}
-        }
-        yield outputs
+        chatbot_output = history.copy()
+        download_output = None
+        final_summary = ""
+        progress_text = {"value": "Starting analysis...", "visible": True}
         
         try:
             # Add user message to history
-            history.append({"role": "user", "content": message})
-            outputs["chatbot"] = history
-            yield outputs
+            chatbot_output.append({"role": "user", "content": message})
+            yield (chatbot_output, download_output, final_summary, progress_text)
             
             # Process uploaded files
+.ArrayIndexOutOfBoundsException: Index -1 out of bounds for length 0
             extracted = []
             file_hash_value = ""
             
             if files:
-                with ThreadPoolExecutor(max_workers=4) as executor:
+                with ThreadPoolExecutor(max_workers=2) as executor:
                     futures = []
                     for f in files:
                         file_type = f.name.split(".")[-1].lower()
-                        futures.append(executor.submit(self.file_processor.process_file, f.name, file_type))
+                        futures.append(executor.submit(self.file_processor.process_file, f.name, file_type, self.text_processor.cache))
                     
                     for i, future in enumerate(as_completed(futures), 1):
                         try:
                             extracted.extend(future.result())
-                            outputs["progress_text"] = self._update_progress(i, len(files), "Processing files")
-                            yield outputs
+                            progress_text = self._update_progress(i, len(files), "Processing files")
+                            yield (chatbot_output, download_output, final_summary, progress_text)
                         except Exception as e:
                             logger.error(f"File processing error: {e}")
                             extracted.append({"error": f"Error processing file: {str(e)}"})
                 
                 file_hash_value = get_file_hash(files[0].name) if files else ""
-                history.append({"role": "assistant", "content": "✅ File processing complete"})
-                outputs.update({
-                    "chatbot": history,
-                    "progress_text": self._update_progress(len(files), len(files), "Files processed")
-                })
-                yield outputs
+                chatbot_output.append({"role": "assistant", "content": "✅ File processing complete"})
+                progress_text = self._update_progress(len(files), len(files), "Files processed")
+                yield (chatbot_output, download_output, final_summary, progress_text)
 
             # Analyze content
             text_content = "\n".join(json.dumps(item) for item in extracted)
@@ -360,159 +371,258 @@ with their potential implications and urgent review recommendations. If no misse
 are found, state 'No missed diagnoses identified'.
 
 Patient Record (Chunk {chunk_idx}/{len(chunks)}):
-{chunk[:1800]}
+{chunk[:1200]}
 """
-                history.append({"role": "assistant", "content": ""})
-                outputs.update({
-                    "chatbot": history,
-                    "progress_text": self._update_progress(chunk_idx, len(chunks), "Analyzing")
-                })
-                yield outputs
+                chatbot_output.append({"role": "assistant", "content": ""})
+                progress_text = self._update_progress(chunk_idx, len(chunks), "Analyzing")
+                yield (chatbot_output, download_output, final_summary, progress_text)
                 
                 # Stream response
                 chunk_response = ""
-                for update in self.process_response_stream(prompt, history):
-                    history[-1] = update
+                for update in self.process_response_stream(prompt, chatbot_output):
+                    chatbot_output[-1] = update
                     chunk_response = update["content"]
-                    outputs.update({
-                        "chatbot": history,
-                        "progress_text": self._update_progress(chunk_idx, len(chunks), "Analyzing")
-                    })
-                    yield outputs
+                    progress_text = self._update_progress(chunk_idx, len(chunks), "Analyzing")
+                    yield (chatbot_output, download_output, final_summary, progress_text)
                 
                 combined_response += f"--- Analysis for Chunk {chunk_idx} ---\n{chunk_response}\n"
                 torch.cuda.empty_cache()
                 gc.collect()
 
             # Generate final outputs
-            summary = self.text_processor.summarize_results(combined_response)
+            final_summary = self.text_processor.summarize_results(combined_response)
             report_path = os.path.join(DIRECTORIES["reports"], 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)
+                    f.write(combined_response + "\n\n" + final_summary)
             
-            outputs.update({
-                "download_output": report_path if report_path else None,
-                "final_summary": summary,
-                "progress_text": {"visible": False}
-            })
-            yield outputs
+            download_output = report_path if report_path and os.path.exists(report_path) else None
+            progress_text = {"visible": False}
+            yield (chatbot_output, download_output, final_summary, progress_text)
 
         except Exception as e:
             logger.error(f"Analysis error: {e}")
-            history.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
-            outputs.update({
-                "chatbot": history,
-                "final_summary": f"Error occurred: {str(e)}",
-                "progress_text": {"visible": False}
-            })
-            yield outputs
+            chatbot_output.append({"role": "assistant", "content": f"❌ Error: {str(e)}"})
+            final_summary = f"Error occurred: {str(e)}"
+            progress_text = {"visible": False}
+            yield (chatbot_output, download_output, final_summary, progress_text)
 
     def _update_progress(self, current: int, total: int, stage: str = "") -> Dict[str, Any]:
         """Format progress update for UI"""
         progress = f"{stage} - {current}/{total}" if stage else f"{current}/{total}"
-        return {"value": progress, "visible": True, "label": f"Progress: {progress}"}
+        return {"value": progress, "visible": True}
 
     def create_interface(self):
-        """Create Gradio interface with improved layout"""
-        with gr.Blocks(
-            theme=gr.themes.Soft(
-                primary_hue="indigo",
-                secondary_hue="blue",
-                neutral_hue="slate"
-            ),
-            title="Clinical Oversight Assistant",
-            css="""
-            .diagnosis-summary {
-                border-left: 4px solid #4f46e5;
-                padding: 12px;
-                background: #f8fafc;
-                border-radius: 4px;
+        """Create Gradio interface with ChatGPT-like design"""
+        css = """
+        body, .gradio-container {
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+            background: var(--background);
+            color: var(--text-color);
+        }
+        .gradio-container {
+            max-width: 800px;
+            margin: 0 auto;
+            padding: 20px;
+        }
+        .chat-container {
+            background: var(--chat-bg);
+            border-radius: 12px;
+            padding: 20px;
+            height: 80vh;
+            overflow-y: auto;
+            box-shadow: 0 2px 8px rgba(0,0,0,0.1);
+        }
+        .message {
+            margin: 10px 0;
+            padding: 12px 16px;
+            border-radius: 12px;
+            max-width: 80%;
+            transition: all 0.2s ease;
+        }
+        .message.user {
+            background: #007bff;
+            color: white;
+            margin-left: auto;
+        }
+        .message.assistant {
+            background: var(--message-bg);
+            color: var(--text-color);
+        }
+        .input-container {
+            display: flex;
+            align-items: center;
+            margin-top: 20px;
+            background: var(--chat-bg);
+            padding: 10px 20px;
+            border-radius: 25px;
+            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+        }
+        .input-textbox {
+            flex-grow: 1;
+            border: none;
+            background: transparent;
+            color: var(--text-color);
+            outline: none;
+        }
+        .send-btn {
+            background: #007bff;
+            color: white;
+            border: none;
+            border-radius: 20px;
+            padding: 8px 16px;
+            margin-left: 10px;
+        }
+        .send-btn:hover {
+            background: #0056b3;
+        }
+        .sidebar {
+            background: var(--sidebar-bg);
+            padding: 20px;
+            border-radius: 12px;
+            margin-top: 20px;
+            box-shadow: 0 2px 8px rgba(0,0,0,0.1);
+        }
+        .sidebar-hidden {
+            display: none;
+        }
+        .header {
+            text-align: center;
+            margin-bottom: 20px;
+        }
+        .theme-toggle {
+            position: absolute;
+            top: 20px;
+            right: 20px;
+            background: #007bff;
+            color: white;
+            border: none;
+            border-radius: 20px;
+            padding: 8px 16px;
+        }
+        :root {
+            --background: #ffffff;
+            --text-color: #333333;
+            --chat-bg: #f7f7f8;
+            --message-bg: #e5e5ea;
+            --sidebar-bg: #f1f1f1;
+        }
+        @media (prefers-color-scheme: dark) {
+            :root {
+                --background: #1e2a44;
+                --text-color: #ffffff;
+                --chat-bg: #2d3b55;
+                --message-bg: #3e4c6a;
+                --sidebar-bg: #2a3650;
             }
-            .file-upload {
-                border: 2px dashed #cbd5e1;
-                border-radius: 8px;
-                padding: 20px;
+        }
+        @media (max-width: 600px) {
+            .gradio-container {
+                padding: 10px;
+            }
+            .chat-container {
+                height: 70vh;
+            }
+            .input-container {
+                flex-direction: column;
+                gap: 10px;
+            }
+            .send-btn {
+                width: 100%;
+                margin-left: 0;
             }
-            """
-        ) as app:
-            # Header Section
-            gr.Markdown("""
-            <div style='text-align: center; margin-bottom: 20px;'>
-                <h1 style='color: #4f46e5;'>🩺 Clinical Oversight Assistant</h1>
-                <p style='color: #64748b;'>
-                    AI-powered analysis of patient records for potential missed diagnoses
+        }
+        """
+        
+        js = """
+        function toggleTheme() {
+            const root = document.documentElement;
+            const isDark = root.style.getPropertyValue('--background') === '#1e2a44';
+            root.style.setProperty('--background', isDark ? '#ffffff' : '#1e2a44');
+            root.style.setProperty('--text-color', isDark ? '#333333' : '#ffffff');
+            root.style.setProperty('--chat-bg', isDark ? '#f7f7f8' : '#2d3b55');
+            root.style.setProperty('--message-bg', isDark ? '#e5e5ea' : '#3e4c6a');
+            root.style.setProperty('--sidebar-bg', isDark ? '#f1f1f1' : '#2a3650');
+            localStorage.setItem('theme', isDark ? 'light' : 'dark');
+        }
+
+        function toggleSidebar() {
+            const sidebar = document.querySelector('.sidebar');
+            sidebar.classList.toggle('sidebar-hidden');
+        }
+
+        document.addEventListener('DOMContentLoaded', () => {
+            const savedTheme = localStorage.getItem('theme');
+            if (savedTheme === 'dark') toggleTheme();
+            document.querySelector('.sidebar').classList.add('sidebar-hidden');
+        });
+        """
+
+        with gr.Blocks(theme=gr.themes.Default(), css=css, js=js, title="Clinical Oversight Assistant") as app:
+            gr.HTML("""
+            <div class='header'>
+                <h1 style='color: var(--text-color);'>🩺 Clinical Oversight Assistant</h1>
+                <p style='color: var(--text-color); opacity: 0.7;'>
+                    AI-powered analysis of patient records for missed diagnoses
                 </p>
             </div>
             """)
-
-            with gr.Row(equal_height=False):
-                # Main Chat Column
-                with gr.Column(scale=3):
-                    chatbot = gr.Chatbot(
-                        label="Clinical Analysis",
-                        height=600,
-                        show_copy_button=True,
-                        avatar_images=(
-                            "assets/user.png", 
-                            "assets/assistant.png"
-                        ) if os.path.exists("assets/user.png") else None,
-                        bubble_full_width=False,
-                        type="messages",
-                        elem_classes=["chat-container"]
-                    )
-
-                # Results Column
-                with gr.Column(scale=1):
-                    with gr.Group():
-                        gr.Markdown("### 📝 Summary of Findings")
-                        final_summary = gr.Markdown(
-                            "Analysis results will appear here...",
-                            elem_classes=["diagnosis-summary"]
-                        )
-                    
-                    with gr.Group():
-                        gr.Markdown("### 📂 Report Download")
-                        download_output = gr.File(
-                            label="Full Report",
-                            visible=False,
-                            interactive=False
-                        )
-
-            # Input Section
+            gr.Button("Toggle Light/Dark Mode", elem_classes="theme-toggle").click(
+                None, None, None, _js="toggleTheme"
+            )
+            
+            with gr.Column(elem_classes="chat-container"):
+                chatbot = gr.Chatbot(
+                    label="Clinical Analysis",
+                    height="100%",
+                    show_copy_button=True,
+                    type="messages",
+                    elem_classes="chatbot"
+                )
+            
             with gr.Row():
+                gr.Button("Show/Hide Tools", variant="secondary").click(
+                    None, None, None, _js="toggleSidebar"
+                )
+            
+            with gr.Column(elem_classes="sidebar"):
                 file_upload = gr.File(
                     file_types=[".pdf", ".csv", ".xls", ".xlsx"],
                     file_count="multiple",
-                    label="Upload Patient Records",
-                    elem_classes=["file-upload"]
+                    label="Upload Patient Records"
                 )
-
-            # Interaction Section
-            with gr.Row():
+                gr.Markdown("### 📝 Summary of Findings")
+                final_summary = gr.Markdown(
+                    "Analysis results will appear here..."
+                )
+                gr.Markdown("### 📂 Report Download")
+                download_output = gr.File(
+                    label="Full Report",
+                    visible=False,
+                    interactive=False
+                )
+            
+            with gr.Row(elem_classes="input-container"):
                 msg_input = gr.Textbox(
                     placeholder="Ask about potential oversights or upload files...",
                     show_label=False,
                     container=False,
-                    scale=7,
+                    elem_classes="input-textbox",
                     autofocus=True
                 )
                 send_btn = gr.Button(
                     "Analyze",
                     variant="primary",
-                    scale=1,
-                    min_width=100
+                    elem_classes="send-btn"
                 )
-
-            # Progress Indicator
+            
             progress_text = gr.Textbox(
                 label="Progress Status",
                 visible=False,
                 interactive=False
             )
 
-            # Event Handlers
             send_btn.click(
                 self.analyze,
                 inputs=[msg_input, chatbot, file_upload],