Update app.py

app.py

@@ -63,7 +63,7 @@ def extract_all_pages(file_path: str, progress_callback=None) -> str:
             if total_pages == 0:
                 return ""
 
-        batch_size = 10  # Process 10 pages per thread
+        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
@@ -143,25 +143,13 @@ def clean_response(text: str) -> str:
     text = re.sub(r"\n{3,}", "\n\n", text)
     text = re.sub(r"[^\n#\-\*\w\s\.\,\:\(\)]+", "", text)
 
-    tool_to_heading = {
-        "get_abuse_info_by_drug_name": "Drugs",
-        "get_dependence_info_by_drug_name": "Drugs",
-        "get_abuse_types_and_related_adverse_reactions_and_controlled_substance_status_by_drug_name": "Drugs",
-        "get_info_for_patients_by_drug_name": "Drugs",
-    }
-
     sections = {}
     current_section = None
-    current_tool = None
     lines = text.splitlines()
     for line in lines:
         line = line.strip()
         if not line:
             continue
-        tool_match = re.match(r"\[TOOL:\s*(\w+)\]", line)
-        if tool_match:
-            current_tool = tool_match.group(1)
-            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)
@@ -170,13 +158,7 @@ def clean_response(text: str) -> str:
             continue
         finding_match = re.match(r"-\s*.+", line)
         if finding_match and current_section and not re.match(r"-\s*No issues identified", line):
-            if current_tool and current_tool in tool_to_heading:
-                heading = tool_to_heading[current_tool]
-                if heading not in sections:
-                    sections[heading] = []
-                sections[heading].append(line)
-            else:
-                sections[current_section].append(line)
+            sections[current_section].append(line)
 
     cleaned = []
     for heading, findings in sections.items():
@@ -186,6 +168,38 @@ def clean_response(text: str) -> str:
     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
+    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:
+        return "### Summary of Clinical Oversights\nNo critical oversights identified."
+
+    return "### Summary of Clinical Oversights\n" + "\n".join(summary_lines)
+
 def init_agent():
     logger.info("Initializing model...")
     log_system_usage("Before Load")
@@ -199,7 +213,7 @@ 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,  # Disabled for speed
+        enable_checker=False,
         step_rag_num=4,
         seed=100,
         additional_default_tools=[],
@@ -212,7 +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="Analysis", height=600, type="messages")
+        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")
@@ -220,7 +235,7 @@ def create_ui(agent):
         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 tool-derived headings (e.g., 'Drugs'). For each finding, include clinical context, risks, and recommendations. Precede findings with a tool tag (e.g., [TOOL: get_abuse_info_by_drug_name]). Output only markdown bullet points under headings. If no issues, state "No issues identified".
+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}
@@ -228,14 +243,14 @@ Patient Record Excerpt (Chunk {0} of {1}):
 
         def analyze(message: str, history: List[dict], files: List, progress=gr.Progress()):
             history.append({"role": "user", "content": message})
-            yield history, None
+            yield history, None, ""
 
             extracted = ""
             file_hash_value = ""
             if files:
                 def update_extraction_progress(current, total):
                     progress(current / total, desc=f"Extracting text... Page {current}/{total}")
-                    return history, None
+                    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]
@@ -244,7 +259,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
                     file_hash_value = file_hash(files[0].name) if files else ""
 
                 history.append({"role": "assistant", "content": "✅ Text extraction complete."})
-                yield history, None
+                yield history, None, ""
 
             chunk_size = 6000
             chunks = [extracted[i:i + chunk_size] for i in range(0, len(extracted), chunk_size)]
@@ -273,7 +288,7 @@ Patient Record Excerpt (Chunk {0} of {1}):
                                             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)
+                                    cleaned = clean_response(m.content)
                                     if cleaned and re.search(r"###\s*\w+", cleaned):
                                         chunk_response += cleaned + "\n\n"
                             batch_responses.append(chunk_response)
@@ -286,26 +301,27 @@ Patient Record Excerpt (Chunk {0} of {1}):
                         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
+                        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(combined_response)
-                yield history, report_path if report_path and os.path.exists(report_path) else None
+                        f.write(combined_response + "\n\n" + summary)
+                yield history, report_path if report_path and os.path.exists(report_path) else None, summary
 
             except Exception as e:
                 logger.error("Analysis error: %s", e)
                 history.append({"role": "assistant", "content": f"❌ Error occurred: {str(e)}"})
-                yield history, None
+                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])
-        msg_input.submit(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])
+        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])
     return demo
 
 if __name__ == "__main__":