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

# 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
@st.cache_resource
def get_s3_client():
    return boto3.client(
        "s3",
        region_name="us-east-2"
    )

@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_runs(_s3_client, bucket, batch_id=None):
    """Get all run directories for a specific batch or all batches"""
    try:
        prefix = f"run_"
        response = _s3_client.list_objects_v2(
            Bucket=bucket,
            Prefix=prefix
        )
        
        runs = []
        for obj in response.get('Contents', []):
            if batch_id is None or batch_id in obj['Key']:
                run_dir = '/'.join(obj['Key'].split('/')[:-1])
                if run_dir and run_dir not in runs:
                    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_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.error(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_prefix):
    """Get all files in a run directory"""
    try:
        response = s3_client.list_objects_v2(
            Bucket=bucket,
            Prefix=run_prefix
        )
        
        files = {
            'individual': [],
            'consolidated': [],
            'logs': []
        }
        
        for obj in response.get('Contents', []):
            key = obj['Key']
            if 'individual/' in key:
                files['individual'].append(key)
            elif 'consolidated/' in key:
                files['consolidated'].append(key)
            elif key.endswith('.csv'):  # Log files in root of run directory
                files['logs'].append(key)
                
        return files
    except Exception as e:
        st.error(f"Error listing files: {str(e)}")
        return None


def get_active_runs(df):
    """
    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.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)
            total_runs = len(df_filtered)
            completed_runs = len(df_filtered[df_filtered['status'] == 'Completed'])
            success_rate = (completed_runs / total_runs * 100) if total_runs > 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()
                display_df['progress'] = (display_df['files_processed'] / display_df['total_input_files'] * 100).round(2)
                display_df['progress'] = display_df['progress'].astype(str) + '%'
                st.dataframe(
                    display_df[['batch_id', 'start_time', 'input_dir', 'progress']],
                    use_container_width=True
                )
        else:
            st.warning("No data available for the selected filters.")
        
        # Graphs row
        col1, col2 = st.columns(2)
        
        with col1:
            st.subheader("Processing Speed Over Time")
            fig = px.line(
                df_filtered.sort_values('start_time'),
                x='start_time',
                y='contracts_per_hour',
                title="Processing Speed Trend",
                line_shape='linear'
            )
            st.plotly_chart(fig, use_container_width=True)
        
        with col2:
            st.subheader("Cumulative Files Processed")
            df_cumulative = df_filtered.sort_values('start_time').copy()
            df_cumulative['cumulative_files'] = df_cumulative['files_processed'].cumsum()
            
            fig = px.line(
                df_cumulative,
                x='start_time',
                y='cumulative_files',
                title="Cumulative Files Processed",
                line_shape='linear'
            )
            st.plotly_chart(fig, use_container_width=True)

        # Cost Analysis section
        st.subheader("Cost Analysis")
        selected_batch = st.selectbox(
            "Select Batch ID",
            options=df_filtered['batch_id'].unique(),
            key="cost_analysis_batch"
        )
        
        if selected_batch:
            runs = get_batch_runs(s3_client, bucket, selected_batch)
            if runs:
                selected_run = st.selectbox("Select Run", runs)
                
                # Load cost log
                cost_log_key = f"{selected_run}/costlog.csv"
                df_cost = load_csv_from_s3(s3_client, bucket, cost_log_key)
                
                if df_cost is not None:
                    col1, col2 = st.columns(2)
                    
                    with col1:
                        st.subheader("Cost Distribution by Function")
                        fig = px.pie(
                            df_cost,
                            values='Total Cost',
                            names='Caller Function',
                            title="Cost Distribution by Function"
                        )
                        st.plotly_chart(fig, use_container_width=True)
                    
                    with col2:
                        st.subheader("Token Usage")
                        fig = px.bar(
                            df_cost,
                            x='Caller Function',
                            y=['Input Prompt Tokens', 'Output Response Tokens'],
                            title="Token Usage by Function",
                            barmode='group'
                        )
                        st.plotly_chart(fig, use_container_width=True)
                    
                    st.subheader("Cost Details")
                    st.dataframe(df_cost, use_container_width=True)

    
    with tab2:
        st.header("Batch Run Explorer")
        
        # Batch selection
        selected_batch_explore = st.selectbox(
            "Select Batch ID to Explore",
            options=df_tracker['batch_id'].unique(),
            key="explorer_batch"
        )
        
        if selected_batch_explore:
            runs = get_batch_runs(s3_client, bucket, selected_batch_explore)
            
            if runs:
                selected_run = st.selectbox("Select Run", runs, key="explorer_run")
                
                # Get all files in the run
                files = get_run_files(s3_client, bucket, selected_run)
                
                if files:
                    # Display file sections with download buttons
                    col1, col2, col3 = st.columns(3)
                    
                    with col1:
                        st.subheader("Individual Files")
                        for key in files['individual']:
                            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"
                                )
                    
                    with col2:
                        st.subheader("Consolidated Files")
                        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"
                                )
                    
                    with col3:
                        st.subheader("Log Files")
                        for key in files['logs']:
                            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"
                                )
                    
                    # Display log file contents
                    st.subheader("Log File Analysis")
                    
                    # File Log Analysis
                    file_log_key = next((k for k in files['logs'] if 'filelog' in k), None)
                    if file_log_key:
                        st.subheader("File Processing Log")
                        df_filelog = load_csv_from_s3(s3_client, bucket, file_log_key)
                        if df_filelog is not None:
                            col1, col2 = st.columns(2)
                            
                            with col1:
                                st.subheader("Processing Time by Status")
                                fig = px.box(
                                    df_filelog,
                                    x='status',
                                    y='processing_time_seconds',
                                    title="Processing Time Distribution by Status"
                                )
                                st.plotly_chart(fig, use_container_width=True)
                            
                            with col2:
                                st.subheader("Resource Usage")
                                fig = px.scatter(
                                    df_filelog,
                                    x='memory_usage_mb',
                                    y='cpu_percent',
                                    color='status',
                                    title="Resource Usage by File"
                                )
                                st.plotly_chart(fig, use_container_width=True)
                            
                            # Show raw data with filters
                            st.subheader("Raw File Log Data")
                            status_filter = st.multiselect(
                                "Filter by Status",
                                options=df_filelog['status'].unique()
                            )
                            if status_filter:
                                df_filelog = df_filelog[df_filelog['status'].isin(status_filter)]
                            st.dataframe(df_filelog, use_container_width=True)
                    
                    # Cost Log Analysis
                    cost_log_key = next((k for k in files['logs'] if 'costlog' in k), None)
                    if cost_log_key:
                        st.subheader("Cost Analysis")
                        df_cost = load_csv_from_s3(s3_client, bucket, cost_log_key)
                        if df_cost is not None:
                            col1, col2 = st.columns(2)
                            
                            with col1:
                                total_cost = df_cost['Total Cost'].sum()
                                avg_cost = df_cost['Total Cost'].mean()
                                st.metric("Total Cost", f"${total_cost:.2f}")
                                st.metric("Average Cost per File", f"${avg_cost:.4f}")
                            
                            with col2:
                                total_tokens = df_cost['Input Prompt Tokens'].sum() + df_cost['Output Response Tokens'].sum()
                                avg_tokens = total_tokens / len(df_cost)
                                st.metric("Total Tokens", f"{total_tokens:,}")
                                st.metric("Average Tokens per File", f"{avg_tokens:,.0f}")

if __name__ == "__main__":
    main()