Update app.py

app.py

@@ -32,10 +32,10 @@ sys.path.insert(0, src_path)
 from txagent.txagent import TxAgent
 
 # Constants
-MAX_MODEL_TOKENS = 32768  # Model's maximum sequence length
-MAX_CHUNK_TOKENS = 8192   # Chunk size aligned with max_num_batched_tokens
-MAX_NEW_TOKENS = 2048     # Maximum tokens for generation
-PROMPT_OVERHEAD = 500     # Estimated tokens for prompt template overhead
+MAX_MODEL_TOKENS = 32768
+MAX_CHUNK_TOKENS = 8192
+MAX_NEW_TOKENS = 2048
+PROMPT_OVERHEAD = 500
 
 def clean_response(text: str) -> str:
     try:
@@ -48,11 +48,9 @@ def clean_response(text: str) -> str:
     return text.strip()
 
 def estimate_tokens(text: str) -> int:
-    """Estimate the number of tokens based on character length."""
-    return len(text) // 3.5 + 1  # Add 1 to avoid zero estimates
+    return len(text) // 3.5 + 1
 
 def extract_text_from_excel(file_path: str) -> str:
-    """Extract text from all sheets in an Excel file."""
     all_text = []
     try:
         xls = pd.ExcelFile(file_path)
@@ -67,10 +65,6 @@ def extract_text_from_excel(file_path: str) -> str:
     return "\n".join(all_text)
 
 def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> List[str]:
-    """
-    Split text into chunks, ensuring each chunk is within token limits,
-    accounting for prompt overhead.
-    """
     effective_max_tokens = max_tokens - PROMPT_OVERHEAD
     if effective_max_tokens <= 0:
         raise ValueError(f"Effective max tokens ({effective_max_tokens}) must be positive.")
@@ -83,7 +77,7 @@ def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> Lis
     for line in lines:
         line_tokens = estimate_tokens(line)
         if current_tokens + line_tokens > effective_max_tokens:
-            if current_chunk:  # Save the current chunk if it's not empty
+            if current_chunk:
                 chunks.append("\n".join(current_chunk))
             current_chunk = [line]
             current_tokens = line_tokens
@@ -97,7 +91,6 @@ def split_text_into_chunks(text: str, max_tokens: int = MAX_CHUNK_TOKENS) -> Lis
     return chunks
 
 def build_prompt_from_text(chunk: str) -> str:
-    """Build a prompt for analyzing a chunk of clinical data."""
     return f"""
 ### Unstructured Clinical Records
 
@@ -118,7 +111,6 @@ Please analyze the above and provide:
 """
 
 def init_agent():
-    """Initialize the TxAgent with model and tool configurations."""
     default_tool_path = os.path.abspath("data/new_tool.json")
     target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
 
@@ -139,148 +131,55 @@ def init_agent():
     return agent
 
 def process_final_report(agent, file, chatbot_state: List[Dict[str, str]]) -> Tuple[List[Dict[str, str]], Union[str, None]]:
