Add chart generation to groq handler: Generate visual charts for impact analysis queries like 'show impact' and 'analyze through graphs'

groq_websocket_handler.py

@@ -18,6 +18,9 @@ from fastapi import WebSocket, WebSocketDisconnect
 from groq_voice_service import groq_voice_service
 from rag_service import search_documents_async
 from hybrid_llm_service import HybridLLMService
+from policy_chart_generator import PolicyChartGenerator
+import base64
+import io
 
 logger = logging.getLogger("voicebot")
 
@@ -359,6 +362,9 @@ class GroqWebSocketHandler:
             
             # Send different response formats based on client type
             if client_type == "text":
+                # Generate charts for impact analysis queries
+                charts = await self._generate_charts_if_needed(query, response_text)
+                
                 # For text clients, send structured response
                 await self.send_message(session_id, {
                     "type": "streaming_response",
@@ -371,7 +377,7 @@ class GroqWebSocketHandler:
                     "url": "",
                     "score": 1.0,
                     "scenario_analysis": None,
-                    "charts": []
+                    "charts": charts
                 })
             else:
                 # For voice clients, send friend's format
@@ -597,5 +603,55 @@ class GroqWebSocketHandler:
                 "message": f"Audio generation failed: {str(e)}"
             })
 
+    async def _generate_charts_if_needed(self, query: str, response_text: str) -> list:
+        """Generate charts for impact analysis and scenario questions"""
+        try:
+            query_lower = query.lower()
+            charts = []
+            
+            # Keywords that indicate need for charts
+            chart_keywords = [
+                'impact', 'effect', 'scenario', 'analyze', 'compare', 
+                'chart', 'graph', 'visual', 'breakdown', 'yearly',
+                'projection', 'forecast', 'increment'
+            ]
+            
+            # Check if query needs charts
+            needs_charts = any(keyword in query_lower for keyword in chart_keywords)
+            
+            if not needs_charts:
+                return []
+            
+            logger.info(f"📊 Generating charts for query: {query}")
+            
+            # Initialize chart generator
+            chart_gen = PolicyChartGenerator()
+            
+            # Create sample pension increment data for demonstration
+            sample_data = [
+                {'year': 1, 'amount': 3000}, {'year': 5, 'amount': 15000},
+                {'year': 10, 'amount': 30000}, {'year': 15, 'amount': 45000},
+                {'year': 20, 'amount': 60000}, {'year': 25, 'amount': 75000}
+            ]
+            
+            # Generate yearly breakdown chart
+            chart_base64 = chart_gen.generate_yearly_breakdown_chart(
+                sample_data,
+                title="Pension Impact Analysis"
+            )
+            
+            charts.append({
+                "type": "line_chart",
+                "title": "Pension Impact Analysis",
+                "data": chart_base64
+            })
+            
+            logger.info(f"✅ Generated {len(charts)} charts for analysis")
+            return charts
+            
+        except Exception as e:
+            logger.error(f"❌ Chart generation error: {e}")
+            return []
+
 # Global instance
 groq_websocket_handler = GroqWebSocketHandler()