Update app.py

app.py

@@ -33,17 +33,16 @@ sys.path.insert(0, src_path)
 from txagent.txagent import TxAgent
 
 # Constants
-MAX_TOKENS = 32768  # TxAgent's maximum token limit
-CHUNK_SIZE = 3000  # Target chunk size to stay under token limit
-MAX_NEW_TOKENS = 1024
+MAX_TOKENS = 32768
+CHUNK_SIZE = 10000
+MAX_NEW_TOKENS = 2048
+MAX_BOOKINGS_PER_CHUNK = 5
 
 def file_hash(path: str) -> str:
-    """Generate MD5 hash of file contents"""
     with open(path, "rb") as f:
         return hashlib.md5(f.read()).hexdigest()
 
 def clean_response(text: str) -> str:
-    """Clean and normalize text output"""
     try:
         text = text.encode('utf-8', 'surrogatepass').decode('utf-8')
     except UnicodeError:
@@ -55,21 +54,19 @@ def clean_response(text: str) -> str:
     return text.strip()
 
 def estimate_tokens(text: str) -> int:
-    """Approximate token count (1 token ~ 4 characters)"""
-    return len(text) // 4
+    return len(text) // 3.5
 
 def process_patient_data(df: pd.DataFrame) -> Dict[str, Any]:
-    """Process raw patient data into structured format"""
     data = {
         'bookings': defaultdict(list),
         'medications': defaultdict(list),
         'diagnoses': defaultdict(list),
         'tests': defaultdict(list),
+        'procedures': defaultdict(list),
         'doctors': set(),
         'timeline': []
     }
     
-    # Sort by date and group by booking
     df = df.sort_values('Interview Date')
     for booking, group in df.groupby('Booking Number'):
         for _, row in group.iterrows():
@@ -87,121 +84,78 @@ def process_patient_data(df: pd.DataFrame) -> Dict[str, Any]:
             data['timeline'].append(entry)
             data['doctors'].add(entry['doctor'])
             
-            # Categorize entries
             form_lower = entry['form'].lower()
             if 'medication' in form_lower or 'drug' in form_lower:
                 data['medications'][entry['item']].append(entry)
-            elif 'diagnosis' in form_lower:
+            elif 'diagnosis' in form_lower or 'condition' in form_lower:
                 data['diagnoses'][entry['item']].append(entry)
-            elif 'test' in form_lower or 'lab' in form_lower:
+            elif 'test' in form_lower or 'lab' in form_lower or 'result' in form_lower:
                 data['tests'][entry['item']].append(entry)
+            elif 'procedure' in form_lower or 'surgery' in form_lower:
+                data['procedures'][entry['item']].append(entry)
     
     return data
 
-def generate_analysis_prompt(patient_data: Dict[str, Any], booking: str) -> str:
-    """Generate focused analysis prompt for a booking"""
-    booking_entries = patient_data['bookings'][booking]
+def generate_analysis_prompt(patient_data: Dict[str, Any], bookings: List[str]) -> str:
+    prompt_lines = [
+        "**Comprehensive Patient Analysis**",
+        f"Analyzing {len(bookings)} bookings",
+        "",
+        "**Key Analysis Points:**",
+        "- Chronological progression of symptoms",
+        "- Medication changes and interactions",
+        "- Diagnostic consistency across providers",
+        "- Missed diagnostic opportunities",
+        "- Gaps in follow-up",
+        "",
+        "**Patient Timeline:**"
+    ]
     
-    # Build timeline string
-    timeline = "\n".join(
-        f"- {entry['date']}: {entry['form']} - {entry['item']} = {entry['response']} (by {entry['doctor']})"
-        for entry in booking_entries
-    )
-    
-    # Get current medications
-    current_meds = []
-    for med, entries in patient_data['medications'].items():
-        if any(e['booking'] == booking for e in entries):
-            latest = max((e for e in entries if e['booking'] == booking), key=lambda x: x['date'])
-            current_meds.append(f"- {med}: {latest['response']} (as of {latest['date']})")
+    for entry in patient_data['timeline']:
+        if entry['booking'] in bookings:
+            prompt_lines.append(
+                f"- {entry['date']}: {entry['form']} - {entry['item']} = {entry['response']} (by {entry['doctor']})"
+            )
     
-    # Get current diagnoses
-    current_diags = []
-    for diag, entries in patient_data['diagnoses'].items():
-        if any(e['booking'] == booking for e in entries):
-            latest = max((e for e in entries if e['booking'] == booking), key=lambda x: x['date'])
-            current_diags.append(f"- {diag}: {latest['response']} (as of {latest['date']})")
+    prompt_lines.extend([
+        "",
+        "**Medication History:**",
+        *[f"- {med}: " + " → ".join(
+            f"{e['date']}: {e['response']}" 
+            for e in entries if e['booking'] in bookings
+        ) for med, entries in patient_data['medications'].items()],
+        "",
+        "**Required Analysis Format:**",
+        "### Diagnostic Patterns",
+        "### Medication Analysis",
+        "### Provider Consistency",
+        "### Missed Opportunities",
+        "### Recommendations"
+    ])
     
