Update agents/planning.py

agents/planning.py

@@ -5,11 +5,10 @@ This agent creates a data pipeline plan based on user intent.
 
 import time
 from typing import Dict, Any
-from agents.state import AgentState
-from anthropic import Anthropic
+from langgraph.graph.message import add_messages
 from anthropic.types import MessageParam
 
-def planning_agent(anthropic_client: Anthropic, state: AgentState) -> AgentState:
+def planning_agent(anthropic_client, state: Dict[str, Any]) -> Dict[str, Any]:
     """
     Agent that creates a data pipeline plan based on user intent.
     
@@ -22,7 +21,6 @@ def planning_agent(anthropic_client: Anthropic, state: AgentState) -> AgentState
     """
     # Get current messages and user intent
     messages = state.get("messages", [])
-    current_agent = state.get("current_agent", ["planning_agent"])
     user_intent = state.get("user_intent", {})
     
     # Add agent-specific instructions
@@ -52,66 +50,44 @@ def planning_agent(anthropic_client: Anthropic, state: AgentState) -> AgentState
     - Data products: DP_SALES_DASHBOARD, DP_HCP_TARGETING
     """
     
-    # Convert messages to the format expected by Anthropic API
-    prompt_messages = []
-    for msg in messages:
-        prompt_messages.append(MessageParam(
-            role=msg["role"],
-            content=msg["content"]
-        ))
+    # Prepare prompt for Claude
+    prompt_messages = [
+        MessageParam(role="system", content=system_message),
+        *[MessageParam(role=m["role"], content=m["content"]) for m in messages],
+        MessageParam(role="user", content=f"Based on my request, create a data pipeline plan. Here is information about available data: {context}")
+    ]
     
-    # Add final user message with context
-    prompt_messages.append(
-        MessageParam(
-            role="user", 
-            content=f"Based on my request, create a data pipeline plan. Here is information about available data: {context}"
-        )
+    # Call Claude API
+    response = anthropic_client.messages.create(
+        model="claude-3-7-sonnet-20250219",
+        messages=prompt_messages,
+        max_tokens=3000
     )
     
-    try:
-        # Call Claude API
-        response = anthropic_client.messages.create(
-            model="claude-3-7-sonnet-20250219",
-            system=system_message,
-            messages=prompt_messages,
-            max_tokens=3000
-        )
-        
-        # Extract the response
-        agent_response = response.content[0].text
-        
-        # Check if plan is complete
-        plan_complete = "PLAN_COMPLETE" in agent_response
-        
-        # Prepare the new state updates
-        new_state = state.copy()
-        
-        # Add agent's response to messages
-        new_state["messages"].append({
-            "role": "assistant", 
-            "content": agent_response.replace("PLAN_COMPLETE", "").strip()
-        })
-        
-        # Add the pipeline plan if complete
-        if plan_complete:
-            new_state["pipeline_plan"] = {
-                "description": agent_response.replace("PLAN_COMPLETE", "").strip(),
-                "created_at": time.time()
-            }
-            # Use .append() to add next agent to the list
-            new_state["current_agent"].append("sql_generator_agent")
-        else:
-            # Use .append() to add current agent if plan is not complete
-            new_state["current_agent"].append("planning_agent")
-        
-        return new_state
+    # Extract the response
+    agent_response = response.content[0].text
     
-    except Exception as e:
-        # Handle any errors
-        print(f"Error in planning_agent: {str(e)}")
-        return {
-            "messages": state.get("messages", []) + [
-                {"role": "assistant", "content": f"I encountered an error: {str(e)}"}
-            ],
-            "current_agent": current_agent + ["planning_agent"]
-        }
+    # Check if plan is complete
+    plan_complete = "PLAN_COMPLETE" in agent_response
+    
+    # Update state based on plan completeness
+    new_state = state.copy()
+    if plan_complete:
+        # Parse the plan from the response
+        pipeline_plan = {
+            "description": agent_response.replace("PLAN_COMPLETE", "").strip(),
+            "created_at": time.time()
+        }
+        new_state["pipeline_plan"] = pipeline_plan
+        new_state["current_agent"] = "sql_generator_agent"
+    else:
+        # Need more information or planning, stay with planning agent
+        new_state["current_agent"] = "planning_agent"
+    
+    # Add agent's response to messages
+    new_messages = add_messages(state, [
+        {"role": "assistant", "content": agent_response.replace("PLAN_COMPLETE", "").strip()}
+    ])
+    new_state["messages"] = new_messages
+    
+    return new_state