Update agents/sql_generator.py

agents/sql_generator.py

@@ -20,8 +20,8 @@ def sql_generator_agent(anthropic_client, state: Dict[str, Any]) -> Dict[str, An
         Updated state
     """
     # Get current messages and pipeline plan
-    messages = state["messages"]
-    pipeline_plan = state["pipeline_plan"]
+    messages = state.get("messages", [])
+    pipeline_plan = state.get("pipeline_plan", {})
     
     # Add agent-specific instructions
     system_message = """
@@ -39,7 +39,7 @@ def sql_generator_agent(anthropic_client, state: Dict[str, Any]) -> Dict[str, An
     
     # Prepare context for Claude
     context = f"""
-    Pipeline plan: {pipeline_plan['description']}
+    Pipeline plan: {pipeline_plan.get('description', 'No plan provided')}
     
     Available tables:
     - Raw data: RAW_SALES_TRANSACTIONS, RAW_HCP_DATA, RAW_PRODUCT_DATA
@@ -48,17 +48,28 @@ def sql_generator_agent(anthropic_client, state: Dict[str, Any]) -> Dict[str, An
     - Data products: DP_SALES_DASHBOARD, DP_HCP_TARGETING
     """
     
-    # Prepare prompt for Claude
-    prompt_messages = [
-        *[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}")
-    ]
+    # Convert messages to the format expected by Anthropic API
+    anthropic_messages = []
+    for msg in messages:
+        if isinstance(msg, dict) and "role" in msg and "content" in msg:
+            anthropic_messages.append(MessageParam(
+                role=msg["role"],
+                content=msg["content"]
+            ))
+    
+    # Add final user message with context
+    anthropic_messages.append(
+        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",
         system=system_message,
-        messages=prompt_messages,
+        messages=anthropic_messages,
         max_tokens=3000
     )
     
@@ -76,74 +87,4 @@ def sql_generator_agent(anthropic_client, state: Dict[str, Any]) -> Dict[str, An
         sql_queries = _extract_sql_queries(agent_response)
         
         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]]:
-    """
-    Extract SQL queries from the agent's response.
-    
-    Args:
-        response: The text response containing SQL queries
-    
-    Returns:
-        List of dictionaries containing query information
-    """
-    # Remove the SQL_COMPLETE tag if present
-    clean_response = response.replace("SQL_COMPLETE", "")
-    
-    # Extract SQL code blocks
-    # This is a simple extraction that looks for ```sql ... ``` blocks
-    # In production, you would want a more robust parser
-    sql_blocks = []
-    current_pos = 0
-    
-    while True:
-        start_marker = "```sql"
-        end_marker = "```"
-        
-        start_pos = clean_response.find(start_marker, current_pos)
-        if start_pos == -1:
-            break
-            
-        # Find the end of this code block
-        end_pos = clean_response.find(end_marker, start_pos + len(start_marker))
-        if end_pos == -1:
-            break
-            
-        # Extract the SQL query
-        sql_content = clean_response[start_pos + len(start_marker):end_pos].strip()
-        
-        # Add to our list
-        sql_blocks.append(sql_content)
-        
-        # Move position forward
-        current_pos = end_pos + len(end_marker)
-    
-    # Convert to query objects
-    sql_queries = []
-    for i, sql in enumerate(sql_blocks):
-        # Try to extract a purpose comment from the SQL
-        purpose = "Data transformation"  # Default
-        lines = sql.split('\n')
-        if lines and lines[0].strip().startswith('--'):
-            purpose = lines[0].strip()[2:].strip()
-        
-        sql_queries.append({
-            "name": f"Query {i+1}",
-            "sql": sql,
-            "purpose": purpose,
-            "created_at": time.time()
-        })
-    
-    return sql_queries
+        new_st