-    """Process the Excel file and generate a final report."""
-    messages = chatbot_state if chatbot_state else []
-    report_path = None
-
-    if file is None or not hasattr(file, "name"):
-        messages.append({"role": "assistant", "content": "❌ Please upload a valid Excel file before analyzing."})
-        return messages, report_path
-
-    try:
-        messages.append({"role": "user", "content": f"Processing Excel file: {os.path.basename(file.name)}"})
-        messages.append({"role": "assistant", "content": "⏳ Extracting and analyzing data..."})
-
-        # Extract text and split into chunks
-        extracted_text = extract_text_from_excel(file.name)
-        chunks = split_text_into_chunks(extracted_text, max_tokens=MAX_CHUNK_TOKENS)
-        chunk_responses = []
-
-        # Process each chunk
-        for i, chunk in enumerate(chunks):
-            messages.append({"role": "assistant", "content": f"🔍 Analyzing chunk {i+1}/{len(chunks)}..."})
-            
-            prompt = build_prompt_from_text(chunk)
-            prompt_tokens = estimate_tokens(prompt)
-            if prompt_tokens > MAX_MODEL_TOKENS:
-                messages.append({"role": "assistant", "content": f"❌ Chunk {i+1} prompt too long ({prompt_tokens} tokens). Skipping..."})
-                continue
-
-            response = ""
-            try:
-                for result in agent.run_gradio_chat(
-                    message=prompt,
-                    history=[],
-                    temperature=0.2,
-                    max_new_tokens=MAX_NEW_TOKENS,
-                    max_token=MAX_MODEL_TOKENS,
-                    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
-            except Exception as e:
-                messages.append({"role": "assistant", "content": f"❌ Error analyzing chunk {i+1}: {str(e)}"})
-                continue
-            
-            chunk_responses.append(clean_response(response))
-            messages.append({"role": "assistant", "content": f"✅ Chunk {i+1} analysis complete"})
-
-        if not chunk_responses:
-            messages.append({"role": "assistant", "content": "❌ No valid chunk responses to summarize."})
-            return messages, report_path
-
-        # Summarize chunk responses incrementally to avoid token limit
-        summary = ""
-        current_summary_tokens = 0
-        for i, response in enumerate(chunk_responses):
-            response_tokens = estimate_tokens(response)
-            if current_summary_tokens + response_tokens > MAX_MODEL_TOKENS - PROMPT_OVERHEAD - MAX_NEW_TOKENS:
-                # Summarize current summary
-                summary_prompt = f"Summarize the following analysis:\n\n{summary}\n\nProvide a concise summary."
-                summary_response = ""
-                try:
-                    for result in agent.run_gradio_chat(
-                        message=summary_prompt,
-                        history=[],
-                        temperature=0.2,
-                        max_new_tokens=MAX_NEW_TOKENS,
-                        max_token=MAX_MODEL_TOKENS,
-                        call_agent=False,
-                        conversation=[],
-                    ):
-                        if isinstance(result, str):
-                            summary_response += result
-                        elif hasattr(result, "content"):
-                            summary_response += result.content
-                        elif isinstance(result, list):
-                            for r in result:
-                                if hasattr(r, "content"):
-                                    summary_response += r.content
-                    summary = clean_response(summary_response)
-                    current_summary_tokens = estimate_tokens(summary)
-                except Exception as e:
-                    messages.append({"role": "assistant", "content": f"❌ Error summarizing intermediate results: {str(e)}"})
-                    return messages, report_path
-
-            summary += f"\n\n### Chunk {i+1} Analysis\n{response}"
-            current_summary_tokens += response_tokens
-
-        # Final summarization
-        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 = ""
-        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"# \U0001f9e0 Final Patient Report\n\n{clean_response(final_report_text)}"
-        messages[-1]["content"] = f"📊 Final Report:\n\n{clean_response(final_report_text)}"
-
-        # Save the report
-        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
-        report_path = os.path.join(report_dir, f"report_{timestamp}.md")
-        
-        with open(report_path, 'w') as f:
-            f.write(final_report)
-
-        messages.append({"role": "assistant", "content": f"✅ Report generated and saved: report_{timestamp}.md"})
-
-    except Exception as e:
-        messages.append({"role": "assistant", "content": f"❌ Error processing file: {str(e)}"})
-
-    return messages, report_path
+    from app_backend import run_analysis_pipeline  # Or keep as is if internal
+    return run_analysis_pipeline(agent, file, chatbot_state)
 
 def create_ui(agent):
     """Create the Gradio UI for the patient history analysis tool."""
-    with gr.Blocks(title="Patient History Chat", css=".gradio-container {max-width: 900px !important}") as demo:
-        gr.Markdown("## 🏥 Patient History Analysis Tool")
-        
+    with gr.Blocks(
+        title="Patient History Chat",
+        css="""
+        .gradio-container {
+            max-width: 900px !important;
+            margin: auto;
+            font-family: 'Segoe UI', sans-serif;
+            background-color: #f8f9fa;
+        }
+        .gr-button.primary {
+            background: linear-gradient(to right, #4b6cb7, #182848);
+            color: white;
+            border: none;
+            border-radius: 8px;
+        }
+        .gr-button.primary:hover {
+            background: linear-gradient(to right, #3552a3, #101a3e);
+        }
+        .gr-file-upload, .gr-chatbot, .gr-markdown {
+            background-color: white;
+            border-radius: 10px;
+            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+            padding: 1rem;
+        }
+        .gr-chatbot {
+            border-left: 4px solid #4b6cb7;
+        }
+        .gr-file-upload input {
+            font-size: 0.95rem;
+        }
+        .chat-message-content p {
+            margin: 0.3em 0;
+        }
+        .chat-message-content ul {
+            padding-left: 1.2em;
+            margin: 0.4em 0;
+        }
+        """
+    ) as demo:
+        gr.Markdown("""
+        <h2 style='color:#182848'>🏥 Patient History Analysis Tool</h2>
+        <p style='color:#444;'>Upload an Excel file containing clinical data. The assistant will analyze it for patterns, inconsistencies, and recommendations.</p>
+        """)
+
         with gr.Row():
             with gr.Column(scale=3):
                 chatbot = gr.Chatbot(
@@ -291,7 +190,8 @@ def create_ui(agent):
                     avatar_images=(
                         None,
                         "https://i.imgur.com/6wX7Zb4.png"
-                    )
+                    ),
+                    render_markdown=True
                 )
             with gr.Column(scale=1):
                 file_upload = gr.File(
@@ -301,7 +201,8 @@ def create_ui(agent):
                 )
                 analyze_btn = gr.Button(
                     "🧠 Analyze Patient History",
-                    variant="primary"
+                    variant="primary",
+                    elem_classes="primary"
                 )
                 report_output = gr.File(
                     label="Download Report",
@@ -309,13 +210,20 @@ def create_ui(agent):
                     interactive=False
                 )
 
-        # State to maintain chatbot messages
         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":
+                    content = content.replace("- ", "\n- ")  # Ensure bullet formatting
+                    content = f"<div class='chat-message-content'>{content}</div>"
+                formatted_messages.append({"role": role, "content": content})
             report_update = gr.update(visible=report_path is not None, value=report_path)
-            return messages, report_update, messages
+            return formatted_messages, report_update, formatted_messages
 
         analyze_btn.click(
             fn=update_ui,
@@ -326,6 +234,7 @@ def create_ui(agent):
 
     return demo
 
+
 if __name__ == "__main__":
     try:
         agent = init_agent()
@@ -339,4 +248,4 @@ if __name__ == "__main__":
         )
     except Exception as e:
         print(f"Error: {str(e)}")
-        sys.exit(1)
+        sys.exit(1)