Update workflows/sales_analysis.py

workflows/sales_analysis.py

@@ -6,8 +6,7 @@ from typing import Dict, List, Any, Tuple, Optional, TypedDict, Annotated, Liter
 from datetime import datetime
 import uuid
 from pydantic import BaseModel, Field
-import langgraph.graph as lg
-from langgraph.graph import END, StateGraph
+from langgraph.graph import StateGraph, END
 from langchain_anthropic import ChatAnthropic
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
 
@@ -76,12 +75,12 @@ class SalesAnalysisWorkflow:
         logger.info("Initializing SalesAnalysisWorkflow")
         
         try:
-            # Initialize agents with comprehensive error handling
-            self.planning_agent = self._safe_agent_init(PlanningAgent, "Planning")
-            self.data_agent = self._safe_agent_init(DataAgent, "Data", db_path)
-            self.analytics_agent = self._safe_agent_init(AnalyticsAgent, "Analytics")
-            self.qa_agent = self._safe_agent_init(QAAgent, "QA")
-            self.insights_agent = self._safe_agent_init(InsightsAgent, "Insights")
+            # 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
             logger.info("Building workflow graph")
@@ -94,32 +93,6 @@ class SalesAnalysisWorkflow:
             logger.error(traceback.format_exc())
             raise
     
-    def _safe_agent_init(self, agent_class, agent_name, *args, **kwargs):
-        """
-        Safely initialize an agent with logging and error handling
-        
-        Args:
-            agent_class (type): The agent class to initialize
-            agent_name (str): Name of the agent for logging
-            *args: Positional arguments for agent initialization
-            **kwargs: Keyword arguments for agent initialization
-        
-        Returns:
-            Initialized agent
-        
-        Raises:
-            Exception if agent initialization fails
-        """
-        try:
-            logger.info(f"Initializing {agent_name} Agent")
-            agent = agent_class(*args, **kwargs)
-            logger.info(f"{agent_name} Agent initialized successfully")
-            return agent
-        except Exception as e:
-            logger.error(f"Failed to initialize {agent_name} Agent: {e}")
-            logger.error(traceback.format_exc())
-            raise RuntimeError(f"{agent_name} Agent initialization failed: {str(e)}")
-    
     def _build_workflow(self) -> StateGraph:
         """Build the workflow graph with comprehensive logging"""
         logger.info("Building workflow graph")
@@ -128,12 +101,277 @@ class SalesAnalysisWorkflow:
             # 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 node methods directly in the workflow
+            def planning_node(state):
+                """Planning agent node"""
+                try:
+                    logger.info("Entering Planning Node")
+                    
+                    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:
+                    logger.error(f"Planning node error: {e}")
+                    logger.error(traceback.format_exc())
+                    
+                    # Update state to error
+                    state["status"] = "error"
+                    state["error"] = f"Planning error: {str(e)}"
+                    
+                    return state
+            
+            def data_collection_node(state):
+                """Data collection agent node"""
+                try:
+                    logger.info("Entering Data Collection Node")
+                    
+                    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:
+                    logger.error(f"Data collection node error: {e}")
+                    logger.error(traceback.format_exc())
+                    
+                    # Update state to error
+                    state["status"] = "error"
+                    state["error"] = f"Data collection error: {str(e)}"
+                    
+                    return state
+            
+            def analysis_node(state):
+                """Analytics agent node"""
+                try:
+                    logger.info("Entering Analysis Node")
+                    
+                    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:
+                    logger.error(f"Analysis node error: {e}")
+                    logger.error(traceback.format_exc())
+                    
+                    # Update state to error
+                    state["status"] = "error"
+                    state["error"] = f"Analysis error: {str(e)}"
+                    
+                    return state
+            
+            def validation_node(state):
+                """QA agent node"""
+                try:
+                    logger.info("Entering Validation Node")
+                    
+                    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:
+                    logger.error(f"Validation node error: {e}")
+                    logger.error(traceback.format_exc())
+                    
+                    # Update state to error
+                    state["status"] = "error"
+                    state["error"] = f"Validation error: {str(e)}"
+                    
+                    return state
+            
+            def insights_node(state):
+                """Insights agent node"""
+                try:
+                    logger.info("Entering Insights Node")
+                    
+                    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:
+                    logger.error(f"Insights node error: {e}")
+                    logger.error(traceback.format_exc())
+                    
+                    # Update state to error
+                    state["status"] = "error"
+                    state["error"] = f"Insights error: {str(e)}"
+                    
+                    return state
+            
+            def handle_error_node(state):
+                """Error handling node"""
+                logger.error(f"Workflow error: {state['error']}")
+                
+                # Log the error
+                state["logs"].append({
+                    "timestamp": datetime.now().isoformat(),
+                    "step": "error",
+                    "message": f"Workflow error: {state['error']}"
+                })
+                
+                return state
+            
+            # Add nodes to the workflow
+            workflow.add_node("planning", planning_node)
+            workflow.add_node("data_collection", data_collection_node)
+            workflow.add_node("analysis", analysis_node)
+            workflow.add_node("validation", validation_node)
+            workflow.add_node("insights", insights_node)
+            workflow.add_node("handle_error", handle_error_node)
             
             # Define the workflow edges
             workflow.add_edge("planning", "data_collection")
@@ -142,14 +380,10 @@ class SalesAnalysisWorkflow:
             workflow.add_edge("validation", "insights")
             workflow.add_edge("insights", END)
             
-            # Add error handling
-            workflow.add_node("handle_error", self.handle_error_node)
-            
             # Add conditional edges for error handling
             def route_on_error(state):
-                logger.warning(f"Routing to error handler. Current status: {state['status']}")
                 return "handle_error" if state["status"] == "error" else None
-                
+            
             workflow.add_conditional_edges("planning", route_on_error)
             workflow.add_conditional_edges("data_collection", route_on_error)
             workflow.add_conditional_edges("analysis", route_on_error)
@@ -208,12 +442,6 @@ class SalesAnalysisWorkflow:
             
             return error_state
 
-    # Keep the existing node methods, adding more logging 
-    # (planning_node, data_collection_node, etc.) would follow the same pattern
-    # of adding comprehensive logging and error tracking
-
-# Keep the rest of the existing implementation the same
-
 # For testing
 if __name__ == "__main__":
     # Configure logging to show all details
@@ -235,8 +463,8 @@ if __name__ == "__main__":
         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")
+            print(f"Generated {len(result.get('insight_cards', {}))} insight cards")
+            print(f"Generated {len(result.get('visualizations', []))} visualizations")
     
     except Exception as e:
         print(f"Test run failed: {e}")