ui/agent_workflow.py

"""
Agent workflow UI component for the pharmaceutical data management agent.
"""

import streamlit as st

def render_workflow_tab(session_state):
    """
    Render the agent workflow visualization tab in the UI.
    
    Args:
        session_state: Streamlit session state
    """
    st.subheader("Agent Workflow Visualization")
    
    # Display current agent and status - ensure it's consistent
    current_agent = session_state.conversation.get("current_agent", "understanding_agent")
    status = session_state.conversation.get("status", "planning")
    
    st.markdown(f"**Current State:** {status.title()}")
    st.markdown(f"**Current Agent:** {current_agent.replace('_', ' ').title()}")
    
    # Visualize the workflow
    col1, col2, col3, col4 = st.columns(4)
    
    # Determine which agent is active
    understanding_active = current_agent == "understanding_agent"
    planning_active = current_agent == "planning_agent"
    sql_active = current_agent == "sql_generator_agent"
    executor_active = current_agent == "executor_agent"
    
    # Show the workflow visualization
    with col1:
        if understanding_active:
            st.markdown("### 🔍 **Understanding**")
        else:
            st.markdown("### 🔍 Understanding")
        st.markdown("Extracts user intent and asks clarification questions")
        
        if "user_intent" in session_state.conversation and session_state.conversation["user_intent"]:
            st.success("Completed")
        elif understanding_active:
            st.info("In Progress")
        else:
            st.warning("Not Started")
    
    with col2:
        if planning_active:
            st.markdown("### 📋 **Planning**")
        else:
            st.markdown("### 📋 Planning")
        st.markdown("Creates data pipeline plan with sources and transformations")
        
        if "pipeline_plan" in session_state.conversation and session_state.conversation["pipeline_plan"]:
            st.success("Completed")
        elif planning_active:
            st.info("In Progress")
        elif understanding_active:
            st.warning("Not Started")
        else:
            st.success("Completed")
    
    with col3:
        if sql_active:
            st.markdown("### 💻 **SQL Generation**")
        else:
            st.markdown("### 💻 SQL Generation")
        st.markdown("Converts plan into executable SQL queries")
        
        if "sql_queries" in session_state.conversation and session_state.conversation["sql_queries"]:
            st.success("Completed")
        elif sql_active:
            st.info("In Progress")
        elif understanding_active or planning_active:
            st.warning("Not Started")
        else:
            st.success("Completed")
    
    with col4:
        if executor_active:
            st.markdown("### ⚙️ **Execution**")
        else:
            st.markdown("### ⚙️ Execution")
        st.markdown("Executes queries and reports results")
        
        if "execution_results" in session_state.conversation and session_state.conversation["execution_results"]:
            st.success("Completed")
        elif executor_active:
            st.info("In Progress")
        elif understanding_active or planning_active or sql_active:
            st.warning("Not Started")
        else:
            st.success("Completed")
    
    # Overall confidence score
    if "confidence_scores" in session_state.conversation and "overall" in session_state.conversation["confidence_scores"]:
        st.markdown("### Overall Pipeline Confidence")
        score = session_state.conversation["confidence_scores"]["overall"] * 100
        st.progress(score / 100, text=f"{score:.1f}%")
        
        # Workflow decision points
        if status == "complete":
            if score > 80:
                st.success("✅ High confidence - Pipeline can be deployed automatically")
            else:
                st.warning("⚠️ Medium confidence - Human review recommended before deployment")
    
    # Add human review section for pending approval status
    if status == "pending_approval":
        st.markdown("### 👤 Human Review Required")
        st.info("This pipeline requires human review before deployment")
        
        col1, col2 = st.columns(2)
        with col1:
            if st.button("✅ Approve Pipeline"):
                # Update state to approved
                session_state.conversation["status"] = "approved"
                # Trigger execution to continue
                st.rerun()
        with col2:
            if st.button("❌ Reject Pipeline"):
                # Update state to rejected
                session_state.conversation["status"] = "rejected"
                st.error("Pipeline rejected. Please provide feedback to refine the pipeline.")