Update app.py

app.py

@@ -32,7 +32,6 @@ sys.path.insert(0, src_path)
 
 from txagent.txagent import TxAgent
 
-# Constants
 MAX_MODEL_TOKENS = 32768
 MAX_CHUNK_TOKENS = 8192
 MAX_NEW_TOKENS = 2048
@@ -68,7 +67,7 @@ def extract_text_from_excel(file_path: str) -> str:
 def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> List[str]:
     effective_max_tokens = max_tokens - PROMPT_OVERHEAD
     if effective_max_tokens <= 0:
-        raise ValueError(f"Effective max tokens ({effective_max_tokens}) must be positive.")
+        raise ValueError("Effective max tokens must be positive.")
     lines = text.split("\n")
     chunks, current_chunk, current_tokens = [], [], 0
     for line in lines:
@@ -131,9 +130,7 @@ def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tu
         return messages, report_path
 
     try:
-        messages.append({"role": "user", "content": f"📄 Processing Excel file: {os.path.basename(file.name)}"})
-        messages.append({"role": "assistant", "content": "🔍 Analyzing clinical data... This may take a moment."})
-        
+        messages.append({"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"})
         extracted_text = extract_text_from_excel(file.name)
         chunks = split_text_into_chunks(extracted_text)
         chunk_responses = [None] * len(chunks)
@@ -142,7 +139,7 @@ def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tu
             prompt = build_prompt_from_text(chunk)
             prompt_tokens = estimate_tokens(prompt)
             if prompt_tokens > MAX_MODEL_TOKENS:
-                return index, f"❌ Chunk {index+1} prompt too long ({prompt_tokens} tokens). Skipping..."
+                return index, f"❌ Chunk {index+1} prompt too long. Skipping..."
             response = ""
             try:
                 for result in agent.run_gradio_chat(
@@ -154,79 +151,47 @@ def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tu
                     call_agent=False,
                     conversation=[],
                 ):
-                    if isinstance(result, str):
-                        response += result
-                    elif hasattr(result, "content"):
-                        response += result.content
-                    elif isinstance(result, list):
-                        for r in result:
-                            if hasattr(r, "content"):
-                                response += r.content
+                    response += getattr(result, "content", result) if isinstance(result, (str, list)) else ""
             except Exception as e:
                 return index, f"❌ Error analyzing chunk {index+1}: {str(e)}"
             return index, clean_response(response)
 
-        # Process chunks silently without displaying progress
         with ThreadPoolExecutor(max_workers=1) as executor:
             futures = [executor.submit(analyze_chunk, i, chunk) for i, chunk in enumerate(chunks)]
             for future in as_completed(futures):
                 i, result = future.result()
                 chunk_responses[i] = result
+                if result.startswith("❌"):
+                    messages.append({"role": "assistant", "content": result})
 
         valid_responses = [res for res in chunk_responses if not res.startswith("❌")]
         if not valid_responses:
-            messages.append({"role": "assistant", "content": "❌ No valid analysis results to summarize."})
+            messages.append({"role": "assistant", "content": "❌ No valid chunk responses to summarize."})
             return messages, report_path
 
         summary = "\n\n".join(valid_responses)
-        final_prompt = f"""Please synthesize the following clinical analyses into a concise, well-structured report:
-        
-{summary}
-
-Structure your response with clear sections:
-1. Key Diagnostic Patterns
-2. Medication Concerns
-3. Potential Missed Opportunities
-4. Notable Inconsistencies
-5. Recommended Follow-ups
-
-Use bullet points for clarity and professional medical terminology."""
-        
-        final_report_text = ""
-        try:
-            for result in agent.run_gradio_chat(
-                message=final_prompt,
-                history=[],
-                temperature=0.2,
-                max_new_tokens=MAX_NEW_TOKENS,
-                max_token=MAX_MODEL_TOKENS,
-                call_agent=False,
-                conversation=[],
-            ):
-                if isinstance(result, str):
-                    final_report_text += result
-                elif hasattr(result, "content"):
-                    final_report_text += result.content
-                elif isinstance(result, list):
-                    for r in result:
-                        if hasattr(r, "content"):
-                            final_report_text += r.content
-        except Exception as e:
-            messages.append({"role": "assistant", "content": f"❌ Error generating final report: {str(e)}"})
-            return messages, report_path
-
-        final_report = f"# 🧠 Clinical Analysis Report\n\n{clean_response(final_report_text)}"
-        
-        # Update the last message with the final report
-        messages[-1]["content"] = f"## 📋 Clinical Analysis Report\n\n{clean_response(final_report_text)}"
-
-        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
-        report_path = os.path.join(report_dir, f"clinical_report_{timestamp}.md")
+        final_prompt = f"Summarize the key findings from the following analyses:\n\n{summary}"
+        messages.append({"role": "assistant", "content": "📊 Generating final report..."})
 