-    prompt = """
-**Comprehensive Patient Analysis - Booking {booking}**
-
-**Patient Timeline:**
-{timeline}
-
-**Current Medications:**
-{meds}
-
-**Current Diagnoses:**
-{diags}
-
-**Analysis Instructions:**
-1. Review the patient's complete history across all visits
-2. Identify any potential missed diagnoses based on symptoms and test results
-3. Check for medication conflicts or inappropriate prescriptions
-4. Note any incomplete assessments or missing tests
-5. Flag any urgent follow-up needs
-6. Compare findings across different doctors for consistency
-
-**Required Output Format:**
-### Missed Diagnoses
-[Potential diagnoses that were not identified]
-
-### Medication Issues
-[Conflicts, side effects, inappropriate prescriptions]
+    return "\n".join(prompt_lines)
 
-### Assessment Gaps
-[Missing tests or incomplete evaluations]
-
-### Follow-up Recommendations
-[Urgent and non-urgent follow-up needs]
-
-### Doctor Consistency
-[Discrepancies between different providers]
-""".format(
-        booking=booking,
-        timeline=timeline,
-        meds='\n'.join(current_meds) if current_meds else "None recorded",
-        diags='\n'.join(current_diags) if current_diags else "None recorded"
-    )
-    return prompt
-
-def chunk_patient_data(patient_data: Dict[str, Any]) -> List[Dict[str, Any]]:
-    """Split patient data into manageable chunks"""
-    chunks = []
-    current_chunk = defaultdict(list)
-    current_size = 0
+def chunk_bookings(patient_data: Dict[str, Any]) -> List[List[str]]:
+    all_bookings = list(patient_data['bookings'].keys())
+    booking_sizes = []
     
-    for booking, entries in patient_data['bookings'].items():
-        booking_size = sum(estimate_tokens(str(e)) for e in entries)
-        
-        if current_size + booking_size > CHUNK_SIZE and current_chunk:
-            chunks.append(dict(current_chunk))
-            current_chunk = defaultdict(list)
-            current_size = 0
-            
-        current_chunk['bookings'][booking] = entries
-        current_size += booking_size
-        
-        # Add related data
-        for med, med_entries in patient_data['medications'].items():
-            if any(e['booking'] == booking for e in med_entries):
-                current_chunk['medications'][med].extend(
-                    e for e in med_entries if e['booking'] == booking
-                )
-        
-        for diag, diag_entries in patient_data['diagnoses'].items():
-            if any(e['booking'] == booking for e in diag_entries):
-                current_chunk['diagnoses'][diag].extend(
-                    e for e in diag_entries if e['booking'] == booking
-                )
+    for booking in all_bookings:
+        entries = patient_data['bookings'][booking]
+        size = sum(estimate_tokens(str(e)) for e in entries)
+        booking_sizes.append((booking, size))
     
-    if current_chunk:
-        chunks.append(dict(current_chunk))
+    booking_sizes.sort(key=lambda x: x[1], reverse=True)
+    chunks = [[] for _ in range(3)]
+    chunk_sizes = [0, 0, 0]
+    
+    for booking, size in booking_sizes:
+        min_chunk = chunk_sizes.index(min(chunk_sizes))
+        chunks[min_chunk].append(booking)
+        chunk_sizes[min_chunk] += size
     
     return chunks
 
 def init_agent():
-    """Initialize TxAgent with proper configuration"""
     default_tool_path = os.path.abspath("data/new_tool.json")
     target_tool_path = os.path.join(tool_cache_dir, "new_tool.json")
     
@@ -216,13 +170,12 @@ def init_agent():
         enable_checker=True,
         step_rag_num=4,
         seed=100,
-        additional_default_tools=[],
+        additional_default_tools=[]
     )
     agent.init_model()
     return agent
 
 def analyze_with_agent(agent, prompt: str) -> str:
