Update app.py

app.py

@@ -24,10 +24,12 @@ os.makedirs(base_dir, exist_ok=True)
 model_cache_dir = os.path.join(base_dir, "txagent_models")
 tool_cache_dir = os.path.join(base_dir, "tool_cache")
 file_cache_dir = os.path.join(base_dir, "cache")
+report_dir = os.path.join(base_dir, "reports")
 
 os.makedirs(model_cache_dir, exist_ok=True)
 os.makedirs(tool_cache_dir, exist_ok=True)
 os.makedirs(file_cache_dir, exist_ok=True)
+os.makedirs(report_dir, exist_ok=True)
 
 os.environ.update({
     "TRANSFORMERS_CACHE": model_cache_dir,
@@ -73,11 +75,7 @@ def convert_file_to_json(file_path: str, file_type: str) -> str:
 
         if file_type == "pdf":
             text = extract_priority_pages(file_path)
-            result = json.dumps({
-                "filename": os.path.basename(file_path),
-                "content": text,
-                "status": "initial"
-            })
+            result = json.dumps({"filename": os.path.basename(file_path), "content": text, "status": "initial"})
             Thread(target=full_pdf_processing, args=(file_path, h)).start()
 
         elif file_type == "csv":
@@ -108,25 +106,19 @@ def full_pdf_processing(file_path: str, file_hash: str):
         cache_path = os.path.join(file_cache_dir, f"{file_hash}_full.json")
         if os.path.exists(cache_path):
             return
-
         with pdfplumber.open(file_path) as pdf:
             full_text = "\n".join([f"=== Page {i+1} ===\n{(page.extract_text() or '').strip()}" for i, page in enumerate(pdf.pages)])
-
-        result = json.dumps({
-            "filename": os.path.basename(file_path),
-            "content": full_text,
-            "status": "complete"
-        })
-
+        result = json.dumps({"filename": os.path.basename(file_path), "content": full_text, "status": "complete"})
         with open(cache_path, "w", encoding="utf-8") as f:
             f.write(result)
+        with open(os.path.join(report_dir, f"{file_hash}_report.txt"), "w", encoding="utf-8") as out:
+            out.write(full_text)
     except Exception as e:
         print(f"Background processing failed: {str(e)}")
 
 def init_agent():
     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)
 
@@ -149,23 +141,31 @@ def create_ui(agent: TxAgent):
         gr.Markdown("<h3 style='text-align: center;'>Identify potential oversights in patient care</h3>")
 
         chatbot = gr.Chatbot(label="Analysis", height=600, type="messages")
-        file_upload = gr.File(label="Upload Medical Records", file_types=[".pdf", ".csv", ".xls", ".xlsx"], file_count="multiple")
+        file_upload = gr.File(
+            label="Upload Medical Records",
+            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")
         conversation_state = gr.State([])
+        download_output = gr.File(label="Download Full Report (after tools finish)")
 
         def analyze_potential_oversights(message: str, history: list, conversation: list, files: list):
             start_time = time.time()
             try:
                 history.append({"role": "user", "content": message})
                 history.append({"role": "assistant", "content": "Analyzing records for potential oversights..."})
-                yield history
+                yield history, None
 
                 extracted_data = ""
+                file_hash_value = ""
                 if files:
                     with ThreadPoolExecutor(max_workers=4) as executor:
                         futures = [executor.submit(convert_file_to_json, f.name, f.name.split(".")[-1].lower()) for f in files if hasattr(f, 'name')]
-                        extracted_data = "\n".join([sanitize_utf8(f.result()) for f in as_completed(futures)])
+                        results = [sanitize_utf8(f.result()) for f in as_completed(futures)]
+                        extracted_data = "\n".join(results)
+                        file_hash_value = file_hash(files[0].name)
 
                 analysis_prompt = f"""Review these medical records and identify EXACTLY what might have been missed:
 1. List potential missed diagnoses
@@ -173,15 +173,9 @@ def create_ui(agent: TxAgent):
 3. Note incomplete assessments
 4. Highlight abnormal results needing follow-up
 
-Medical Records:
-{extracted_data[:15000]}
-
-Provide ONLY the potential oversights in this format:
+Medical Records:\n{extracted_data[:15000]}
 
-### Potential Oversights:
-1. [Missed diagnosis] - [Evidence from records]
-2. [Medication issue] - [Supporting data]
-3. [Assessment gap] - [Relevant findings]"""
+### Potential Oversights:\n"""
 
                 response = []
                 for chunk in agent.run_gradio_chat(
@@ -197,28 +191,29 @@ Provide ONLY the potential oversights in this format:
                         response.append(chunk)
                     elif isinstance(chunk, list):
                         response.extend([c.content for c in chunk if hasattr(c, 'content')])
+
                     if len(response) % 3 == 0:
                         history[-1] = {"role": "assistant", "content": "".join(response).strip()}
-                        yield history
+                        yield history, None
 
                 final_output = "".join(response).strip()
                 if not final_output:
                     final_output = "No clear oversights identified. Recommend comprehensive review."
-
                 if not final_output.startswith(("1.", "-", "*", "#")):
                     final_output = "• " + final_output.replace("\n", "\n• ")
+                history[-1] = {"role": "assistant", "content": final_output}
 
-                history[-1] = {"role": "assistant", "content": f"### Potential Clinical Oversights:\n{final_output}"}
-                print(f"Analysis completed in {time.time()-start_time:.2f}s")
-                yield history
+                report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
+                return history, report_path if os.path.exists(report_path) else None
 
             except Exception as e:
                 history.append({"role": "assistant", "content": f"❌ Analysis failed: {str(e)}"})
-                yield history
+                return history, None
 
         inputs = [msg_input, chatbot, conversation_state, file_upload]
-        send_btn.click(analyze_potential_oversights, inputs=inputs, outputs=chatbot)
-        msg_input.submit(analyze_potential_oversights, inputs=inputs, outputs=chatbot)
+        outputs = [chatbot, download_output]
+        send_btn.click(analyze_potential_oversights, inputs=inputs, outputs=outputs)
+        msg_input.submit(analyze_potential_oversights, inputs=inputs, outputs=outputs)
 
         gr.Examples([
             ["What might have been missed in this patient's treatment?"],
@@ -231,7 +226,6 @@ Provide ONLY the potential oversights in this format:
 if __name__ == "__main__":
     print("Initializing medical analysis agent...")
     agent = init_agent()
-
     print("Launching interface...")
     demo = create_ui(agent)
     demo.queue().launch(
@@ -239,4 +233,4 @@ if __name__ == "__main__":
         server_port=7860,
         show_error=True,
         share=True
-    )
+    )