Update diagnostics.py

diagnostics.py

@@ -1,243 +1,154 @@
 """
-Updated diagnostics module for the Pharmaceutical Analytics application.
-Uses the simplified Planning Agent implementation.
+Simplified Planning Agent for Pharmaceutical Analytics
+This version uses direct API calls instead of LangChain components
 """
 
-import streamlit as st
 import os
 import json
-import traceback
-import time
-from datetime import datetime
-from typing import Dict, Any, Tuple
+import re
+from typing import Dict, List, Any, Tuple
+from pydantic import BaseModel, Field
 
-def test_claude_connectivity() -> Tuple[bool, str]:
-    """
-    Test basic connectivity to Claude API
+# 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")
+    required_data_sources: List[Dict[str, str]] = Field(
+        description="List of data sources needed with table name and purpose")
+    analysis_approaches: List[Dict[str, str]] = Field(
+        description="List of analytical approaches to be used with type and purpose")
+    tasks: List[Dict[str, Any]] = Field(
+        description="Ordered list of tasks to execute with dependencies")
+    expected_insights: List[str] = Field(
+        description="List of expected insights that would answer the problem")
+
+class PlanningAgent:
+    """Agent responsible for planning the analysis workflow"""
     
-    Returns:
-        Tuple of (success_flag, response_or_error)
-    """
-    try:
-        # Check API key
+    def __init__(self):
+        """Initialize the planning agent"""
         api_key = os.getenv("ANTHROPIC_API_KEY")
         if not api_key:
-            return False, "No ANTHROPIC_API_KEY found in environment variables"
-        
-        # Import anthropic
-        try:
-            import anthropic
-        except ImportError:
-            return False, "Failed to import anthropic library. Make sure it's installed (pip install anthropic)"
-        
-        # Create client and send request
-        client = anthropic.Anthropic(api_key=api_key)
-        
-        response = client.messages.create(
-            model="claude-3-haiku-20240307",
-            max_tokens=300,
-            messages=[
-                {"role": "user", "content": "Briefly tell me about Egyptian history in 2-3 sentences."}
-            ]
-        )
+            raise ValueError("ANTHROPIC_API_KEY not found in environment variables")
         
-        if response and response.content:
-            return True, response.content[0].text
-        else:
-            return False, "Received empty response from Anthropic API"
+        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...")
         
-    except Exception as e:
-        error_msg = f"Error testing Claude API: {str(e)}\n\n{traceback.format_exc()}"
-        return False, error_msg
+        # 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.
 