-    """Run analysis with proper error handling"""
     try:
         response = ""
         for result in agent.run_gradio_chat(
@@ -249,124 +202,76 @@ def analyze_with_agent(agent, prompt: str) -> str:
 
 def create_ui(agent):
     with gr.Blocks(theme=gr.themes.Soft(), title="Patient History Analyzer") as demo:
-        gr.Markdown("# 🏥 Comprehensive Patient History Analysis")
+        gr.Markdown("# 🏥 Patient History Analyzer")
         
         with gr.Tabs():
             with gr.TabItem("Analysis"):
                 with gr.Row():
                     with gr.Column(scale=1):
                         file_upload = gr.File(
-                            label="Upload Patient Excel File",
+                            label="Upload Excel File",
                             file_types=[".xlsx"],
                             file_count="single"
                         )
-                        analysis_btn = gr.Button("Analyze Patient History", variant="primary")
-                        status = gr.Markdown("Ready for analysis")
+                        analyze_btn = gr.Button("Analyze", variant="primary")
+                        status = gr.Markdown("Ready")
                     
                     with gr.Column(scale=2):
-                        output_display = gr.Markdown(
-                            label="Analysis Results",
-                            elem_id="results"
-                        )
-                        report_download = gr.File(
-                            label="Download Full Report",
-                            interactive=False
-                        )
+                        output = gr.Markdown()
+                        report = gr.File(label="Download Report")
             
             with gr.TabItem("Instructions"):
                 gr.Markdown("""
-                ## How to Use This Tool
-                
-                1. **Upload Excel File**: Patient history Excel file
-                2. **Click Analyze**: System will process all bookings
-                3. **Review Results**: Comprehensive analysis appears
-                4. **Download Report**: Full report with all findings
+                ## How to Use
+                1. Upload patient history Excel
+                2. Click Analyze
+                3. View/download report
                 
-                ### Excel Requirements
-                Must contain these columns:
+                **Required Columns:**
                 - Booking Number
                 - Interview Date
-                - Interviewer (Doctor)
+                - Interviewer
                 - Form Name
                 - Form Item
                 - Item Response
                 - Description
-                
-                ### Analysis Includes:
-                - Missed diagnoses across visits
-                - Medication conflicts over time
-                - Incomplete assessments
-                - Doctor consistency checks
-                - Follow-up recommendations
                 """)
         
-        def analyze_patient(file) -> Tuple[str, str]:
+        def analyze(file):
             if not file:
-                raise gr.Error("Please upload an Excel file first")
+                raise gr.Error("Please upload a file")
             
             try:
-                # Process Excel file
                 df = pd.read_excel(file.name)
                 patient_data = process_patient_data(df)
-                
-                # Generate and process prompts
+                chunks = chunk_bookings(patient_data)
                 full_report = []
-                bookings_processed = 0
                 
-                for booking in patient_data['bookings']:
-                    prompt = generate_analysis_prompt(patient_data, booking)
+                for i, bookings in enumerate(chunks, 1):
+                    prompt = generate_analysis_prompt(patient_data, bookings)
                     response = analyze_with_agent(agent, prompt)
-                    
-                    if "Error in analysis" not in response:
-                        bookings_processed += 1
-                        full_report.append(f"## Booking {booking}\n{response}\n")
-                    
+                    full_report.append(f"## Chunk {i}\n{response}\n")
                     yield "\n".join(full_report), None
-                    time.sleep(0.1)  # Prevent UI freezing
                 
-                # Generate overall summary
-                if bookings_processed > 1:
-                    summary_prompt = """
-**Comprehensive Patient Summary**
-
-Analyze all bookings ({0} total) to identify:
-1. Patterns across the entire treatment history
-2. Chronic issues that may have been missed
-3. Medication changes over time
-4. Doctor consistency across visits
-5. Long-term recommendations
-
-**Required Format:**
-### Chronic Health Patterns
-[Recurring issues over time]
-
-### Treatment Evolution
-[How treatment has changed]
-
-### Long-term Concerns
-[Issues needing ongoing attention]
-
-### Comprehensive Recommendations
-[Overall care plan]
-""".format(bookings_processed)
+                # Final summary
+                if len(chunks) > 1:
+                    summary_prompt = "Create final summary combining all chunks"
                     summary = analyze_with_agent(agent, summary_prompt)
-                    full_report.append(f"## Overall Patient Summary\n{summary}\n")
+                    full_report.append(f"## Final Summary\n{summary}\n")
                 
-                # Save report
-                report_path = os.path.join(report_dir, f"patient_report_{datetime.now().strftime('%Y%m%d_%H%M%S')}.md")
-                with open(report_path, 'w', encoding='utf-8') as f:
+                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("\n".join(full_report))
                 
                 yield "\n".join(full_report), report_path
-                
+            
             except Exception as e:
-                raise gr.Error(f"Analysis failed: {str(e)}")
+                raise gr.Error(f"Error: {str(e)}")
         
-        analysis_btn.click(
-            analyze_patient,
+        analyze_btn.click(
+            analyze,
             inputs=file_upload,
-            outputs=[output_display, report_download],
-            api_name="analyze"
+            outputs=[output, report]
         )
     
     return demo
@@ -375,17 +280,11 @@ if __name__ == "__main__":
     try:
         agent = init_agent()
         demo = create_ui(agent)
-        
-        demo.queue(
-            api_open=False,
-            max_size=20
-        ).launch(
+        demo.launch(
             server_name="0.0.0.0",
             server_port=7860,
-            show_error=True,
-            allowed_paths=[report_dir],
-            share=False
+            show_error=True
         )
     except Exception as e:
-        print(f"Failed to launch application: {str(e)}")
+        print(f"Error: {str(e)}")
         sys.exit(1)