+        final_report_text = ""
+        for result in agent.run_gradio_chat(
+            message=final_prompt,
+            history=[],
+            temperature=0.2,
+            max_new_tokens=MAX_NEW_TOKENS,
+            max_token=MAX_MODEL_TOKENS,
+            call_agent=False,
+            conversation=[],
+        ):
+            final_report_text += getattr(result, "content", result) if isinstance(result, (str, list)) else ""
+
+        cleaned = clean_response(final_report_text)
+        report_path = os.path.join(report_dir, f"report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md")
         with open(report_path, 'w') as f:
-            f.write(final_report)
+            f.write(f"# 🧠 Final Patient Report\n\n{cleaned}")
 
-        messages.append({"role": "assistant", "content": f"✅ Report generated successfully. You can download it below."})
+        messages.append({"role": "assistant", "content": f"📊 Final Report:\n\n{cleaned}"})
+        messages.append({"role": "assistant", "content": f"✅ Report generated and saved: {os.path.basename(report_path)}"})
 
     except Exception as e:
         messages.append({"role": "assistant", "content": f"❌ Error processing file: {str(e)}"})
@@ -234,155 +199,45 @@ Use bullet points for clarity and professional medical terminology."""
     return messages, report_path
 
 def create_ui(agent):
-    with gr.Blocks(
-        title="Clinical Analysis Tool",
-        css="""
+    with gr.Blocks(css="""
         .gradio-container {
-            max-width: 900px !important;
+            max-width: 900px;
             margin: auto;
-            font-family: 'Inter', sans-serif;
-            background-color: #f9fafb;
+            font-family: 'Segoe UI', sans-serif;
+            background-color: #f9fafc;
         }
         .gr-button.primary {
-            background: linear-gradient(to right, #4f46e5, #7c3aed);
+            background: linear-gradient(to right, #4b6cb7, #182848);
             color: white;
             border: none;
             border-radius: 8px;
-            padding: 12px 24px;
-            font-weight: 500;
-            transition: all 0.2s;
-        }
-        .gr-button.primary:hover {
-            background: linear-gradient(to right, #4338ca, #6d28d9);
-            transform: translateY(-1px);
-            box-shadow: 0 4px 6px rgba(0,0,0,0.1);
-        }
-        .gr-file-upload, .gr-chatbot, .gr-markdown {
-            background-color: white;
-            border-radius: 12px;
-            box-shadow: 0 1px 3px rgba(0,0,0,0.05);
-            padding: 1.5rem;
-            border: 1px solid #e5e7eb;
-        }
-        .gr-chatbot {
-            min-height: 600px;
-            border-left: none;
         }
-        .chat-message-user {
-            background-color: #f3f4f6;
-            border-radius: 12px;
-            padding: 12px 16px;
-            margin: 8px 0;
-        }
-        .chat-message-assistant {
+        .gr-chatbot, .gr-markdown, .gr-file-upload {
             background-color: white;
-            border-radius: 12px;
-            padding: 12px 16px;
-            margin: 8px 0;
-            border: 1px solid #e5e7eb;
-        }
-        .chat-message-content ul, .chat-message-content ol {
-            padding-left: 1.5em;
-            margin: 0.5em 0;
-        }
-        .chat-message-content li {
-            margin: 0.3em 0;
-        }
-        h1, h2, h3 {
-            color: #111827;
-        }
-        .gr-markdown h1 {
-            font-size: 1.8rem;
-            margin-bottom: 1rem;
-            font-weight: 600;
+            border-radius: 10px;
+            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
         }
-        .gr-markdown p {
-            color: #4b5563;
-            line-height: 1.6;
-        }
-        .progress-bar {
-            height: 4px;
-            background: #e5e7eb;
-            border-radius: 2px;
-            margin: 12px 0;
-            overflow: hidden;
-        }
-        .progress-bar-fill {
-            height: 100%;
-            background: linear-gradient(to right, #4f46e5, #7c3aed);
-            transition: width 0.3s ease;
-        }
-        """
-    ) as demo:
+    """) as demo:
         gr.Markdown("""
