doczyai-pipelines/streamlit/runs-dashboard.py

import streamlit as st
import pandas as pd
import boto3
import plotly.express as px
import plotly.graph_objects as go
from datetime import datetime, timedelta
import json
from io import StringIO
from dotenv import load_dotenv
import os
import re

# Load environment variables
load_dotenv()

# Configure page
st.set_page_config(
    page_title="Doczy.AI Batch Processing Dashboard",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Initialize S3 client with fallback to local credentials
@st.cache_resource
def get_s3_client():
    try:
        # First try direct connection
        return boto3.client(
            "s3",
            region_name="us-east-2"
        )
    except Exception as e:
        st.warning("Direct AWS connection failed, attempting to use local credentials...")
        
        try:
            # Load environment variables
            load_dotenv()
            
            # Get temporary credentials from environment variables
            aws_access_key = os.getenv('AWS_ACCESS_KEY_ID')
            aws_secret_key = os.getenv('AWS_SECRET_ACCESS_KEY')
            aws_session_token = os.getenv('AWS_SESSION_TOKEN')
            
            if not all([aws_access_key, aws_secret_key, aws_session_token]):
                raise ValueError("Missing required AWS credentials in .env file")
            
            # Create client with temporary credentials
            return boto3.client(
                "s3",
                region_name="us-east-2",
                aws_access_key_id=aws_access_key,
                aws_secret_access_key=aws_secret_key,
                aws_session_token=aws_session_token
            )
        except Exception as inner_e:
            st.error(f"Failed to connect using local credentials: {str(inner_e)}")
            st.error("Please ensure .env file exists with AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, and AWS_SESSION_TOKEN")
            raise

@st.cache_data
def load_batch_tracker(_s3_client, bucket):
    """Load the main batch tracker file from S3"""
    try:
        response = _s3_client.get_object(
            Bucket=bucket,
            Key='batch-tracker/doczyai-batch-tracker.csv'
        )
        return pd.read_csv(response['Body'])
    except Exception as e:
        st.error(f"Error loading batch tracker: {str(e)}")
        return None

@st.cache_data
def get_batch_folders(_s3_client, bucket):
    """Get all batch folders in the bucket"""
    try:
        response = _s3_client.list_objects_v2(
            Bucket=bucket,
            Delimiter='/'
        )
        
        batch_folders = []
        for prefix in response.get('CommonPrefixes', []):
            folder_name = prefix.get('Prefix', '').rstrip('/')
            if folder_name and not folder_name.startswith('batch-tracker'):
                batch_folders.append(folder_name)
        
        return sorted(batch_folders)
    except Exception as e:
        st.error(f"Error listing batch folders: {str(e)}")
        return []

@st.cache_data
def get_batch_runs(_s3_client, bucket, batch_id):
    """Get all run directories for a specific batch"""
    try:
        prefix = f"{batch_id}/run_"
        response = _s3_client.list_objects_v2(
            Bucket=bucket,
            Prefix=prefix,
            Delimiter='/'
        )
        
        runs = []
        for obj in response.get('CommonPrefixes', []):
            run_dir = obj.get('Prefix', '').rstrip('/')
            if run_dir:
                runs.append(run_dir)
        
        return sorted(runs, reverse=True)
    except Exception as e:
        st.error(f"Error listing runs: {str(e)}")
        return []

@st.cache_data
def load_master_batch_tracking(_s3_client, bucket, batch_id):
    """Load the master batch tracking CSV for a specific batch"""
    try:
        response = _s3_client.get_object(
            Bucket=bucket,
            Key=f"{batch_id}/master_batch_tracking.csv"
        )
        return pd.read_csv(response['Body'])
    except Exception as e:
        st.warning(f"No master batch tracking found for {batch_id}: {str(e)}")
        return None

@st.cache_data
def load_csv_from_s3(_s3_client, bucket, key):
    """Load any CSV file from S3"""
    try:
        response = _s3_client.get_object(
            Bucket=bucket,
            Key=key
        )
        return pd.read_csv(response['Body'])
    except Exception as e:
        st.warning(f"Error loading CSV {key}: {str(e)}")
        return None

def create_download_link(_s3_client, bucket, key):
    """Create a presigned URL for downloading the file"""
    try:
        url = _s3_client.generate_presigned_url(
            'get_object',
            Params={
                'Bucket': bucket,
                'Key': key
            },
            ExpiresIn=3600
        )
        return url
    except Exception as e:
        st.error(f"Error creating download link: {str(e)}")
        return None

def get_run_files(s3_client, bucket, run_path):
    """Get all files in a run directory with the new structure"""
    try:
        # Check for individual files
        individual_prefix = f"{run_path}/individual/"
        individual_response = s3_client.list_objects_v2(
            Bucket=bucket,
            Prefix=individual_prefix
        )
        
        # Check for consolidated files
        consolidated_prefix = f"{run_path}/consolidated/"
        consolidated_response = s3_client.list_objects_v2(
            Bucket=bucket,
            Prefix=consolidated_prefix
        )
        
        # Check for tracking files
        tracking_prefix = f"{run_path}/tracking/"
        tracking_response = s3_client.list_objects_v2(
            Bucket=bucket,
            Prefix=tracking_prefix
        )
        
        files = {
            'individual': [],
            'consolidated': [],
            'tracking': []
        }
        
        # Process individual files - organize by subfolder
        individual_folders = {}
        for obj in individual_response.get('Contents', []):
            key = obj['Key']
            parts = key.split('/')
            if len(parts) > 3:  # batch_id/run_id/individual/folder/file
                folder = parts[-2]
                if folder not in individual_folders:
                    individual_folders[folder] = []
                individual_folders[folder].append(key)
        
        files['individual'] = individual_folders
        
        # Process consolidated files
        for obj in consolidated_response.get('Contents', []):
            files['consolidated'].append(obj['Key'])
                
        # Process tracking files
        for obj in tracking_response.get('Contents', []):
            files['tracking'].append(obj['Key'])
                
        return files
    except Exception as e:
        st.error(f"Error listing files: {str(e)}")
        return None

def extract_timestamp_from_run(run_path):
    """Improved timestamp extraction from run path"""
    # Match patterns like 'run_20250110_09:02' or 'run_20250110_09:02_...'
    match = re.search(r'run_(\d{8}_\d{2}:\d{2})', run_path)
    if match:
        timestamp_str = match.group(1)
        try:
            return pd.to_datetime(timestamp_str, format='%Y%m%d_%H:%M')
        except:
            return None
    return None



def calculate_processing_rate(master_df, run_path):
    """Calculate processing rate based on file timestamps and run start time"""
    if master_df is None or master_df.empty:
        return None
    
    run_start_time = extract_timestamp_from_run(run_path)
    if run_start_time is None:
        return None
    
    # Convert the timestamp to datetime
    master_df['upload_time'] = pd.to_datetime(master_df['upload_time']) if 'upload_time' in master_df.columns else run_start_time
    
    # Calculate time difference in hours
    if 'upload_time' in master_df.columns:
        time_diffs = (master_df['upload_time'] - run_start_time).dt.total_seconds() / 3600
        time_diffs = time_diffs[time_diffs > 0]  # Only consider positive time differences
        
        if len(time_diffs) == 0:
            return 0
            
        # Calculate files per hour
        max_time_diff = time_diffs.max()
        if max_time_diff > 0:
            files_per_hour = len(time_diffs) / max_time_diff
            return files_per_hour
    
    # If no upload_time column or all times are negative, use count and rough estimate
    time_diff = (datetime.now() - run_start_time).total_seconds() / 3600
    if time_diff > 0:
        return len(master_df) / time_diff
    
    return 0

def get_active_runs(df_tracker):
    """
    Determines truly active runs by checking if there are any subsequent 'Completed' 
    entries for the same batch_id after a 'Running' status, only showing runs from today
    """
    active_runs = []
    today = pd.Timestamp.now().date()
    
    # Sort by batch_id and start_time to get chronological order
    df_sorted = df_tracker.sort_values(['batch_id', 'start_time'])
    
    # Group by batch_id to analyze status progression
    for batch_id, group in df_sorted.groupby('batch_id'):
        latest_record = group.iloc[-1]
        if (latest_record['status'] == 'Running' and 
            pd.to_datetime(latest_record['start_time']).date() == today):
            active_runs.append({
                'batch_id': batch_id,
                'start_time': latest_record['start_time'],
                'input_dir': latest_record['input_dir'],
                'files_processed': latest_record['files_processed'],
                'total_input_files': latest_record['total_input_files']
            })
    
    return pd.DataFrame(active_runs) if active_runs else pd.DataFrame()

def main():
    st.title("Doczy.AI Batch Processing Dashboard")
    
    # S3 Configuration
    s3_client = get_s3_client()
    bucket = "doczy-output"
    
    # Load batch tracker data
    df_tracker = load_batch_tracker(s3_client, bucket)
    
    if df_tracker is None:
        st.error("Unable to load batch data. Please check S3 connection.")
        return
    
    # Convert datetime columns
    df_tracker['start_time'] = pd.to_datetime(df_tracker['start_time'])
    df_tracker['end_time'] = pd.to_datetime(df_tracker['end_time'])
    
    # Create a copy for filtering
    df_filtered = df_tracker.copy()
    
    # Add tabs for different views
    tab1, tab2 = st.tabs(["Overview Dashboard", "Batch Explorer"])
    
    with tab1:
        # Sidebar filters
        st.sidebar.header("Filters")
        
        # Batch filter
        batch_filter = st.sidebar.selectbox(
            "Filter by Batch ID",
            options=['All'] + sorted(df_tracker['batch_id'].unique().tolist())
        )
        
        # Apply batch filter if selected
        if batch_filter != 'All':
            df_filtered = df_filtered[df_filtered['batch_id'] == batch_filter]
        
        # Date range filter
        min_date = df_filtered['start_time'].min().date()
        max_date = df_filtered['start_time'].max().date()
        
        date_range = st.sidebar.date_input(
            "Date Range",
            value=(min_date, max_date),
            min_value=min_date,
            max_value=max_date
        )
        
        # Apply date filter
        if len(date_range) == 2:
            start_date, end_date = date_range
            mask = (df_filtered['start_time'].dt.date >= start_date) & (df_filtered['start_time'].dt.date <= end_date)
            df_filtered = df_filtered[mask]
        
        # Status filter
        status_options = ['All'] + sorted(df_filtered['status'].unique().tolist())
        selected_status = st.sidebar.selectbox("Status", status_options)
        
        if selected_status != 'All':
            df_filtered = df_filtered[df_filtered['status'] == selected_status]
        
        # Calculate metrics based on filtered data
        if not df_filtered.empty:
            # Calculate active runs
            active_runs_df = get_active_runs(df_filtered)
            
            # Top metrics row
            col1, col2, col3, col4 = st.columns(4)
            
            # Active runs metric
            with col1:
                st.metric("Active Runs", len(active_runs_df))
            
            # Total files processed
            total_files = df_filtered['files_processed'].sum()
            with col2:
                st.metric("Total Files Processed", f"{total_files:,}")
            
            # Average processing speed (only for rows with positive values)
            speed_data = df_filtered[df_filtered['contracts_per_hour'] > 0]
            avg_speed = speed_data['contracts_per_hour'].mean() if not speed_data.empty else 0
            with col3:
                st.metric("Avg Processing Speed", f"{avg_speed:.2f} files/hour")
            
            # Success rate (considering only completed or failed status)
            completed_runs = len(df_filtered[df_filtered['status'] == 'Completed'])
            failed_runs = len(df_filtered[df_filtered['status'] == 'Failed'])
            total_finished = completed_runs + failed_runs
            success_rate = (completed_runs / total_finished * 100) if total_finished > 0 else 0
            with col4:
                st.metric("Success Rate", f"{success_rate:.1f}%" if success_rate > 0 else "-")
            
            # Display current active runs if any exist
            if not active_runs_df.empty:
                st.subheader("Currently Active Runs")
                display_df = active_runs_df.copy()
                
                st.dataframe(
                    display_df[['batch_id', 'start_time', 'input_dir']],
                    use_container_width=True
                )
        else:
            st.warning("No data available for the selected filters.")

    with tab2:
        st.header("Batch Run Explorer")
        
        # Get all batch folders
        batch_folders = get_batch_folders(s3_client, bucket)
        
        if not batch_folders:
            st.warning("No batch folders found in S3 bucket.")
            return
            
        # Batch selection
        selected_batch_explore = st.selectbox(
            "Select Batch ID to Explore",
            options=batch_folders
        )
        
        if selected_batch_explore:
            # Load master batch tracking data
            master_df = load_master_batch_tracking(s3_client, bucket, selected_batch_explore)
            
            # Get batch details from tracker
            batch_info = df_tracker[df_tracker['batch_id'] == selected_batch_explore].iloc[0] if not df_tracker[df_tracker['batch_id'] == selected_batch_explore].empty else None
            
            # Display batch summary metrics
            if batch_info is not None and master_df is not None:
                st.subheader("Batch Summary")
                
                # Calculate metrics for B and AC output types
                b_files = master_df[master_df['output_type'] == 'B'].shape[0] if not master_df.empty else 0
                ac_files = master_df[master_df['output_type'] == 'AC'].shape[0] if not master_df.empty else 0
                
                total_input = batch_info['total_input_files']
                total_processed = b_files + ac_files
                remaining_b = total_input - b_files
                remaining_ac = total_input - ac_files
                
                # Display metrics
                col1, col2, col3, col4 = st.columns(4)
                with col1:
                    st.metric("Total Input Files", f"{total_input:,}")
                    st.metric("Total Processed", f"{total_processed:,}")
                
                with col2:
                    st.metric("B Files Processed", f"{b_files:,}")
                    st.metric("B Files Remaining", f"{remaining_b:,}")
                
                with col3:
                    st.metric("AC Files Processed", f"{ac_files:,}")
                    st.metric("AC Files Remaining", f"{remaining_ac:,}")
                
                with col4:
                    if batch_info['status'] == 'Running':
                        elapsed_time = (datetime.now() - batch_info['start_time']).total_seconds() / 3600
                        current_rate = total_processed / elapsed_time if elapsed_time > 0 else 0
                        st.metric("Current Rate", f"{current_rate:.2f} files/hour")
                        
                        if current_rate > 0:
                            est_completion_b = remaining_b / current_rate
                            est_completion_ac = remaining_ac / current_rate
                            st.metric("Est. Time to Complete B", f"{est_completion_b:.1f} hours")
                            st.metric("Est. Time to Complete AC(if this was a B only run, the rate for B will be applied here. So it is not reflective)", f"{est_completion_ac:.1f} hours")
            
            # Get all runs for the selected batch
            runs = get_batch_runs(s3_client, bucket, selected_batch_explore)
            
            if runs:
                selected_run = st.selectbox("Select Run", runs)
                
                # If master_df exists, calculate processing rate for this run
                if master_df is not None and not master_df.empty:
                    processing_rate = calculate_processing_rate(master_df, selected_run)
                    if processing_rate is not None:
                        st.metric("Processing Rate", f"{processing_rate:.2f} files/hour")
                
                # Get all files in the run with new structure
                files = get_run_files(s3_client, bucket, selected_run)
                
                if files:
                    # Create tabs for different file types
                    file_tabs = st.tabs(["Individual Files", "Consolidated Files", "Tracking Files"])
                    
                    # Individual files tab - now organized by folder
                    with file_tabs[0]:
                        if files['individual']:
                            for folder, file_keys in files['individual'].items():
                                with st.expander(f"{folder} ({len(file_keys)} files)"):
                                    for key in file_keys:
                                        filename = key.split('/')[-1]
                                        download_url = create_download_link(s3_client, bucket, key)
                                        if download_url:
                                            st.download_button(
                                                f"📥 {filename}",
                                                download_url,
                                                filename,
                                                mime="text/csv",
                                                key=f"individual_{folder}_{filename}"
                                            )
                        else:
                            st.info("No individual files found for this run.")
                    
                    # Consolidated files tab
                    with file_tabs[1]:
                        if files['consolidated']:
                            for key in files['consolidated']:
                                filename = key.split('/')[-1]
                                download_url = create_download_link(s3_client, bucket, key)
                                if download_url:
                                    st.download_button(
                                        f"📥 {filename}",
                                        download_url,
                                        filename,
                                        mime="text/csv",
                                        key=f"consolidated_{filename}"
                                    )
                        else:
                            st.info("No consolidated files found for this run.")
                    
                    # Tracking files tab
                    with file_tabs[2]:
                        if files['tracking']:
                            for key in files['tracking']:
                                filename = key.split('/')[-1]
                                download_url = create_download_link(s3_client, bucket, key)
                                if download_url:
                                    st.download_button(
                                        f"📥 {filename}",
                                        download_url,
                                        filename,
                                        mime="text/csv",
                                        key=f"tracking_{filename}"
                                    )
                                
                                # Simplified cost analysis display
                                if 'costlog' in key.lower():
                                    st.subheader("Cost Analysis")
                                    df_cost = load_csv_from_s3(s3_client, bucket, key)
                                    if df_cost is not None:
                                        # Display raw data in a compact format
                                        st.dataframe(df_cost, use_container_width=True)
                                        
                                        # Basic metrics calculation
                                        if 'Total Cost' in df_cost.columns:
                                            total_cost = df_cost['Total Cost'].sum()
                                            st.metric("Total Cost", f"${total_cost:.2f}")
                
                # Display master batch tracking data
                if master_df is not None and not master_df.empty:
                    st.subheader("Master Batch Tracking")
                    
                    # Add filtering options
                    output_type_filter = st.multiselect(
                        "Filter by Output Type", 
                        options=master_df['output_type'].unique(),
                        default=master_df['output_type'].unique()
                    )
                    
                    upload_status_filter = st.multiselect(
                        "Filter by Upload Status",
                        options=master_df['upload_status'].unique(),
                        default=master_df['upload_status'].unique()
                    )
                    
                    # Apply filters
                    filtered_master_df = master_df
                    if output_type_filter:
                        filtered_master_df = filtered_master_df[filtered_master_df['output_type'].isin(output_type_filter)]
                    
                    if upload_status_filter:
                        filtered_master_df = filtered_master_df[filtered_master_df['upload_status'].isin(upload_status_filter)]
                    
                    # Show data
                    st.dataframe(filtered_master_df, use_container_width=True)
                    
                    # Create visualizations for processing status
                    st.subheader("Processing Status")
                    
                    col1, col2 = st.columns(2)
                    
                    with col1:
                        # Distribution by output type
                        output_counts = master_df['output_type'].value_counts().reset_index()
                        output_counts.columns = ['Output Type', 'Count']
                        
                        fig = px.pie(
                            output_counts,
                            values='Count',
                            names='Output Type',
                            title="Files by Output Type"
                        )
                        st.plotly_chart(fig, use_container_width=True)
                    
                    with col2:
                        # Distribution by upload status
                        status_counts = master_df['upload_status'].value_counts().reset_index()
                        status_counts.columns = ['Upload Status', 'Count']
                        
                        fig = px.pie(
                            status_counts,
                            values='Count',
                            names='Upload Status',
                            title="Files by Upload Status"
                        )
                        st.plotly_chart(fig, use_container_width=True)
                    
                    # If there's a run timestamp, calculate hourly throughput
                    
                    if 'run_timestamp' in master_df.columns:
                        # First extract proper timestamps
                        master_df['extracted_timestamp'] = master_df['run_timestamp'].apply(extract_timestamp_from_run)
                        
                        # Then filter out invalid timestamps
                        master_df = master_df.dropna(subset=['extracted_timestamp'])
                        
                        # Use extracted timestamps for analysis
                        hourly_counts = master_df.resample('H', on='extracted_timestamp').size()
                        
                        if not hourly_counts.empty:
                            st.subheader("Hourly Processing Rate")
                            fig = px.bar(
                                x=hourly_counts.index,
                                y=hourly_counts.values,
                                labels={'x': 'Hour', 'y': 'Files Processed'},
                                title="Files Processed by Hour"
                            )
                            st.plotly_chart(fig, use_container_width=True)

if __name__ == "__main__":
    main()