integrate_process_flow.py

# integrate_process_flow.py
"""
Main integration script to add Process Flow Visualization to Research Assistant
This script modifies the existing app.py to include comprehensive process flow visualization
"""

import gradio as gr
import logging
import uuid
import time
from typing import Dict, Any, List, Tuple
from process_flow_visualizer import (
    create_process_flow_tab,
    update_process_flow_visualization,
    clear_flow_history,
    export_flow_data
)

logger = logging.getLogger(__name__)

def modify_app_py():
    """
    This function contains the modifications needed for app.py
    Copy these modifications into your existing app.py file
    """
    
    # 1. Add this import at the top of app.py
    import_statement = """
# Add this import at the top of app.py
from process_flow_visualizer import (
    create_process_flow_tab,
    update_process_flow_visualization,
    clear_flow_history,
    export_flow_data
)
"""
    
    # 2. Modify the create_mobile_optimized_interface function
    interface_modification = """
# MODIFY the create_mobile_optimized_interface function in app.py
# Add this after the existing tabs (around line 233)

# NEW: Process Flow Tab
process_flow_tab = create_process_flow_tab(interface_components)
interface_components['process_flow_tab'] = process_flow_tab
"""
    
    # 3. Modify the mobile navigation
    navigation_modification = """
# MODIFY the mobile navigation section (around line 235)
# Add the Process Flow button

with gr.Row(visible=False, elem_id="mobile_nav") as mobile_navigation:
    chat_nav_btn = gr.Button("💬 Chat", variant="secondary", size="sm", min_width=0)
    details_nav_btn = gr.Button("🔍 Details", variant="secondary", size="sm", min_width=0)
    flow_nav_btn = gr.Button("🔄 Flow", variant="secondary", size="sm", min_width=0)  # NEW
    settings_nav_btn = gr.Button("⚙️ Settings", variant="secondary", size="sm", min_width=0)

interface_components['flow_nav_btn'] = flow_nav_btn  # NEW
"""
    
    # 4. Add process flow settings
    settings_modification = """
# ADD this to the settings panel (around line 264)

show_process_flow = gr.Checkbox(
    label="Show process flow visualization",
    value=True,
    info="Display detailed LLM inference and agent execution flow"
)
interface_components['show_process_flow'] = show_process_flow
"""
    
    # 5. Enhanced chat handler
    enhanced_handler = """
# REPLACE the existing chat_handler_fn with this enhanced version

def enhanced_chat_handler_fn(message, history, session_id=None, show_reasoning=True, show_agent_trace=False, show_process_flow=True):
    \"\"\"
    Enhanced chat handler with process flow visualization
    \"\"\"
    start_time = time.time()
    
    try:
        # Use existing process_message function
        result = process_message(message, history, session_id)
        updated_history, empty_string, reasoning_data, performance_data, context_data, session_id = result
        
        # Calculate processing time
        processing_time = time.time() - start_time
        
        # Prepare process flow data if enabled
        flow_updates = {}
        if show_process_flow:
            # Extract agent results from the processing
            intent_result = {
                "primary_intent": "information_request",  # Would be extracted from actual processing
                "confidence_scores": {"information_request": 0.8},
                "secondary_intents": [],
                "reasoning_chain": ["Step 1: Analyze user input", "Step 2: Determine intent"],
                "context_tags": ["general"],
                "processing_time": 0.15,
                "agent_id": "INTENT_REC_001"
            }
            
            synthesis_result = {
                "final_response": updated_history[-1]["content"] if updated_history else "",
                "draft_response": "",
                "source_references": ["INTENT_REC_001"],
                "coherence_score": 0.85,
                "synthesis_method": "llm_enhanced",
                "intent_alignment": {"intent_detected": "information_request", "alignment_score": 0.8},
                "processing_time": processing_time - 0.15,
                "agent_id": "RESP_SYNTH_001"
            }
            
            safety_result = {
                "original_response": updated_history[-1]["content"] if updated_history else "",
                "safety_checked_response": updated_history[-1]["content"] if updated_history else "",
                "warnings": [],
                "safety_analysis": {
                    "toxicity_score": 0.1,
                    "bias_indicators": [],
                    "privacy_concerns": [],
                    "overall_safety_score": 0.9,
                    "confidence_scores": {"safety": 0.9}
                },
                "blocked": False,
                "processing_time": 0.1,
                "agent_id": "SAFETY_BIAS_001"
            }
            
            # Update process flow visualization
            flow_updates = update_process_flow_visualization(
                user_input=message,
                intent_result=intent_result,
                synthesis_result=synthesis_result,
                safety_result=safety_result,
                final_response=updated_history[-1]["content"] if updated_history else "",
                session_id=session_id,
                processing_time=processing_time
            )
        
        # Return all updates including process flow data
        return (
            updated_history,  # chatbot
            empty_string,    # message_input
            reasoning_data,  # reasoning_display
            performance_data, # performance_display
            context_data,    # context_display
            session_id,      # session_info
            flow_updates.get("flow_display", ""),  # flow_display
            flow_updates.get("flow_stats", {}),   # flow_stats
            flow_updates.get("performance_metrics", {}), # performance_metrics
            flow_updates.get("intent_details", {}), # intent_details
            flow_updates.get("synthesis_details", {}), # synthesis_details
            flow_updates.get("safety_details", {}) # safety_details
        )
        
    except Exception as e:
        logger.error(f"Error in enhanced chat handler: {e}")
        # Return error state
        error_history = list(history) if history else []
        error_history.append({"role": "user", "content": message})
        error_history.append({"role": "assistant", "content": f"Error: {str(e)}"})
        
        return (
            error_history,  # chatbot
            "",             # message_input
            {"error": str(e)}, # reasoning_display
            {"error": str(e)}, # performance_display
            {"error": str(e)}, # context_display
            session_id,     # session_info
            "",             # flow_display
            {"error": str(e)}, # flow_stats
            {"error": str(e)}, # performance_metrics
            {},             # intent_details
            {},             # synthesis_details
            {}              # safety_details
        )

# Update the chat_handler_fn assignment
chat_handler_fn = enhanced_chat_handler_fn
"""
    
    # 6. Update the send button click handler
    click_handler_modification = """
# MODIFY the send button click handler (around line 303)
# Update the outputs to include process flow components

interface_components['send_btn'].click(
    fn=chat_handler_fn,
    inputs=[
        interface_components['message_input'], 
        interface_components['chatbot'],
        interface_components['session_info'],
        interface_components.get('show_reasoning', gr.Checkbox(value=True)),
        interface_components.get('show_agent_trace', gr.Checkbox(value=False)),
        interface_components.get('show_process_flow', gr.Checkbox(value=True))
    ],
    outputs=[
        interface_components['chatbot'], 
        interface_components['message_input'],
        interface_components.get('reasoning_display', gr.JSON()),
        interface_components.get('performance_display', gr.JSON()),
        interface_components.get('context_display', gr.JSON()),
        interface_components['session_info'],
        interface_components.get('flow_display', gr.HTML()),
        interface_components.get('flow_stats', gr.JSON()),
        interface_components.get('performance_metrics', gr.JSON()),
        interface_components.get('intent_details', gr.JSON()),
        interface_components.get('synthesis_details', gr.JSON()),
        interface_components.get('safety_details', gr.JSON())
    ]
)
"""
    
    # 7. Add process flow event handlers
    event_handlers = """
# ADD these event handlers after the existing ones (around line 340)

# Process Flow event handlers
if 'clear_flow_btn' in interface_components:
    interface_components['clear_flow_btn'].click(
        fn=clear_flow_history,
        outputs=[
            interface_components.get('flow_display', gr.HTML()),
            interface_components.get('flow_stats', gr.JSON()),
            interface_components.get('performance_metrics', gr.JSON()),
            interface_components.get('intent_details', gr.JSON()),
            interface_components.get('synthesis_details', gr.JSON()),
            interface_components.get('safety_details', gr.JSON())
        ]
    )

if 'export_flow_btn' in interface_components:
    interface_components['export_flow_btn'].click(
        fn=export_flow_data,
        outputs=[gr.File(label="Download Flow Data")]
    )

if 'share_flow_btn' in interface_components:
    interface_components['share_flow_btn'].click(
        fn=lambda: gr.Info("Flow sharing feature coming soon!"),
        outputs=[]
    )
"""
    
    return {
        "import_statement": import_statement,
        "interface_modification": interface_modification,
        "navigation_modification": navigation_modification,
        "settings_modification": settings_modification,
        "enhanced_handler": enhanced_handler,
        "click_handler_modification": click_handler_modification,
        "event_handlers": event_handlers
    }

def create_integration_guide():
    """
    Create a step-by-step integration guide
    """
    modifications = modify_app_py()
    
    guide = f"""
# Process Flow Visualization Integration Guide

## Overview
This guide will help you integrate the Process Flow Visualization into your Research Assistant UI.

## Files Created
1. `process_flow_visualizer.py` - Main visualization component
2. `app_integration.py` - Integration utilities
3. `integrate_process_flow.py` - This integration guide

## Step-by-Step Integration

### Step 1: Add Import Statement
Add this import at the top of your `app.py` file:

```python
{modifications['import_statement']}
```

### Step 2: Modify Interface Creation
In your `create_mobile_optimized_interface()` function, add the Process Flow tab after the existing tabs:

```python
{modifications['interface_modification']}
```

### Step 3: Update Mobile Navigation
Modify the mobile navigation section to include the Process Flow button:

```python
{modifications['navigation_modification']}
```

### Step 4: Add Process Flow Settings
Add the process flow checkbox to your settings panel:

```python
{modifications['settings_modification']}
```

### Step 5: Replace Chat Handler
Replace your existing `chat_handler_fn` with the enhanced version:

```python
{modifications['enhanced_handler']}
```

### Step 6: Update Send Button Handler
Modify the send button click handler to include process flow outputs:

```python
{modifications['click_handler_modification']}
```

### Step 7: Add Event Handlers
Add the process flow event handlers after your existing ones:

```python
{modifications['event_handlers']}
```

## Features Added

### 🎯 Process Flow Tab
- **Visual Flow Display**: Shows step-by-step LLM inference process
- **Real-time Updates**: Updates with each user interaction
- **Mobile Optimized**: Responsive design for all devices

### 📊 Flow Statistics
- **Performance Metrics**: Processing time, confidence scores
- **Intent Distribution**: Shows intent classification patterns
- **Agent Performance**: Individual agent execution metrics

### 🔍 Detailed Analysis
- **Intent Recognition Details**: Complete intent analysis data
- **Response Synthesis Details**: Synthesis method and quality metrics
- **Safety Check Details**: Safety analysis and warnings

### 📥 Export & Share
- **Export Flow Data**: Download complete flow history as JSON
- **Share Flow**: Share flow visualizations (coming soon)

## UX Enhancements

### 🎨 Visual Design
- **Gradient Backgrounds**: Modern, professional appearance
- **Smooth Animations**: Hover effects and transitions
- **Color-coded Steps**: Different colors for different process steps
- **Progress Indicators**: Visual confidence and safety score bars

### 📱 Mobile Optimization
- **Responsive Grid**: Adapts to different screen sizes
- **Touch-friendly**: Optimized for mobile interactions
- **Collapsible Sections**: Accordion-style organization
- **Compact Mode**: Option for smaller screens

### ⚡ Performance
- **Efficient Updates**: Only updates changed components
- **Caching**: Stores flow history for analysis
- **Error Handling**: Graceful degradation on errors
- **Loading States**: Visual feedback during processing

## Testing

### Test the Integration
1. Start your Research Assistant
2. Navigate to the "🔄 Process Flow" tab
3. Send a message in the chat
4. Watch the process flow update in real-time
5. Check the statistics and detailed analysis

### Verify Features
- [ ] Process Flow tab appears
- [ ] Flow updates with each message
- [ ] Statistics show correct data
- [ ] Export functionality works
- [ ] Mobile responsive design
- [ ] Settings control visibility

## Troubleshooting

### Common Issues
1. **Import Errors**: Ensure all files are in the same directory
2. **Missing Components**: Check that all interface components are created
3. **Handler Errors**: Verify the enhanced handler is properly assigned
4. **Display Issues**: Check CSS styling and responsive design

### Debug Mode
Enable debug logging to troubleshoot issues:
```python
import logging
logging.basicConfig(level=logging.DEBUG)
```

## Support
If you encounter issues, check the logs and ensure all modifications are applied correctly.
The integration maintains backward compatibility with your existing functionality.
"""
    
    return guide

def main():
    """
    Main function to generate integration files
    """
    print("🔄 Process Flow Visualization Integration")
    print("=" * 50)
    
    # Generate integration guide
    guide = create_integration_guide()
    
    # Save the guide
    with open("INTEGRATION_GUIDE.md", "w", encoding="utf-8") as f:
        f.write(guide)
    
    print("✅ Integration files created:")
    print("   - process_flow_visualizer.py")
    print("   - app_integration.py") 
    print("   - integrate_process_flow.py")
    print("   - INTEGRATION_GUIDE.md")
    print()
    print("📖 Next steps:")
    print("   1. Follow the INTEGRATION_GUIDE.md")
    print("   2. Modify your app.py file")
    print("   3. Test the integration")
    print("   4. Enjoy the enhanced UX!")
    print()
    print("🎯 Features added:")
    print("   - Real-time process flow visualization")
    print("   - Mobile-optimized responsive design")
    print("   - Performance metrics and statistics")
    print("   - Export and sharing capabilities")
    print("   - Enhanced user experience")

if __name__ == "__main__":
    main()