Create ui/agent_workflow.py

ui/agent_workflow.py

@@ -0,0 +1,124 @@
+"""
+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
+    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.")