-        <div style='text-align: center; margin-bottom: 1.5rem;'>
-            <h1 style='margin-bottom: 0.5rem; color: #111827;'>Clinical Documentation Analyzer</h1>
-            <p style='color: #6b7280; margin-top: 0;'>Upload patient records in Excel format for comprehensive clinical analysis</p>
-        </div>
+        <h2 style='color:#182848'>🏥 Patient History Analysis Tool</h2>
+        <p>Upload your clinical Excel file to receive a professional diagnostic summary.</p>
         """)
 
         with gr.Row():
             with gr.Column(scale=3):
-                chatbot = gr.Chatbot(
-                    label="Analysis Results",
-                    show_copy_button=True,
-                    height=600,
-                    bubble_full_width=False,
-                    avatar_images=(None, "https://i.imgur.com/6wX7Zb4.png"),
-                    render_markdown=True
-                )
+                chatbot = gr.Chatbot(label="Clinical Assistant", height=600, type="messages")
             with gr.Column(scale=1):
-                file_upload = gr.File(
-                    label="Upload Patient Records",
-                    file_types=[".xlsx", ".xls"],
-                    height=100,
-                    interactive=True
-                )
-                analyze_btn = gr.Button(
-                    "Analyze Clinical Data",
-                    variant="primary",
-                    elem_classes="primary"
-                )
-                report_output = gr.File(
-                    label="Download Report",
-                    visible=False,
-                    interactive=False
-                )
-                gr.Markdown("""
-                <div style='margin-top: 1rem; padding: 1rem; background-color: #f8fafc; border-radius: 8px;'>
-                    <h3 style='margin-top: 0; margin-bottom: 0.5rem; font-size: 1rem;'>About this tool</h3>
-                    <p style='margin: 0; font-size: 0.9rem; color: #64748b;'>
-                        This tool analyzes clinical documentation to identify patterns, inconsistencies, and opportunities for improved patient care.
-                    </p>
-                </div>
-                """)
+                file_upload = gr.File(label="Upload Excel File", file_types=[".xlsx"])
+                analyze_btn = gr.Button("🧠 Analyze", variant="primary")
+                report_output = gr.File(label="Download Report", visible=False, interactive=False)
 
         chatbot_state = gr.State(value=[])
 
         def update_ui(file, current_state):
             messages, report_path = process_final_report(agent, file, current_state)
-            formatted_messages = []
-            for msg in messages:
-                role = msg.get("role")
-                content = msg.get("content", "")
-                if role == "assistant":
-                    # Format lists and sections for better readability
-                    content = content.replace("- ", "• ")
-                    content = re.sub(r"(\d+\.\s)", r"\n\1", content)
-                    content = f"<div class='chat-message-assistant'>{content}</div>"
-                else:
-                    content = f"<div class='chat-message-user'>{content}</div>"
-                formatted_messages.append({"role": role, "content": content})
-            
-            report_update = gr.update(visible=report_path is not None, value=report_path)
-            return formatted_messages, report_update, formatted_messages
-
-        analyze_btn.click(
-            fn=update_ui,
-            inputs=[file_upload, chatbot_state],
-            outputs=[chatbot, report_output, chatbot_state],
-            api_name="analyze"
-        )
+            return messages, gr.update(visible=report_path is not None, value=report_path), messages
+
+        analyze_btn.click(fn=update_ui, inputs=[file_upload, chatbot_state], outputs=[chatbot, report_output, chatbot_state])
 
     return demo
 
@@ -390,13 +245,7 @@ if __name__ == "__main__":
     try:
         agent = init_agent()
         demo = create_ui(agent)
-        demo.launch(
-            server_name="0.0.0.0",
-            server_port=7860,
-            show_error=True,
-            allowed_paths=["/data/hf_cache/reports"],
-            share=False
-        )
+        demo.launch(server_name="0.0.0.0", server_port=7860, allowed_paths=["/data/hf_cache/reports"], share=False)
     except Exception as e:
         print(f"Error: {str(e)}")
-        sys.exit(1)
+        sys.exit(1)