Update agents/planning_agent.py

agents/planning_agent.py

@@ -1,13 +1,15 @@
+"""
+Simplified Planning Agent for Pharmaceutical Analytics
+This version uses direct API calls instead of LangChain components
+"""
+
 import os
-from typing import Dict, List, Optional, Any, Tuple
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
-from langchain_core.runnables import RunnablePassthrough
-from langchain_anthropic import ChatAnthropic
-from pydantic import BaseModel, Field
 import json
+import re
+from typing import Dict, List, Any, Tuple
+from pydantic import BaseModel, Field
 
-# Define task types and output schema
+# Define analysis plan schema
 class AnalysisPlan(BaseModel):
     """Planning agent output with analysis plan details"""
     problem_statement: str = Field(description="Refined problem statement based on the alert")
@@ -24,21 +26,20 @@ class PlanningAgent:
     """Agent responsible for planning the analysis workflow"""
     
     def __init__(self):
-        """Initialize the planning agent with Claude API"""
-        # Set up Claude API client
+        """Initialize the planning agent"""
         api_key = os.getenv("ANTHROPIC_API_KEY")
         if not api_key:
             raise ValueError("ANTHROPIC_API_KEY not found in environment variables")
         
-        self.llm = ChatAnthropic(
-            model="claude-3-haiku-20240307",  # Changed to haiku to use less tokens
-            anthropic_api_key=api_key,
-            temperature=0.1
-        )
+        self.api_key = api_key
+        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 planning prompt - fixed variable issue in the template
-        self.planning_prompt = ChatPromptTemplate.from_messages([
-            ("system", """You are an expert pharmaceutical analytics planning agent.
+        # 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:
@@ -81,28 +82,62 @@ Your output should be a complete JSON-formatted analysis plan following this str
   ]
 }
 
-Be thorough in your planning but focus on creating a practical analysis workflow.
-Tasks should follow a logical sequence with proper dependencies.
-"""),
-            ("human", "{input}")  # Changed from {alert} to {input}
-        ])
+Be thorough but focus on creating a practical analysis workflow.
+"""
+        
+        user_message = f"Create an analysis plan for the following alert: {alert_description}"
         
-        # Set up the planning chain
-        self.planning_chain = (
-            {"input": RunnablePassthrough()}  # Use input as the key
-            | self.planning_prompt
-            | self.llm
-            | StrOutputParser()
-        )
+        # 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
+            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"""
         try:
-            # First, try to parse the entire text as JSON
+            # First try to parse the entire text as JSON
             return json.loads(text)
         except json.JSONDecodeError:
-            # If that fails, look for JSON block
-            import re
+            # Try to find JSON block with regex
             json_pattern = r'```json\s*([\s\S]*?)\s*```'
             match = re.search(json_pattern, text)
             if match:
@@ -111,7 +146,7 @@ Tasks should follow a logical sequence with proper dependencies.
                 except json.JSONDecodeError:
                     pass
             
-            # Try a more aggressive approach to find JSON-like content
+            # Try to find anything that looks like JSON
             json_pattern = r'({[\s\S]*})'
             match = re.search(json_pattern, text)
             if match:
@@ -120,70 +155,13 @@ Tasks should follow a logical sequence with proper dependencies.
                 except json.JSONDecodeError:
                     pass
             
+            # If all extraction attempts fail
             raise ValueError(f"Could not extract JSON from response: {text}")
-    
-    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...")
-        
-        try:
-            # Execute the planning chain with the alert as input
-            response = self.planning_chain.invoke(alert_description)
-            
-            # Extract and parse the response as JSON
-            plan_dict = self.extract_json_from_text(response)
-            
-            # Convert to Pydantic model for validation and structure
-            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 visualize_plan(self, plan: AnalysisPlan) -> Dict:
-        """Generate visualization data for the analysis plan"""
-        # Create nodes representing tasks
-        nodes = []
-        edges = []
-        
-        for task in plan.tasks:
-            nodes.append({
-                "id": f"task_{task['id']}",
-                "label": task['name'],
-                "type": "task",
-                "agent": task['agent']
-            })
-            
-            # Create edges based on dependencies
-            for dep in task.get('dependencies', []):
-                edges.append({
-                    "source": f"task_{dep}",
-                    "target": f"task_{task['id']}",
-                    "label": "depends on"
-                })
-        
-        # Add data source nodes
-        for i, src in enumerate(plan.required_data_sources):
-            src_id = f"data_{i}"
-            nodes.append({
-                "id": src_id,
-                "label": src['table'],
-                "type": "data_source"
-            })
-            
-            # Connect data sources to the data acquisition task
-            data_task = next((t for t in plan.tasks if t['agent'] == 'data_agent'), None)
-            if data_task:
-                edges.append({
-                    "source": src_id,
-                    "target": f"task_{data_task['id']}",
-                    "label": "input"
-                })
-        
-        return {
-            "nodes": nodes,
-            "edges": edges,
-            "problem_statement": plan.problem_statement,
-            "expected_insights": plan.expected_insights
-        }
+
+# For testing
+if __name__ == "__main__":
+    # Get API key from environment
+    agent = PlanningAgent()
+    alert = "Sales of DrugX down 15% in Northeast region over past 30 days compared to forecast."
+    plan, plan_dict = agent.create_analysis_plan(alert)
+    print(json.dumps(plan_dict, indent=2))