Update src/utils/cocoon_manager.py

src/utils/cocoon_manager.py

@@ -1,137 +1,140 @@
-from datetime import datetime
-import os
-import json
-import logging
-from typing import List, Dict, Any, Optional
-
-logger = logging.getLogger(__name__)
-
-class CocoonManager:
-    """Manages Codette's cocoon data storage and retrieval"""
-    
-    def __init__(self, base_dir: str = "./cocoons"):
-        self.base_dir = base_dir
-        self.cocoon_data = []
-        self.quantum_state = {"coherence": 0.5}
-        self._ensure_cocoon_dir()
-        
-    def _ensure_cocoon_dir(self):
-        """Ensure cocoon directory exists"""
-        os.makedirs(self.base_dir, exist_ok=True)
-        
-    def load_cocoons(self) -> None:
-        """Load all cocoon data from files"""
-        try:
-            # Ensure directory exists
-            os.makedirs(self.base_dir, exist_ok=True)
-            
-            cocoon_files = [
-                f for f in os.listdir(self.base_dir) 
-                if f.endswith('.cocoon')
-            ]
-            logger.info(f"Found {len(cocoon_files)} cocoon files in {self.base_dir}")
-            
-            loaded_data = []
-            for fname in cocoon_files:
-                try:
-                    full_path = os.path.join(self.base_dir, fname)
-                    with open(full_path, 'r') as f:
-                        cocoon = json.load(f)
-                        if self._validate_cocoon(cocoon):
-                            loaded_data.append(cocoon)
-                            # Update quantum state if present and most recent
-                            current_quantum_state = cocoon.get('data', {}).get('quantum_state')
-                            if current_quantum_state:
-                                # Convert list quantum state to dict if needed
-                                if isinstance(current_quantum_state, list):
-                                    current_quantum_state = {
-                                        "coherence": sum(current_quantum_state) / len(current_quantum_state)
-                                        if current_quantum_state else 0.5
-                                    }
-                                # Update if we don't have state yet or if this cocoon is more recent
-                                if (not isinstance(self.quantum_state, dict) or
-                                    not self.quantum_state.get('coherence') or
-                                    not loaded_data[:-1] or
-                                    cocoon['timestamp'] > loaded_data[-2]['timestamp']):
-                                    self.quantum_state = current_quantum_state.copy() if isinstance(current_quantum_state, dict) else {"coherence": 0.5}
-                except Exception as e:
-                    logger.error(f"Error loading cocoon {fname}: {e}")
-                    continue
-                    
-            # Sort by timestamp and handle empty timestamps
-            self.cocoon_data = sorted(
-                [c for c in loaded_data if isinstance(c, dict) and self._validate_cocoon(c)],
-                key=lambda x: x.get('timestamp', '0'),
-                reverse=True
-            )
-            logger.info(f"Successfully loaded {len(self.cocoon_data)} valid cocoons")
-            
-            logger.info(
-                f"Loaded {len(self.cocoon_data)} cocoons "
-                f"with quantum coherence {self.quantum_state.get('coherence', 0.5)}"
-            )
-            
-        except Exception as e:
-            logger.error(f"Error loading cocoons: {e}")
-            self.cocoon_data = []
-            # Ensure we have a valid quantum state
-            if not isinstance(self.quantum_state, dict) or 'coherence' not in self.quantum_state:
-                self.quantum_state = {"coherence": 0.5}
-            
-    def _validate_cocoon(self, cocoon: Dict[str, Any]) -> bool:
-        """Validate cocoon data structure"""
-        required_fields = ['timestamp', 'data']
-        return all(field in cocoon for field in required_fields)
-        
-    def save_cocoon(
-        self, 
-        data: Dict[str, Any],
-        cocoon_type: str = "codette"
-    ) -> bool:
-        """Save new cocoon data to file"""
-        try:
-            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-            filename = f"{cocoon_type}_cocoon_{timestamp}.cocoon"
-            filepath = os.path.join(self.base_dir, filename)
-            
-            # If data has its own quantum state, update our state
-            if "quantum_state" in data:
-                self.quantum_state = data["quantum_state"]
-            
-            cocoon = {
-                "timestamp": timestamp,
-                "data": {
-                    **data,
-                    "quantum_state": self.quantum_state.copy()
-                }
-            }
-            
-            with open(filepath, 'w') as f:
-                json.dump(cocoon, f, indent=2)
-                
-            self.cocoon_data.insert(0, cocoon)
-            logger.info(f"Saved cocoon: {filename}")
-            return True
-            
-        except Exception as e:
-            logger.error(f"Error saving cocoon: {e}")
-            return False
-            
-    def get_latest_quantum_state(self) -> Dict[str, float]:
-        """Get the most recent quantum state"""
-        return self.quantum_state.copy()
-        
-    def get_latest_cocoons(self, limit: int = 5) -> List[Dict[str, Any]]:
-        """Get the most recent cocoons"""
-        return self.cocoon_data[:limit]
-        
-    def update_quantum_state(self, new_state: Dict[str, float]) -> None:
-        """Update the current quantum state and save it"""
-        self.quantum_state.update(new_state)
-        logger.info(f"Updated quantum state: {self.quantum_state}")
-        
-        # Save the new state in a cocoon
-        self.save_cocoon({
-            "type": "quantum_update",
-            "quantum_state": self.quantum_state.copy()
+from datetime import datetime
+import os
+import json
+import logging
+from typing import List, Dict, Any, Optional
+
+logger = logging.getLogger(__name__)
+
+class CocoonManager:
+    """Manages Codette's cocoon data storage and retrieval"""
+    
+    def __init__(self, base_dir: str = "./cocoons"):
+        self.base_dir = base_dir
+        self.cocoon_data = []
+        self.quantum_state = {"coherence": 0.5}
+        self._ensure_cocoon_dir()
+        
+    def _ensure_cocoon_dir(self):
+        """Ensure cocoon directory exists"""
+        os.makedirs(self.base_dir, exist_ok=True)
+        
+    def load_cocoons(self) -> None:
+        """Load all cocoon data from files"""
+        try:
+            # Ensure directory exists
+            os.makedirs(self.base_dir, exist_ok=True)
+            
+            cocoon_files = [
+                f for f in os.listdir(self.base_dir) 
+                if f.endswith('.cocoon')
+            ]
+            logger.info(f"Found {len(cocoon_files)} cocoon files in {self.base_dir}")
+            
+            loaded_data = []
+            for fname in cocoon_files:
+                try:
+                    full_path = os.path.join(self.base_dir, fname)
+                    with open(full_path, 'r') as f:
+                        cocoon = json.load(f)
+                        if self._validate_cocoon(cocoon):
+                            loaded_data.append(cocoon)
+                            # Update quantum state if present and most recent
+                            current_quantum_state = cocoon.get('data', {}).get('quantum_state')
+                            if current_quantum_state:
+                                # Convert list quantum state to dict if needed
+                                if isinstance(current_quantum_state, list):
+                                    current_quantum_state = {
+                                        "coherence": sum(current_quantum_state) / len(current_quantum_state)
+                                        if current_quantum_state else 0.5
+                                    }
+                                # Update if we don't have state yet or if this cocoon is more recent
+                                if (not isinstance(self.quantum_state, dict) or
+                                    not self.quantum_state.get('coherence') or
+                                    not loaded_data[:-1] or
+                                    cocoon['timestamp'] > loaded_data[-2]['timestamp']):
+                                    self.quantum_state = current_quantum_state.copy() if isinstance(current_quantum_state, dict) else {"coherence": 0.5}
+                except Exception as e:
+                    logger.error(f"Error loading cocoon {fname}: {e}")
+                    continue
+                    
+            # Sort by timestamp and handle empty timestamps
+            self.cocoon_data = sorted(
+                [c for c in loaded_data if isinstance(c, dict) and self._validate_cocoon(c)],
+                key=lambda x: x.get('timestamp', '0'),
+                reverse=True
+            )
+            logger.info(f"Successfully loaded {len(self.cocoon_data)} valid cocoons")
+            
+            logger.info(
+                f"Loaded {len(self.cocoon_data)} cocoons "
+                f"with quantum coherence {self.quantum_state.get('coherence', 0.5)}"
+            )
+            
+        except Exception as e:
+            logger.error(f"Error loading cocoons: {e}")
+            self.cocoon_data = []
+            # Ensure we have a valid quantum state
+            if not isinstance(self.quantum_state, dict) or 'coherence' not in self.quantum_state:
+                self.quantum_state = {"coherence": 0.5}
+            
+    def _validate_cocoon(self, cocoon: Dict[str, Any]) -> bool:
+        """Validate cocoon data structure"""
+        required_fields = ['timestamp', 'data']
+        return all(field in cocoon for field in required_fields)
+        
+    def save_cocoon(
+        self, 
+        data: Dict[str, Any],
+        cocoon_type: str = "codette"
+    ) -> bool:
+        """Save new cocoon data to file"""
+        try:
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            filename = f"{cocoon_type}_cocoon_{timestamp}.cocoon"
+            filepath = os.path.join(self.base_dir, filename)
+            
+            # If data has its own quantum state, update our state
+            if "quantum_state" in data:
+                self.quantum_state = data["quantum_state"]
+            
+            cocoon = {
+                "timestamp": timestamp,
+                "data": {
+                    **data,
+                    "quantum_state": self.quantum_state.copy()
+                }
+            }
+            
+            with open(filepath, 'w') as f:
+                json.dump(cocoon, f, indent=2)
+                
+            self.cocoon_data.insert(0, cocoon)
+            logger.info(f"Saved cocoon: {filename}")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Error saving cocoon: {e}")
+            return False
+            
+    def get_latest_quantum_state(self) -> Dict[str, float]:
+        """Get the most recent quantum state"""
+        # Ensure quantum_state is always a proper dict
+        if not isinstance(self.quantum_state, dict):
+            self.quantum_state = {"coherence": 0.5}
+        return self.quantum_state.copy()
+        
+    def get_latest_cocoons(self, limit: int = 5) -> List[Dict[str, Any]]:
+        """Get the most recent cocoons"""
+        return self.cocoon_data[:limit]
+        
+    def update_quantum_state(self, new_state: Dict[str, float]) -> None:
+        """Update the current quantum state and save it"""
+        self.quantum_state.update(new_state)
+        logger.info(f"Updated quantum state: {self.quantum_state}")
+        
+        # Save the new state in a cocoon
+        self.save_cocoon({
+            "type": "quantum_update",
+            "quantum_state": self.quantum_state.copy()
         })