Update agents/sql_generator.py

agents/sql_generator.py

@@ -5,11 +5,10 @@ This agent converts a data pipeline plan into executable SQL queries.
 
 import time
 from typing import Dict, Any, List
-from agents.state import AgentState
-from anthropic import Anthropic
+from langgraph.graph.message import add_messages
 from anthropic.types import MessageParam
 
-def sql_generator_agent(anthropic_client: Anthropic, state: AgentState) -> AgentState:
+def sql_generator_agent(anthropic_client, state: Dict[str, Any]) -> Dict[str, Any]:
     """
     Agent that generates SQL queries based on the pipeline plan.
     
@@ -22,7 +21,6 @@ def sql_generator_agent(anthropic_client: Anthropic, state: AgentState) -> Agent
     """
     # Get current messages and pipeline plan
     messages = state.get("messages", [])
-    current_agent = state.get("current_agent", ["sql_generator_agent"])
     pipeline_plan = state.get("pipeline_plan", {})
     
     # Add agent-specific instructions
@@ -50,68 +48,46 @@ def sql_generator_agent(anthropic_client: Anthropic, state: AgentState) -> Agent
     - 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"]
-        ))
-    
-    # Add final user message with context
-    prompt_messages.append(
-        MessageParam(
-            role="user", 
-            content=f"Based on this pipeline plan, generate the SQL queries needed. {context}"
-        )
+    # 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 this pipeline plan, generate the SQL queries needed. {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 SQL generation is complete
-        sql_complete = "SQL_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("SQL_COMPLETE", "").strip()
-        })
-        
+    # Extract the response
+    agent_response = response.content[0].text
+    
+    # Check if SQL generation is complete
+    sql_complete = "SQL_COMPLETE" in agent_response
+    
+    # Update state based on SQL completeness
+    new_state = state.copy()
+    if sql_complete:
         # Extract SQL queries from the response
-        if sql_complete:
-            sql_queries = _extract_sql_queries(agent_response)
-            
-            new_state["sql_queries"] = sql_queries
-            # Use .append() to add next agent to the list
-            new_state["current_agent"].append("executor_agent")
-        else:
-            # Use .append() to add current agent if SQL generation is not complete
-            new_state["current_agent"].append("sql_generator_agent")
+        # In a production system, you would parse the SQL more carefully
+        sql_queries = _extract_sql_queries(agent_response)
         
-        return new_state
-    
-    except Exception as e:
-        # Handle any errors
-        print(f"Error in sql_generator_agent: {str(e)}")
-        return {
-            "messages": state.get("messages", []) + [
-                {"role": "assistant", "content": f"I encountered an error: {str(e)}"}
-            ],
-            "current_agent": current_agent + ["sql_generator_agent"]
-        }
+        new_state["sql_queries"] = sql_queries
+        new_state["current_agent"] = "executor_agent"
+    else:
+        # Need more information or work, stay with SQL generator agent
+        new_state["current_agent"] = "sql_generator_agent"
+    
+    # Add agent's response to messages
+    new_messages = add_messages(state, [
+        {"role": "assistant", "content": agent_response.replace("SQL_COMPLETE", "").strip()}
+    ])
+    new_state["messages"] = new_messages
+    
+    return new_state
 
 def _extract_sql_queries(response: str) -> List[Dict[str, Any]]:
     """