Create workflows/sales_analysis.py

workflows/sales_analysis.py

@@ -0,0 +1,384 @@
+import os
+import json
+from typing import Dict, List, Any, Tuple, Optional, TypedDict, Annotated, Literal, cast
+from datetime import datetime
+import uuid
+from pydantic import BaseModel, Field
+import langgraph.graph as lg
+from langgraph.graph import END, StateGraph
+from langchain_anthropic import ChatAnthropic
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+
+# Import the agents
+from agents.planning_agent import PlanningAgent, AnalysisPlan
+from agents.data_agent import DataAgent, DataRequest, DataSource
+from agents.analytics_agent import AnalyticsAgent, AnalysisRequest, AnalysisResult
+from agents.qa_agent import QAAgent, ValidationRequest, ValidationResult
+from agents.insights_agent import InsightsAgent, InsightRequest, InsightCard
+
+# Define the workflow state
+class WorkflowState(TypedDict):
+    alert: str
+    status: Literal["planning", "data_collection", "analysis", "validation", "insights", "complete", "error"]
+    plan: Optional[AnalysisPlan]
+    data_requests: List[DataRequest]
+    data_sources: Dict[str, DataSource]
+    analysis_requests: List[AnalysisRequest]
+    analysis_results: Dict[str, AnalysisResult]
+    validation_requests: List[ValidationRequest]
+    validation_results: Dict[str, ValidationResult]
+    insight_requests: List[InsightRequest]
+    insight_cards: Dict[str, InsightCard]
+    error: Optional[str]
+    logs: List[Dict[str, Any]]
+    visualizations: List[str]
+
+# Function to initialize workflow state
+def create_initial_state(alert: str) -> WorkflowState:
+    """Create initial workflow state from alert"""
+    return WorkflowState(
+        alert=alert,
+        status="planning",
+        plan=None,
+        data_requests=[],
+        data_sources={},
+        analysis_requests=[],
+        analysis_results={},
+        validation_requests=[],
+        validation_results={},
+        insight_requests=[],
+        insight_cards={},
+        error=None,
+        logs=[],
+        visualizations=[]
+    )
+
+# Define the agent nodes
+class SalesAnalysisWorkflow:
+    """Orchestrates the sales analysis workflow with multiple agents"""
+    
+    def __init__(self, db_path: str = "data/pharma_db.sqlite"):
+        """Initialize the workflow with agents"""
+        # Initialize agents
+        self.planning_agent = PlanningAgent()
+        self.data_agent = DataAgent(db_path=db_path)
+        self.analytics_agent = AnalyticsAgent()
+        self.qa_agent = QAAgent()
+        self.insights_agent = InsightsAgent()
+        
+        # Create the workflow graph
+        self.workflow = self._build_workflow()
+    
+    def _build_workflow(self) -> StateGraph:
+        """Build the workflow graph"""
+        # Create the state graph
+        workflow = StateGraph(WorkflowState)
+        
+        # Add nodes for each agent
+        workflow.add_node("planning", self.planning_node)
+        workflow.add_node("data_collection", self.data_collection_node)
+        workflow.add_node("analysis", self.analysis_node)
+        workflow.add_node("validation", self.validation_node)
+        workflow.add_node("insights", self.insights_node)
+        
+        # Define the workflow edges
+        workflow.add_edge("planning", "data_collection")
+        workflow.add_edge("data_collection", "analysis")
+        workflow.add_edge("analysis", "validation")
+        workflow.add_edge("validation", "insights")
+        workflow.add_edge("insights", END)
+        
+        # Add error handling
+        workflow.add_node("handle_error", self.handle_error_node)
+        workflow.add_edge_condition(
+            "planning", "handle_error", 
+            lambda x: x["status"] == "error"
+        )
+        workflow.add_edge_condition(
+            "data_collection", "handle_error", 
+            lambda x: x["status"] == "error"
+        )
+        workflow.add_edge_condition(
+            "analysis", "handle_error", 
+            lambda x: x["status"] == "error"
+        )
+        workflow.add_edge_condition(
+            "validation", "handle_error", 
+            lambda x: x["status"] == "error"
+        )
+        workflow.add_edge_condition(
+            "insights", "handle_error", 
+            lambda x: x["status"] == "error"
+        )
+        workflow.add_edge("handle_error", END)
+        
+        # Set the entry point
+        workflow.set_entry_point("planning")
+        
+        return workflow.compile()
+    
+    def planning_node(self, state: WorkflowState) -> WorkflowState:
+        """Planning agent node"""
+        try:
+            # Log the step
+            state["logs"].append({
+                "timestamp": datetime.now().isoformat(),
+                "step": "planning",
+                "message": "Planning analysis approach"
+            })
+            
+            # Update status
+            state["status"] = "planning"
+            
+            # Create analysis plan
+            analysis_plan, plan_dict = self.planning_agent.create_analysis_plan(state["alert"])
+            
+            # Store the plan in the state
+            state["plan"] = analysis_plan
+            
+            # Create data requests from plan
+            data_requests = []
+            for i, data_source in enumerate(analysis_plan.required_data_sources):
+                data_requests.append(DataRequest(
+                    request_id=f"data_{i}",
+                    description=data_source["purpose"],
+                    tables=[data_source["table"]],
+                    purpose=data_source["purpose"]
+                ))
+            
+            # Store data requests in the state
+            state["data_requests"] = data_requests
+            
+            # Update status for next step
+            state["status"] = "data_collection"
+            
+            return state
+        except Exception as e:
+            # Handle error
+            state["status"] = "error"
+            state["error"] = f"Planning error: {str(e)}"
+            return state
+    
+    def data_collection_node(self, state: WorkflowState) -> WorkflowState:
+        """Data collection agent node"""
+        try:
+            # Log the step
+            state["logs"].append({
+                "timestamp": datetime.now().isoformat(),
+                "step": "data_collection",
+                "message": f"Collecting data from {len(state['data_requests'])} sources"
+            })
+            
+            # Update status
+            state["status"] = "data_collection"
+            
+            # Process data requests
+            data_sources = self.data_agent.get_data_for_analysis(state["data_requests"])
+            
+            # Store data sources in the state
+            state["data_sources"] = data_sources
+            
+            # Create analysis requests based on plan
+            analysis_requests = []
+            if state["plan"]:
+                analysis_approaches = state["plan"].analysis_approaches
+                for i, approach in enumerate(analysis_approaches):
+                    analysis_requests.append(AnalysisRequest(
+                        request_id=f"analysis_{i}",
+                        description=approach["purpose"],
+                        data_sources=list(data_sources.keys()),
+                        analysis_type=approach["type"],
+                        purpose=approach["purpose"]
+                    ))
+            
+            # Store analysis requests in the state
+            state["analysis_requests"] = analysis_requests
+            
+            # Update status for next step
+            state["status"] = "analysis"
+            
+            return state
+        except Exception as e:
+            # Handle error
+            state["status"] = "error"
+            state["error"] = f"Data collection error: {str(e)}"
+            return state
+    
+    def analysis_node(self, state: WorkflowState) -> WorkflowState:
+        """Analytics agent node"""
+        try:
+            # Log the step
+            state["logs"].append({
+                "timestamp": datetime.now().isoformat(),
+                "step": "analysis",
+                "message": f"Performing {len(state['analysis_requests'])} analysis tasks"
+            })
+            
+            # Update status
+            state["status"] = "analysis"
+            
+            # Process analysis requests
+            analysis_results = {}
+            for request in state["analysis_requests"]:
+                result = self.analytics_agent.perform_analysis(request, state["data_sources"])
+                analysis_results[request.request_id] = result
+            
+            # Store analysis results in the state
+            state["analysis_results"] = analysis_results
+            
+            # Create validation requests
+            validation_requests = []
+            for analysis_id, result in analysis_results.items():
+                validation_requests.append(ValidationRequest(
+                    request_id=analysis_id,
+                    original_problem=state["alert"],
+                    analysis_results=result.model_dump(),
+                    data_sources=list(state["data_sources"].keys())
+                ))
+            
+            # Store validation requests in the state
+            state["validation_requests"] = validation_requests
+            
+            # Update status for next step
+            state["status"] = "validation"
+            
+            return state
+        except Exception as e:
+            # Handle error
+            state["status"] = "error"
+            state["error"] = f"Analysis error: {str(e)}"
+            return state
+    
+    def validation_node(self, state: WorkflowState) -> WorkflowState:
+        """QA agent node"""
+        try:
+            # Log the step
+            state["logs"].append({
+                "timestamp": datetime.now().isoformat(),
+                "step": "validation",
+                "message": f"Validating {len(state['validation_requests'])} analysis results"
+            })
+            
+            # Update status
+            state["status"] = "validation"
+            
+            # Process validation requests
+            validation_results = {}
+            for request in state["validation_requests"]:
+                result = self.qa_agent.validate_analysis(request, state["data_sources"])
+                validation_results[request.request_id] = result
+            
+            # Store validation results in the state
+            state["validation_results"] = validation_results
+            
+            # Create insight requests
+            insight_requests = []
+            for analysis_id, validation in validation_results.items():
+                analysis_result = state["analysis_results"][analysis_id]
+                insight_requests.append(InsightRequest(
+                    request_id=analysis_id,
+                    original_problem=state["alert"],
+                    analysis_results=analysis_result.model_dump(),
+                    validation_results=validation.model_dump(),
+                    target_audience="executive"
+                ))
+            
+            # Store insight requests in the state
+            state["insight_requests"] = insight_requests
+            
+            # Update status for next step
+            state["status"] = "insights"
+            
+            return state
+        except Exception as e:
+            # Handle error
+            state["status"] = "error"
+            state["error"] = f"Validation error: {str(e)}"
+            return state
+    
+    def insights_node(self, state: WorkflowState) -> WorkflowState:
+        """Insights agent node"""
+        try:
+            # Log the step
+            state["logs"].append({
+                "timestamp": datetime.now().isoformat(),
+                "step": "insights",
+                "message": f"Generating insights from {len(state['insight_requests'])} validated analyses"
+            })
+            
+            # Update status
+            state["status"] = "insights"
+            
+            # Process insight requests
+            insight_cards = {}
+            visualizations = []
+            for request in state["insight_requests"]:
+                # Generate insight card
+                insight_card = self.insights_agent.generate_insights(request)
+                insight_cards[request.request_id] = insight_card
+                
+                # Generate visualizations for the insight card
+                viz_files = self.insights_agent.generate_visualizations(insight_card, state["data_sources"])
+                visualizations.extend(viz_files)
+            
+            # Store insight cards and visualizations in the state
+            state["insight_cards"] = insight_cards
+            state["visualizations"] = visualizations
+            
+            # Update status for completion
+            state["status"] = "complete"
+            
+            # Log completion
+            state["logs"].append({
+                "timestamp": datetime.now().isoformat(),
+                "step": "complete",
+                "message": f"Analysis workflow completed with {len(insight_cards)} insight cards"
+            })
+            
+            return state
+        except Exception as e:
+            # Handle error
+            state["status"] = "error"
+            state["error"] = f"Insights error: {str(e)}"
+            return state
+    
+    def handle_error_node(self, state: WorkflowState) -> WorkflowState:
+        """Error handling node"""
+        # Log the error
+        state["logs"].append({
+            "timestamp": datetime.now().isoformat(),
+            "step": "error",
+            "message": f"Workflow error: {state['error']}"
+        })
+        
+        return state
+    
+    def run_workflow(self, alert: str) -> WorkflowState:
+        """Run the workflow with a given alert"""
+        # Create initial state
+        initial_state = create_initial_state(alert)
+        
+        # Run the workflow
+        final_state = self.workflow.invoke(initial_state)
+        
+        return final_state
+
+# For testing
+if __name__ == "__main__":
+    # Set API key for testing
+    os.environ["ANTHROPIC_API_KEY"] = "your_api_key_here"
+    
+    # Create workflow
+    workflow = SalesAnalysisWorkflow(db_path="data/pharma_db.sqlite")
+    
+    # Run workflow with test alert
+    alert = "Sales of DrugX down 15% in Northeast region over past 30 days compared to forecast."
+    result = workflow.run_workflow(alert)
+    
+    # Print workflow results
+    print(f"Workflow status: {result['status']}")
+    if result["status"] == "error":
+        print(f"Error: {result['error']}")
+    else:
+        print(f"Generated {len(result['insight_cards'])} insight cards")
+        print(f"Generated {len(result['visualizations'])} visualizations")