Update ui/conversation.py

ui/conversation.py

@@ -1,95 +1,85 @@
-"""
-Understanding agent for pharmaceutical data management.
-This agent extracts user intent from natural language conversations.
-"""
+import streamlit as st
 
-import time
-from typing import Dict, Any
-from langgraph.graph.message import add_messages
-from anthropic.types import MessageParam
-
-def understanding_agent(anthropic_client, state: Dict[str, Any]) -> Dict[str, Any]:
-    """
-    Agent that understands user intent and requests clarification if needed.
-    
-    Args:
-        anthropic_client: The Anthropic client for calling Claude API
-        state: Current state of the agent workflow
-    
-    Returns:
-        Updated state
-    """
-    # Ensure messages is properly initialized as a list
-    messages = state.get("messages", [])
-    if not isinstance(messages, list):
-        messages = []
-    
-    # Format all messages to ensure they have role and content
-    formatted_messages = []
-    for msg in messages:
-        if isinstance(msg, dict) and "role" in msg and "content" in msg:
-            formatted_messages.append(msg)
-    
-    # Add agent-specific instructions
-    system_message = """
-    You are an AI assistant specializing in understanding pharmaceutical data needs.
-    Your job is to understand what data pipeline the user needs to create.
-    Extract key information:
-    1. Business question/goal
-    2. Required data sources/tables
-    3. Needed transformations
-    4. Output format/visualization
-    5. Time frame/frequency
-    
-    If information is missing, ask clarifying questions.
-    If you understand the request fully, summarize the user's intent and tag it as INTENT_COMPLETE.
-    """
-    
-    # Convert messages to the format expected by Anthropic API
-    anthropic_messages = []
-    for msg in formatted_messages:
-        anthropic_messages.append(MessageParam(
-            role=msg["role"],
-            content=msg["content"]
-        ))
-    
-    # Call Claude API
-    response = anthropic_client.messages.create(
-        model="claude-3-7-sonnet-20250219",
-        system=system_message,
-        messages=anthropic_messages,
-        max_tokens=2000
-    )
-    
-    # Extract the response
-    agent_response = response.content[0].text
-    
-    # Check if intent is complete
-    intent_complete = "INTENT_COMPLETE" in agent_response
+def render_conversation_tab(session_state, agent_graph, update_state_dict):
+    """Render the conversation tab in the UI."""
+    st.subheader("Conversation with Data Management Agent")
     
-    # Update state based on intent completeness
-    new_state = state.copy()
-    if intent_complete:
-        # Parse the structured intent from the response
-        # This would be more sophisticated in production
-        user_intent = {
-            "understood": True,
-            "description": agent_response.replace("INTENT_COMPLETE", "").strip(),
-            "time": time.time()
-        }
-        new_state["user_intent"] = user_intent
-        new_state["current_agent"] = "planning_agent"
-    else:
-        # Need more information, stay with understanding agent
-        new_state["current_agent"] = "understanding_agent"
+    # Initialize messages in conversation if not already present
+    if "messages" not in session_state.conversation:
+        session_state.conversation["messages"] = []
     
-    # Add agent's response to messages safely
-    if "messages" not in new_state or not isinstance(new_state["messages"], list):
-        new_state["messages"] = []
+    # Display conversation history
+    for message in session_state.conversation["messages"]:
+        if isinstance(message, dict) and "role" in message and "content" in message:
+            if message["role"] == "user":
+                st.markdown(f"**You:** {message['content']}")
+            else:
+                st.markdown(f"**Agent:** {message['content']}")
     
-    new_state["messages"].append({
-        "role": "assistant", 
-        "content": agent_response.replace("INTENT_COMPLETE", "").strip()
-    })
+    # Input for new message
+    with st.form(key="user_input_form"):
+        user_input = st.text_area("What data pipeline do you need to create?", 
+                        placeholder="e.g., I need a sales performance dashboard showing regional performance by product for the last 2 years")
+        submit_button = st.form_submit_button("Submit")
     
-    return new_state
+    if submit_button and user_input:
+        # Add user message to conversation
+        new_message = {"role": "user", "content": user_input}
+        session_state.conversation["messages"].append(new_message)
+        
+        # Update agent state - ensure messages is a list
+        agent_state = session_state.agent_state.copy()
+        if "messages" not in agent_state or not isinstance(agent_state["messages"], list):
+            agent_state["messages"] = []
+        
+        # Add new message to agent state
+        agent_state["messages"].append(new_message)
+        
+        # Run the agent graph
+        with st.spinner("Agent is processing..."):
+            try:
+                # Update the state dictionary for tools
+                update_state_dict(agent_state)
+                
+                # Execute the agent workflow
+                result = agent_graph.invoke(agent_state)
+                
+                # Update session state with result
+                session_state.agent_state = result
+                
+                # Update the state dictionary for tools again with the result
+                update_state_dict(result)
+                
+                # Ensure messages is a list in the result
+                if not isinstance(result.get("messages", []), list):
+                    result["messages"] = []
+                
+                # Update conversation with agent responses
+                for message in result.get("messages", []):
+                    if isinstance(message, dict) and "role" in message and "content" in message:
+                        # Check if this message isn't already in the conversation
+                        message_exists = False
+                        for existing_msg in session_state.conversation["messages"]:
+                            if (isinstance(existing_msg, dict) and 
+                                existing_msg.get("role") == message.get("role") and 
+                                existing_msg.get("content") == message.get("content")):
+                                message_exists = True
+                                break
+                                
+                        if not message_exists:
+                            session_state.conversation["messages"].append(message)
+                
+                # Update other state properties
+                session_state.conversation["user_intent"] = result.get("user_intent", {})
+                session_state.conversation["pipeline_plan"] = result.get("pipeline_plan", {})
+                session_state.conversation["sql_queries"] = result.get("sql_queries", [])
+                session_state.conversation["execution_results"] = result.get("execution_results", {})
+                session_state.conversation["confidence_scores"] = result.get("confidence_scores", {})
+                session_state.conversation["status"] = result.get("status", "planning")
+                session_state.conversation["current_agent"] = result.get("current_agent", "understanding_agent")
+                
+                # Force refresh
+                st.rerun()
+                
+            except Exception as e:
+                st.error(f"Error executing agent workflow: {str(e)}")