agents/understanding.py

"""
Understanding agent for pharmaceutical data management.
This agent extracts user intent from natural language conversations.
"""

import time
from typing import Dict, Any, List
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 values only (not the entire state)
    """
    # Get current messages - ensure it's a list
    messages = state.get("messages", [])
    
    # System message as a string
    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.
    """
    
    # Format messages for the Anthropic API - each must 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(MessageParam(
                role=msg["role"],
                content=msg["content"]
            ))
    
    try:
        # Call Claude API with system parameter separately (not in messages)
        response = anthropic_client.messages.create(
            model="claude-3-7-sonnet-20250219",
            system=system_message,  # System message as separate parameter
            messages=formatted_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
        
        # Clean response text
        clean_response = agent_response.replace("INTENT_COMPLETE", "").strip()
        
        # Return only the STATE UPDATES, not the entire state
        result = {}
        
        # Add a new message to the list (will be combined with existing via operator.add)
        result["messages"] = [{"role": "assistant", "content": clean_response}]
        
        # Add the user intent if complete
        if intent_complete:
            result["user_intent"] = {
                "understood": True,
                "description": clean_response,
                "time": time.time()
            }
            result["current_agent"] = "planning_agent"
        else:
            result["current_agent"] = "understanding_agent"
        
        return result
    
    except Exception as e:
        # Handle any errors - return only state updates
        print(f"Error in understanding_agent: {str(e)}")
        return {
            "messages": [{"role": "assistant", "content": f"I encountered an error: {str(e)}"}],
            "current_agent": "understanding_agent"
        }