Update app.py

app.py

@@ -146,8 +146,6 @@ def init_agent():
             seed=100,
             additional_default_tools=[],
         )
-        # This call attempts to load the models. If device inference fails,
-        # it will now produce DEBUG-level logs.
         agent.init_model()
         log_system_usage("After Load")
         print("✅ Agent Ready")
@@ -168,10 +166,12 @@ def create_ui(agent):
 
         def analyze(message: str, history: list, files: list):
             try:
+                # Initialize response with loading message
                 history.append({"role": "user", "content": message})
                 history.append({"role": "assistant", "content": "⏳ Analyzing records for potential oversights..."})
                 yield history, None
 
+                # Process files in parallel
                 extracted = ""
                 file_hash_value = ""
                 if files:
@@ -206,7 +206,12 @@ Medical Records:
                 print("🔎 Generated prompt:")
                 print(prompt)
 
+                # Initialize response tracking
                 full_response = ""
+                last_update_time = 0
+                response_chunks = []
+
+                # Process streaming response
                 for chunk in agent.run_gradio_chat(
                     message=prompt,
                     history=[],
@@ -219,26 +224,40 @@ Medical Records:
                     try:
                         if chunk is None:
                             continue
+                            
+                        # Handle different chunk types
                         if isinstance(chunk, str):
-                            print("DEBUG: Received string chunk")
-                            full_response += chunk
+                            chunk_content = chunk
                         elif isinstance(chunk, list):
                             chunk_content = "".join([c.content for c in chunk if hasattr(c, "content") and c.content])
-                            print("DEBUG: Received list chunk:", chunk_content)
-                            full_response += chunk_content
                         else:
                             print("DEBUG: Received unknown type chunk", type(chunk))
+                            continue
+
+                        if not chunk_content:
+                            continue
+
+                        response_chunks.append(chunk_content)
+                        full_response = "".join(response_chunks)
+
+                        # Update the chat history with the latest response
+                        if len(history) > 0 and history[-1]["role"] == "assistant":
+                            history[-1]["content"] = full_response
+                        else:
+                            history.append({"role": "assistant", "content": full_response})
+
+                        yield history, None
+
                     except Exception as e:
                         print("❌ Error processing chunk:", str(e))
                         traceback.print_exc()
+                        continue
 
-                    # Yield intermediate raw response for debugging purposes
-                    history[-1] = {"role": "assistant", "content": full_response}
-                    yield history, None
-
+                # Final response handling
                 if not full_response:
                     full_response = "⚠️ No clear oversights identified or model output was invalid."
 
+                # Save report if we have files
                 report_path = None
                 if file_hash_value:
                     report_path = os.path.join(report_dir, f"{file_hash_value}_report.txt")
@@ -249,14 +268,19 @@ Medical Records:
                         print("❌ Error saving report:", str(e))
                         traceback.print_exc()
 
-                history[-1] = {"role": "assistant", "content": full_response}
+                # Ensure the final response is in the history
+                if len(history) > 0 and history[-1]["role"] == "assistant":
+                    history[-1]["content"] = full_response
+                else:
+                    history.append({"role": "assistant", "content": full_response})
+
                 yield history, report_path if report_path and os.path.exists(report_path) else None
 
             except Exception as e:
                 error_message = f"❌ An error occurred in analyze: {str(e)}"
                 print(error_message)
                 traceback.print_exc()
-                history[-1] = {"role": "assistant", "content": error_message}
+                history.append({"role": "assistant", "content": error_message})
                 yield history, None
 
         send_btn.click(analyze, inputs=[msg_input, gr.State([]), file_upload], outputs=[chatbot, download_output])
@@ -277,4 +301,4 @@ if __name__ == "__main__":
         )
     except Exception as e:
         print("❌ Fatal error during launch:", str(e))
-        traceback.print_exc()
+        traceback.print_exc()