Update storage/storage_manager.py

storage/storage_manager.py

@@ -5,36 +5,184 @@ from pathlib import Path
 from datetime import datetime
 import shutil
 import streamlit as st
-import base64
 
 class UserStorageManager:
     def __init__(self, storage_paths):
         self.paths = storage_paths
-        # Validate storage access
-        if not self._validate_storage():
-            raise RuntimeError("Cannot access persistent storage. Please ensure /data directory is available and writable.")
         self.ensure_directories()
 
-    def _validate_storage(self):
-        """Validate that we can access and write to storage"""
+    def ensure_directories(self):
+        """Ensure all required directories exist"""
+        for path in self.paths.values():
+            path.mkdir(parents=True, exist_ok=True)
+
+    def save_chat(self, chat_data, images=None, filename=None):
+        """Save chat with associated images"""
+        try:
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            chat_id = f"chat_{timestamp}"
+            
+            if filename:
+                chat_id = filename.replace('.json', '')
+            
+            # Save chat data
+            chat_file = self.paths["chats"] / f"{chat_id}.json"
+            
+            # Add metadata
+            chat_metadata = {
+                'id': chat_id,
+                'timestamp': timestamp,
+                'data': chat_data
+            }
+            
+            with open(chat_file, "w") as f:
+                json.dump(chat_metadata, f)
+            
+            # Save associated images
+            if images:
+                image_dir = self.paths["images"] / chat_id
+                image_dir.mkdir(exist_ok=True)
+                
+                if isinstance(images, list):
+                    for idx, img_data in enumerate(images):
+                        img_file = image_dir / f"image_{idx}.png"
+                        with open(img_file, "wb") as f:
+                            f.write(img_data)
+                else:
+                    img_file = image_dir / "image_0.png"
+                    with open(img_file, "wb") as f:
+                        f.write(images)
+                
+                # Update chat data with image references
+                chat_metadata['image_paths'] = [str(p) for p in image_dir.glob("*.png")]
+                with open(chat_file, "w") as f:
+                    json.dump(chat_metadata, f)
+
+            return chat_id
+            
+        except Exception as e:
+            st.error(f"Error saving chat: {str(e)}")
+            return None
+
+    def get_all_chats(self):
+        """Get all user's chats with metadata"""
+        try:
+            chats = []
+            for chat_file in sorted(self.paths["chats"].glob("*.json"), reverse=True):
+                with open(chat_file, "r") as f:
+                    chat_data = json.load(f)
+                    # Load associated images if they exist
+                    chat_id = chat_data.get('id', chat_file.stem)
+                    image_dir = self.paths["images"] / chat_id
+                    if image_dir.exists():
+                        images = []
+                        for img_file in sorted(image_dir.glob("*.png")):
+                            with open(img_file, "rb") as f:
+                                images.append(f.read())
+                        chat_data['images'] = images
+                    chats.append(chat_data)
+            return chats
+        except Exception as e:
+            st.error(f"Error loading chats: {str(e)}")
+            return []
+
+    def get_chat(self, chat_id):
+        """Get specific chat with images"""
+        try:
+            chat_file = self.paths["chats"] / f"{chat_id}.json"
+            if not chat_file.exists():
+                return None
+                
+            with open(chat_file, "r") as f:
+                chat_data = json.load(f)
+                
+            # Load associated images
+            image_dir = self.paths["images"] / chat_id
+            if image_dir.exists():
+                images = []
+                for img_file in sorted(image_dir.glob("*.png")):
+                    with open(img_file, "rb") as f:
+                        images.append(f.read())
+                chat_data['images'] = images
+                
+            return chat_data
+        except Exception as e:
+            st.error(f"Error loading chat: {str(e)}")
+            return None
+
+    def save_context(self, context_data):
+        """Save conversation context"""
+        try:
+            context_file = self.paths["context"] / "context.json"
+            with open(context_file, "w") as f:
+                json.dump(context_data, f)
+        except Exception as e:
+            st.error(f"Error saving context: {str(e)}")
+
+    def get_context(self):
+        """Get saved conversation context"""
+        try:
+            context_file = self.paths["context"] / "context.json"
+            if context_file.exists():
+                with open(context_file, "r") as f:
+                    return json.load(f)
+            return None
+        except Exception as e:
+            st.error(f"Error loading context: {str(e)}")
+            return None
+
+    def clear_context(self):
+        """Clear conversation context"""
+        try:
+            context_file = self.paths["context"] / "context.json"
+            if context_file.exists():
+                context_file.unlink()
+        except Exception as e:
+            st.error(f"Error clearing context: {str(e)}")
+
+    def delete_chat(self, chat_id):
+        """Delete a chat and its associated images"""
         try:
-            test_file = Path("/data/storage_test")
-            test_file.touch()
-            test_file.unlink()
+            # Delete chat file
+            chat_file = self.paths["chats"] / f"{chat_id}.json"
+            if chat_file.exists():
+                chat_file.unlink()
+            
+            # Delete associated images
+            image_dir = self.paths["images"] / chat_id
+            if image_dir.exists():
+                shutil.rmtree(image_dir)
+            
             return True
         except Exception as e:
-            st.error(f"Storage validation failed: {str(e)}")
+            st.error(f"Error deleting chat: {str(e)}")
             return False
 
-    def ensure_directories(self):
-        """Ensure all required directories exist"""
+    def get_storage_stats(self):
+        """Get storage usage statistics"""
         try:
-            for path in self.paths.values():
-                path.mkdir(parents=True, exist_ok=True)
-                # Test write access
-                test_file = path / ".write_test"
-                test_file.touch()
-                test_file.unlink()
+            stats = {
+                'total_chats': 0,
+                'total_images': 0,
+                'storage_used': 0  # in bytes
+            }
+            
+            # Count chats
+            stats['total_chats'] = len(list(self.paths["chats"].glob("*.json")))
+            
+            # Count images
+            for image_dir in self.paths["images"].glob("*"):
+                if image_dir.is_dir():
+                    stats['total_images'] += len(list(image_dir.glob("*.png")))
+            
+            # Calculate storage used
+            for path_type, path in self.paths.items():
+                if path.exists():
+                    for file_path in path.rglob("*"):
+                        if file_path.is_file():
+                            stats['storage_used'] += file_path.stat().st_size
+            
+            return stats
         except Exception as e:
-            st.error(f"Failed to create or access storage directories: {str(e)}")
-            raise RuntimeError("Storage directories are not accessible")
+            st.error(f"Error getting storage stats: {str(e)}")
+            return None