Update simplified_planning_agent.py

simplified_planning_agent.py

@@ -35,11 +35,11 @@ class PlanningAgent:
         print("Planning Agent initialized successfully")
     
     def create_analysis_plan(self, alert_description: str) -> Tuple[AnalysisPlan, Dict]:
-    """Generate an analysis plan based on the alert description"""
-    print("Planning Agent: Creating analysis plan...")
-    
-    # Create the system prompt and user message
-    system_prompt = """You are an expert pharmaceutical analytics planning agent.
+        """Generate an analysis plan based on the alert description"""
+        print("Planning Agent: Creating analysis plan...")
+        
+        # Create the system prompt and user message
+        system_prompt = """You are an expert pharmaceutical analytics planning agent.
 Your task is to create a detailed analysis plan to investigate sales anomalies.
 
 For pharmaceutical sales analysis:
@@ -84,32 +84,52 @@ Your output should be a complete JSON-formatted analysis plan following this str
 
 Be thorough but focus on creating a practical analysis workflow.
 """
-    
-    user_message = f"Create an analysis plan for the following alert: {alert_description}"
-    
-    # Make direct API call to Claude
-    try:
-        import anthropic
-        client = anthropic.Anthropic(api_key=self.api_key)
         
-        # Use the correct API structure based on the Anthropic Python SDK version
+        user_message = f"Create an analysis plan for the following alert: {alert_description}"
+        
+        # Make direct API call to Claude
         try:
-            # For newer versions of the Anthropic SDK
-            response = client.messages.create(
-                model="claude-3-haiku-20240307",
-                max_tokens=2000,
-                temperature=0.2,
-                system=system_prompt,
-                messages=[
-                    {"role": "user", "content": user_message}
-                ]
-            )
-        except TypeError:
-            # Fallback for older versions of the Anthropic SDK
-            response = client.messages.create(
-                model="claude-3-haiku-20240307",
-                max_tokens=2000,
-                temperature=0.2,
+            import anthropic
+            client = anthropic.Anthropic(api_key=self.api_key)
+            
+            # Use the correct API structure based on the Anthropic Python SDK version
+            try:
+                # For newer versions of the Anthropic SDK
+                response = client.messages.create(
+                    model="claude-3-haiku-20240307",
+                    max_tokens=2000,
+                    temperature=0.2,
+                    system=system_prompt,
+                    messages=[
+                        {"role": "user", "content": user_message}
+                    ]
+                )
+            except TypeError:
+                # Fallback for older versions of the Anthropic SDK
+                response = client.messages.create(
+                    model="claude-3-haiku-20240307",
+                    max_tokens=2000,
+                    temperature=0.2,
+                    messages=[
+                        {"role": "system", "content": system_prompt},
+                        {"role": "user", "content": user_message}
+                    ]
+                )
+            
+            # Extract response content
+            response_text = response.content[0].text
+            
+            # Extract JSON from the response
+            plan_dict = self.extract_json_from_text(response_text)
+            
+            # Convert to Pydantic model for validation
+            analysis_plan = AnalysisPlan.model_validate(plan_dict)
+            
+            return analysis_plan, plan_dict
+            
+        except Exception as e:
+            print(f"Error creating analysis plan: {e}")
+            raise
     
     def extract_json_from_text(self, text: str) -> Dict:
         """Extract JSON from text that might contain additional content"""