Update extensions/visualization.py

extensions/visualization.py

@@ -1,20 +1,21 @@
 """
-Data Visualization Extension
+Enhanced Data Visualization Extension
 Create charts, graphs, and visualizations from data
-Enhanced for stock market data integration
+Now with better state management and orchestrator integration
 """
 
 from base_extension import BaseExtension
 from google.genai import types
-from typing import Dict, Any, List
+from typing import Dict, Any, List, Optional
 import json
 import matplotlib
-matplotlib.use('Agg')  # Use non-GUI backend
+matplotlib.use('Agg')
 import matplotlib.pyplot as plt
 import io
 import base64
 from pathlib import Path
 import numpy as np
+import datetime
 
 
 class VisualizationExtension(BaseExtension):
@@ -35,6 +36,10 @@ class VisualizationExtension(BaseExtension):
     def icon(self) -> str:
         return "📊"
     
+    @property
+    def version(self) -> str:
+        return "2.0.0"
+    
     def get_system_context(self) -> str:
         return """
 You have access to a Data Visualization system for creating charts and graphs.
@@ -107,7 +112,33 @@ create_bar_chart(
     def _get_default_state(self) -> Dict[str, Any]:
         return {
             "charts": [],
-            "output_dir": "visualizations"
+            "output_dir": "visualizations",
+            "total_created": 0,
+            "created_at": datetime.datetime.now().isoformat(),
+            "last_updated": datetime.datetime.now().isoformat()
+        }
+    
+    def get_state_summary(self, user_id: str) -> Optional[str]:
+        """Provide state summary for system prompt"""
+        state = self.get_state(user_id)
+        chart_count = len(state.get("charts", []))
+        if chart_count > 0:
+            return f"{chart_count} visualizations created this session"
+        return None
+    
+    def get_metrics(self, user_id: str) -> Dict[str, Any]:
+        """Provide usage metrics"""
+        state = self.get_state(user_id)
+        charts = state.get("charts", [])
+        chart_types = {}
+        for chart in charts:
+            chart_type = chart.get("type", "unknown")
+            chart_types[chart_type] = chart_types.get(chart_type, 0) + 1
+        
+        return {
+            "total_created": state.get("total_created", 0),
+            "current_session": len(charts),
+            "by_type": chart_types
         }
     
     def get_tools(self) -> List[types.Tool]:
@@ -220,7 +251,6 @@ create_bar_chart(
         output_dir.mkdir(exist_ok=True)
         
         # Generate filename
-        import datetime
         timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
         safe_title = "".join(c for c in title if c.isalnum() or c in (' ', '-', '_')).strip()
         safe_title = safe_title.replace(' ', '_')[:50]
@@ -247,22 +277,17 @@ create_bar_chart(
             "timestamp": timestamp
         }
         state["charts"].append(chart_info)
+        state["total_created"] = state.get("total_created", 0) + 1
         self.update_state(user_id, state)
         
-        # Return as tuple: (filepath_string, base64_string)
-        filepath_str = str(filepath)
-        return filepath_str, img_base64
+        return str(filepath), img_base64
     
-    def handle_tool_call(self, user_id: str, tool_name: str, args: Dict[str, Any]) -> Any:
-        # Ensure state is initialized
-        if user_id not in self.state:
-            self.initialize_state(user_id)
-        
+    def _execute_tool(self, user_id: str, tool_name: str, args: Dict[str, Any]) -> Any:
+        """Execute tool logic"""
         try:
             if tool_name == "create_line_chart":
                 fig, ax = plt.subplots(figsize=(12, 7))
                 
-                # Plot each series
                 data = args["data"]
                 
                 print(f"📊 Creating line chart with {len(data)} series")
@@ -281,15 +306,18 @@ create_bar_chart(
                 ax.legend(fontsize=10)
                 ax.grid(True, alpha=0.3, linestyle='--')
                 
-                # Rotate x-axis labels for readability (especially for dates)
                 plt.xticks(rotation=45, ha='right')
-                
-                # Tight layout to prevent label cutoff
                 plt.tight_layout()
                 
-                # Save chart and get both filepath and base64
                 filepath, img_base64 = self._save_chart(fig, user_id, "line_chart", args["title"])
                 
+                # Log activity
+                self.log_activity(user_id, "chart_created", {
+                    "type": "line",
+                    "title": args["title"],
+                    "series_count": len(data)
+                })
+                
                 print(f"✅ Line chart saved: {filepath}")
                 
                 return {
@@ -320,7 +348,6 @@ create_bar_chart(
                 ax.set_title(args["title"], fontsize=14, fontweight='bold')
                 ax.grid(True, alpha=0.3, axis='y')
                 
-                # Rotate x-labels if needed
                 if len(categories) > 5:
                     plt.xticks(rotation=45, ha='right')
                 
@@ -328,6 +355,12 @@ create_bar_chart(
                 
                 filepath, img_base64 = self._save_chart(fig, user_id, "bar_chart", args["title"])
                 
+                self.log_activity(user_id, "chart_created", {
+                    "type": "bar",
+                    "title": args["title"],
+                    "categories": len(categories)
+                })
+                
                 return {
                     "success": True,
                     "message": f"Bar chart created: {args['title']}",
@@ -360,6 +393,12 @@ create_bar_chart(
                 
                 filepath, img_base64 = self._save_chart(fig, user_id, "scatter_plot", args["title"])
                 
+                self.log_activity(user_id, "chart_created", {
+                    "type": "scatter",
+                    "title": args["title"],
+                    "points": len(x_vals)
+                })
+                
                 return {
                     "success": True,
                     "message": f"Scatter plot created: {args['title']}",
@@ -384,7 +423,6 @@ create_bar_chart(
                     textprops={'fontsize': 11, 'fontweight': 'bold'}
                 )
                 
-                # Enhance text
                 for text in texts:
                     text.set_fontsize(11)
                 for autotext in autotexts:
@@ -396,6 +434,12 @@ create_bar_chart(
                 
                 filepath, img_base64 = self._save_chart(fig, user_id, "pie_chart", args["title"])
                 
+                self.log_activity(user_id, "chart_created", {
+                    "type": "pie",
+                    "title": args["title"],
+                    "slices": len(labels)
+                })
+                
                 return {
                     "success": True,
                     "message": f"Pie chart created: {args['title']}",
@@ -410,7 +454,7 @@ create_bar_chart(
                 
                 return {
                     "total_charts": len(charts),
-                    "charts": charts
+                    "charts": charts[-10:]  # Last 10 charts
                 }
             
         except Exception as e:
@@ -427,4 +471,27 @@ create_bar_chart(
     
     def on_enable(self, user_id: str) -> str:
         self.initialize_state(user_id)
-        return "📊 Data Visualization enabled! I can create charts from stock data and other numerical data. Just ask me to visualize something!"
+        return "📊 Data Visualization enabled! I can create charts from stock data and other numerical data. Just ask me to visualize something!"
+    
+    def on_disable(self, user_id: str) -> str:
+        state = self.get_state(user_id)
+        total = state.get("total_created", 0)
+        return f"📊 Data Visualization disabled. You created {total} visualizations this session."
+    
+    def health_check(self, user_id: str) -> Dict[str, Any]:
+        """Check extension health"""
+        try:
+            import matplotlib
+            return {
+                "healthy": True,
+                "extension": self.name,
+                "version": self.version,
+                "matplotlib_available": True
+            }
+        except ImportError:
+            return {
+                "healthy": False,
+                "extension": self.name,
+                "version": self.version,
+                "issues": ["matplotlib library not installed"]
+            }