Update app.py

app.py

@@ -179,4 +179,267 @@ def render_workflow_tab(session_state):
     executor_active = current_agent == "executor_agent"
     
     # Show the workflow visualization
-    with col1:
+    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.")
+""")
+
+    # DB Explorer UI
+    db_explorer_py = "ui/db_explorer.py"
+    if not os.path.exists(db_explorer_py):
+        with open(db_explorer_py, "w") as f:
+            f.write("""
+import streamlit as st
+import pandas as pd
+
+def render_db_explorer_tab(session_state):
+    \"\"\"Render the database explorer tab in the UI.\"\"\"
+    st.subheader("Database Explorer")
+    
+    # Get tables by category
+    tables = session_state.db.get_tables()
+    
+    # Display tables by category
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        st.markdown("### Raw Data Tables")
+        for table in tables["raw_tables"]:
+            with st.expander(table):
+                sample = session_state.db.get_table_sample(table, 3)
+                st.dataframe(pd.DataFrame(sample))
+        
+        st.markdown("### Staging Tables")
+        for table in tables["staging_tables"]:
+            with st.expander(table):
+                sample = session_state.db.get_table_sample(table, 3)
+                st.dataframe(pd.DataFrame(sample))
+    
+    with col2:
+        st.markdown("### Analytics Ready Data")
+        for table in tables["ard_tables"]:
+            with st.expander(table):
+                sample = session_state.db.get_table_sample(table, 3)
+                st.dataframe(pd.DataFrame(sample))
+        
+        st.markdown("### Data Products")
+        for table in tables["data_products"]:
+            with st.expander(table):
+                sample = session_state.db.get_table_sample(table, 3)
+                st.dataframe(pd.DataFrame(sample))
+    
+    # SQL Query Executor
+    st.markdown("### Query Explorer")
+    with st.form(key="sql_form"):
+        sql_query = st.text_area("Enter SQL Query", height=100, 
+                                placeholder="SELECT * FROM ARD_SALES_PERFORMANCE WHERE region = 'North' LIMIT 5")
+        run_sql = st.form_submit_button("Run Query")
+    
+    if run_sql and sql_query:
+        with st.spinner("Executing query..."):
+            result = session_state.db.execute_query(sql_query)
+            
+            if "error" in result:
+                st.error(f"Error executing query: {result['error']}")
+            elif "data" in result:
+                st.dataframe(pd.DataFrame(result["data"]))
+                st.success(f"Query returned {len(result['data'])} rows")
+            elif "tables" in result:
+                st.write(result["tables"])
+            elif "schema" in result:
+                st.write(f"Schema for {result['table']}:")
+                st.dataframe(pd.DataFrame(result["schema"]))
+""")
+
+# Now import UI modules after ensuring they exist
+ensure_ui_files_exist()
+
+# Import UI modules
+from ui.conversation import render_conversation_tab
+from ui.pipeline import render_pipeline_tab
+from ui.agent_workflow import render_workflow_tab
+from ui.db_explorer import render_db_explorer_tab
+
+# Import agent state
+from agents.state import AgentState
+
+# Load environment variables from .env file if it exists
+load_dotenv()
+
+def main():
+    """Main application entry point."""
+    # Set page configuration
+    st.set_page_config(page_title="Pharmaceutical Data Agent", page_icon="💊", layout="wide")
+    
+    # Sidebar for configuration
+    st.sidebar.title("Configuration")
+    
+    # API key input (in a real app, use more secure methods)
+    api_key = st.sidebar.text_input(
+        "Anthropic API Key", 
+        value=os.getenv("ANTHROPIC_API_KEY", ""),
+        type="password"
+    )
+    
+    if not api_key:
+        st.warning("Please enter your Anthropic API Key in the sidebar to continue.")
+        return
+    
+    # Set the API key as an environment variable
+    os.environ["ANTHROPIC_API_KEY"] = api_key
+    
+    # Initialize Anthropic client using default settings
+    anthropic_client = Anthropic()  # Uses API key from environment variable
+    
+    # Initialize database with synthetic data
+    if "db" not in st.session_state:
+        st.session_state.db = SyntheticDatabase()
+    
+    # Initialize agent graph
+    if "agent_graph" not in st.session_state or "update_state_dict" not in st.session_state:
+        st.session_state.agent_graph, st.session_state.update_state_dict = create_agent_graph(
+            anthropic_client, 
+            st.session_state.db
+        )
+    
+    # Initialize conversation state
+    if "conversation" not in st.session_state:
+        st.session_state.conversation = {
+            "messages": [],
+            "user_intent": {},
+            "data_context": {},
+            "pipeline_plan": {},
+            "sql_queries": [],
+            "execution_results": {},
+            "confidence_scores": {},
+            "status": "planning",
+            "current_agent": "understanding_agent"
+        }
+    
+    # Initialize agent state for new runs
+    if "agent_state" not in st.session_state:
+        st.session_state.agent_state = AgentState(
+            messages=[],
+            user_intent={},
+            data_context={},
+            pipeline_plan={},
+            sql_queries=[],
+            execution_results={},
+            confidence_scores={},
+            status="planning",
+            current_agent="understanding_agent"
+        )
+    
+    # Main content area
+    st.title("Pharmaceutical Data Management Agent")
+    
+    # Tabs for different views
+    tab1, tab2, tab3, tab4 = st.tabs(["Conversation", "Pipeline Details", "Agent Workflow", "Database Explorer"])
+    
+    with tab1:
+        render_conversation_tab(
+            st.session_state, 
+            st.session_state.agent_graph,
+            st.session_state.update_state_dict
+        )
+    
+    with tab2:
+        render_pipeline_tab(st.session_state)
+    
+    with tab3:
+        render_workflow_tab(st.session_state)
+    
+    with tab4:
+        render_db_explorer_tab(st.session_state)
+
+if __name__ == "__main__":
+    main()