Update app.py

app.py

@@ -35,7 +35,6 @@ from txagent.txagent import TxAgent
 MAX_TOKENS = 32768
 MAX_NEW_TOKENS = 2048
 
-
 def clean_response(text: str) -> str:
     try:
         text = text.encode('utf-8', 'surrogatepass').decode('utf-8')
@@ -46,11 +45,9 @@ def clean_response(text: str) -> str:
     text = re.sub(r"[^\n#\-\*\w\s\.,:\(\)]+", "", text)
     return text.strip()
 
-
 def estimate_tokens(text: str) -> int:
     return len(text) // 3.5
 
-
 def extract_text_from_excel(file_path: str) -> str:
     all_text = []
     xls = pd.ExcelFile(file_path)
@@ -62,7 +59,6 @@ def extract_text_from_excel(file_path: str) -> str:
         all_text.extend(sheet_text)
     return "\n".join(all_text)
 
-
 def split_text_into_chunks(text: str, max_tokens: int = MAX_TOKENS) -> List[str]:
     lines = text.split("\n")
     chunks = []
@@ -83,7 +79,6 @@ def split_text_into_chunks(text: str, max_tokens: int = MAX_TOKENS) -> List[str]
         chunks.append("\n".join(current_chunk))
     return chunks
 
-
 def build_prompt_from_text(chunk: str) -> str:
     return f"""
 ### Unstructured Clinical Records
@@ -104,7 +99,6 @@ Please analyze the above and provide:
 - Follow-up Recommendations
 """
 
-
 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")
@@ -125,31 +119,29 @@ def init_agent():
     agent.init_model()
     return agent
 
-
-def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Union[str, None]], None, None]:
-    # Initialize with empty values
+def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Dict[str, Any]], None, None]:
     messages = []
-    report_path = None
-    
+    report_output = {"visible": False, "value": None}
+
     if file is None or not hasattr(file, "name"):
         messages = [{"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."}]
-        yield messages, None
+        yield messages, report_output
         return
 
     try:
-        # Initial processing message
-        messages = [{"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"},
-                   {"role": "assistant", "content": "⏳ Extracting and analyzing data..."}]
-        yield messages, None
+        messages = [
+            {"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"},
+            {"role": "assistant", "content": "⏳ Extracting and analyzing data..."}
+        ]
+        yield messages, report_output
 
         extracted_text = extract_text_from_excel(file.name)
         chunks = split_text_into_chunks(extracted_text)
         chunk_responses = []
 
-        # Process each chunk
         for i, chunk in enumerate(chunks):
             messages.append({"role": "assistant", "content": f"🔍 Analyzing chunk {i+1}/{len(chunks)}..."})
-            yield messages, None
+            yield messages, report_output
             
             prompt = build_prompt_from_text(chunk)
             response = ""
@@ -173,12 +165,11 @@ def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Un
             
             chunk_responses.append(clean_response(response))
             messages.append({"role": "assistant", "content": f"✅ Chunk {i+1} analysis complete"})
-            yield messages, None
+            yield messages, report_output
 
-        # Final summarization
         final_prompt = "\n\n".join(chunk_responses) + "\n\nSummarize the key findings above."
         messages.append({"role": "assistant", "content": "📊 Generating final report..."})
-        yield messages, None
+        yield messages, report_output
 
         stream_text = ""
         for result in agent.run_gradio_chat(
@@ -200,9 +191,8 @@ def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Un
                         stream_text += r.content
             
             messages[-1]["content"] = f"📊 Generating final report...\n\n{clean_response(stream_text)}"
-            yield messages, None
+            yield messages, report_output
 
-        # Save final report
         final_report = f"# \U0001f9e0 Final Patient Report\n\n{clean_response(stream_text)}"
         timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
         report_path = os.path.join(report_dir, f"report_{timestamp}.md")
@@ -211,12 +201,12 @@ def stream_final_report(agent, file) -> Generator[Tuple[List[Dict[str, str]], Un
             f.write(final_report)
 
         messages.append({"role": "assistant", "content": f"✅ Report generated and saved: report_{timestamp}.md"})
-        yield messages, report_path
+        report_output = {"visible": True, "value": report_path}
+        yield messages, report_output
 
     except Exception as e:
         messages.append({"role": "assistant", "content": f"❌ Error processing file: {str(e)}"})
-        yield messages, None
-
+        yield messages, report_output
 
 def create_ui(agent):
     with gr.Blocks(title="Patient History Chat", css=".gradio-container {max-width: 900px !important}") as demo:
@@ -230,8 +220,8 @@ def create_ui(agent):
                     height=600,
                     type="messages",
                     avatar_images=(
-                        None,  # User avatar
-                        "https://i.imgur.com/6wX7Zb4.png"  # Bot avatar
+                        None,
+                        "https://i.imgur.com/6wX7Zb4.png"
                     )
                 )
             with gr.Column(scale=1):
@@ -257,14 +247,8 @@ def create_ui(agent):
             api_name="analyze"
         )
 
-        def show_report(report_path):
-            if report_path:
-                return gr.File(visible=True, value=report_path)
-            return gr.File(visible=False)
-
     return demo
 
-
 if __name__ == "__main__":
     try:
         agent = init_agent()