-def test_simplified_agent() -> Tuple[bool, Dict[str, Any]]:
-    """
-    Test the simplified Planning Agent functionality
-    
-    Returns:
-        Tuple of (success_flag, results_dict)
-    """
-    results = {
-        "steps": [],
-        "errors": [],
-        "plan": None
+For pharmaceutical sales analysis:
+- Consider product performance, competitor activities, prescriber behavior
+- Include geographic, temporal, and demographic dimensions in your analysis
+- Consider both internal factors (supply, marketing) and external factors (market events, seasonality)
+
+Your output should be a complete JSON-formatted analysis plan following this structure:
+{
+  "problem_statement": "Clear definition of the problem to solve",
+  "required_data_sources": [
+    {"table": "sales", "purpose": "Core sales metrics analysis"},
+    {"table": "regions", "purpose": "Geographic segmentation"}
+  ],
+  "analysis_approaches": [
+    {"type": "time_series_decomposition", "purpose": "Separate trend from seasonality"},
+    {"type": "comparative_analysis", "purpose": "Compare performance across regions"}
+  ],
+  "tasks": [
+    {
+      "id": 1,
+      "name": "Data acquisition",
+      "description": "Pull relevant data from sources",
+      "agent": "data_agent",
+      "dependencies": [],
+      "expected_output": "Cleaned datasets for analysis"
+    },
+    {
+      "id": 2,
+      "name": "Analysis execution",
+      "description": "Perform statistical analysis",
+      "agent": "analytics_agent",
+      "dependencies": [1],
+      "expected_output": "Analysis results"
     }
-    
-    try:
-        # Test Planning Agent import
-        results["steps"].append({"step": "import", "status": "started"})
-        try:
-            from simplified_planning_agent import PlanningAgent
-            results["steps"].append({"step": "import", "status": "success"})
-        except Exception as import_error:
-            error_msg = f"Error importing simplified PlanningAgent: {str(import_error)}\n\n{traceback.format_exc()}"
-            results["errors"].append({"phase": "import", "error": error_msg})
-            return False, results
+  ],
+  "expected_insights": [
+    "Primary factors contributing to sales decline",
+    "Regional variations in performance"
+  ]
+}
+
+Be thorough but focus on creating a practical analysis workflow.
+"""
         
-        # Test agent initialization
-        results["steps"].append({"step": "initialization", "status": "started"})
-        try:
-            planning_agent = PlanningAgent()
-            results["steps"].append({"step": "initialization", "status": "success"})
-        except Exception as init_error:
-            error_msg = f"Error initializing PlanningAgent: {str(init_error)}\n\n{traceback.format_exc()}"
-            results["errors"].append({"phase": "initialization", "error": error_msg})
-            return False, results
+        user_message = f"Create an analysis plan for the following alert: {alert_description}"
         
-        # Test plan creation
-        results["steps"].append({"step": "plan_creation", "status": "started"})
+        # Make direct API call to Claude
         try:
-            test_alert = "Sales of DrugX down 15% in Northeast region over past 30 days."
-            analysis_plan, plan_dict = planning_agent.create_analysis_plan(test_alert)
+            import anthropic
+            client = anthropic.Anthropic(api_key=self.api_key)
             
-            results["steps"].append({"step": "plan_creation", "status": "success"})
-            results["plan"] = plan_dict
-            return True, results
-        except Exception as plan_error:
-            error_msg = f"Error creating analysis plan: {str(plan_error)}\n\n{traceback.format_exc()}"
-            results["errors"].append({"phase": "plan_creation", "error": error_msg})
-            return False, results
+            # FIXED: Pass system as a parameter, not as a message
+            response = client.messages.create(
+                model="claude-3-haiku-20240307",
+                max_tokens=2000,
+                temperature=0.2,
+                system=system_prompt,  # Pass as system parameter
+                messages=[
+                    {"role": "user", "content": user_message}
+                ]
+            )
             
-    except Exception as e:
-        error_msg = f"Unexpected error in agent functionality test: {str(e)}\n\n{traceback.format_exc()}"
-        results["errors"].append({"phase": "unknown", "error": error_msg})
-        return False, results
-
-def render_diagnostics_tab():
-    """
-    Render the diagnostics tab in the Streamlit app
-    """
-    st.header("🔧 System Diagnostics")
-    
-    # Environment information
-    st.subheader("Environment Information")
-    
-    # Check API key
-    api_key = os.getenv("ANTHROPIC_API_KEY")
-    api_status = "✅ Set" if api_key else "❌ Not Set"
-    
-    environment_data = {
-        "API Key": api_status,
-        "Python Version": os.sys.version,
-        "Timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
-    }
-    
-    # Try to get library versions
-    try:
-        import anthropic
-        environment_data["Anthropic Library"] = anthropic.__version__
-    except (ImportError, AttributeError):
-        environment_data["Anthropic Library"] = "Not found"
-    
-    try:
-        import langgraph
-        environment_data["LangGraph Library"] = langgraph.__version__
-    except (ImportError, AttributeError):
-        environment_data["LangGraph Library"] = "Not found"
-    
-    # Display environment table
-    for key, value in environment_data.items():
-        cols = st.columns([1, 3])
-        with cols[0]:
-            st.markdown(f"**{key}:**")
-        with cols[1]:
-            st.markdown(f"{value}")
-    
-    # Claude connectivity test
-    st.subheader("Claude LLM Test")
-    if st.button("Test Claude Connectivity"):
-        with st.spinner("Testing connection to Claude..."):
-            success, response = test_claude_connectivity()
-        
-        if success:
-            st.success("✅ Claude API is working properly!")
-            st.markdown("**Response:**")
-            st.info(response)
-        else:
-            st.error("❌ Claude API test failed")
-            st.expander("Error Details", expanded=True).error(response)
-    
-    # Simplified Agent functionality test
-    st.subheader("Simplified Planning Agent Test")
-    if st.button("Test Simplified Planning Agent"):
-        with st.spinner("Testing Simplified Planning Agent..."):
-            success, results = test_simplified_agent()
-        
-        if success:
-            st.success("✅ Simplified Planning Agent is working properly!")
-            st.json(results["plan"])
-        else:
-            st.error("�� Simplified Planning Agent test failed")
-            for error in results["errors"]:
-                st.expander(f"Error in {error['phase']} phase", expanded=True).error(error["error"])
-    
-    # Original Agent functionality test
-    st.subheader("Original Planning Agent Test (Likely to Fail)")
-    if st.button("Test Original Planning Agent"):
-        with st.spinner("Testing Original Planning Agent..."):
-            try:
-                from agents.planning_agent import PlanningAgent
-                planning_agent = PlanningAgent()
-                test_alert = "Sales of DrugX down 15% in Northeast region over past 30 days."
-                analysis_plan, plan_dict = planning_agent.create_analysis_plan(test_alert)
-                
-                st.success("✅ Original Planning Agent is working!")
-                st.json(plan_dict)
-            except Exception as e:
-                st.error("❌ Original Planning Agent test failed")
-                st.expander("Error Details", expanded=True).error(f"{str(e)}\n\n{traceback.format_exc()}")
-    
-    # Database Connectivity
-    st.subheader("Database Connectivity")
-    if st.button("Test Database Connection"):
-        with st.spinner("Testing database connection..."):
-            try:
-                import sqlite3
-                conn = sqlite3.connect("data/pharma_db.sqlite")
-                
-                # Test query
-                cursor = conn.cursor()
-                cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
-                tables = cursor.fetchall()
-                
-                # Close connection
-                conn.close()
-                
-                st.success(f"✅ Successfully connected to database")
-                st.write(f"Found {len(tables)} tables:")
-                st.json([table[0] for table in tables])
-                
-            except Exception as e:
-                st.error(f"❌ Database connection failed: {str(e)}")
-                st.expander("Error Details", expanded=True).error(traceback.format_exc())
-    
-    # Add guidance
-    st.markdown("---")
-    st.markdown("""
-    ### Using the Simplified Agent
-    
-    The original Planning Agent implementation has issues with the LangChain prompt templates. 
-    We've created a simplified version that uses direct API calls to Claude.
-    
-    To use the simplified agent:
-    
-    1. Copy the `simplified_planning_agent.py` file to your project
-    2. Update your workflow to import from `simplified_planning_agent` instead of `agents.planning_agent`
-    
-    ### Troubleshooting Tips
-    
-    1. If Claude test fails:
-       - Check your API key in environment variables
-       - Verify internet connectivity
-       - Confirm your API key has sufficient credits
-    
-    2. If Agent test fails:
-       - Check import errors (are all libraries installed?)
-       - Verify agent implementation
-       - Check for syntax errors in agent code
+            # 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
     
-    3. If Database test fails:
-       - Verify database file exists
-       - Check file permissions
-       - Run setup_environment() function again
-    """)
+    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
+            return json.loads(text)
+        except json.JSONDecodeError:
+            # Try to find JSON block with regex
+            json_pattern = r'```json\s*([\s\S]*?)\s*```'
+            match = re.search(json_pattern, text)
+            if match:
+                try:
+                    return json.loads(match.group(1))
+                except json.JSONDecodeError:
+                    pass
+            
+            # Try to find anything that looks like JSON
+            json_pattern = r'({[\s\S]*})'
+            match = re.search(json_pattern, text)
+            if match:
+                try:
+                    return json.loads(match.group(1))
+                except json.JSONDecodeError:
+                    pass
+            
+            # If all extraction attempts fail
+            raise ValueError(f"Could not extract JSON from response: {text}")
+
+# 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))