Add 89999999999999999999

enhanced_websocket_handler.py

@@ -11,12 +11,15 @@ import uuid
 import tempfile
 import base64
 from pathlib import Path
+import io
+import matplotlib.pyplot as plt
 
 from llm_service import create_graph, create_basic_graph
 from lancedb_service import lancedb_service
 from hybrid_llm_service import HybridLLMService
 from voice_service import voice_service
 from rag_service import search_government_docs
+from policy_chart_generator import PolicyChartGenerator
 
 # Initialize hybrid LLM service
 hybrid_llm_service = HybridLLMService()
@@ -257,7 +260,8 @@ async def handle_text_message(websocket: WebSocket, data: dict, session_data: di
                     "date": chunk.get("date", ""),
                     "url": chunk.get("url", ""),
                     "score": chunk.get("score", 1.0),
-                    "scenario_analysis": chunk.get("scenario_analysis", None)
+                    "scenario_analysis": chunk.get("scenario_analysis", None),
+                    "charts": chunk.get("charts", [])
                 })
             # Optionally, aggregate or select the best chunk for final response
             # Here, just use the first chunk for context update and provider
@@ -451,6 +455,27 @@ async def get_hybrid_response(user_message: str, context: str, config: dict, kno
                     'inflation': 0.05
                 }
                 scenario_result = run_scenario_analysis(params)
+                # Generate charts for scenario_result
+                chart_gen = PolicyChartGenerator()
+                charts = []
+                # Example: line chart for yearly results
+                if "yearly_results" in scenario_result:
+                    years = [r['year'] for r in scenario_result['yearly_results']]
+                    base_costs = [r['base_cost'] for r in scenario_result['yearly_results']]
+                    scenario_costs = [r['scenario_cost'] for r in scenario_result['yearly_results']]
+                    # Generate chart and append to charts list
+                    fig, ax = plt.subplots()
+                    ax.plot(years, base_costs, label='Base Cost')
+                    ax.plot(years, scenario_costs, label='Scenario Cost')
+                    ax.legend()
+                    ax.set_title('Scenario Analysis: Cost Over Years')
+                    buf = io.BytesIO()
+                    fig.savefig(buf, format='png')
+                    buf.seek(0)
+                    chart_base64 = base64.b64encode(buf.read()).decode('utf-8')
+                    plt.close(fig)
+                    charts.append({"type": "line_chart", "data": chart_base64})
+                scenario_result["charts"] = charts
             else:
                 scenario_result = None
             for doc in docs:
@@ -463,7 +488,8 @@ async def get_hybrid_response(user_message: str, context: str, config: dict, kno
                     "date": doc.get("date", ""),
                     "url": doc.get("url", ""),
                     "score": doc.get("score", 1.0),
-                    "scenario_analysis": scenario_result
+                    "scenario_analysis": scenario_result,
+                    "charts": scenario_result.get("charts", []) if scenario_result else []
                 }
                 yield response_obj
         else: