Update storage/storage_manager.py

storage/storage_manager.py

@@ -13,14 +13,17 @@ class UserStorageManager:
 
     def ensure_directories(self):
         """Ensure all required directories exist"""
-        for path in self.paths.values():
-            path.mkdir(parents=True, exist_ok=True)
+        try:
+            for path in self.paths.values():
+                path.mkdir(parents=True, exist_ok=True)
+        except Exception as e:
+            st.error(f"Error creating directories: {str(e)}")
 
     def save_chat(self, chat_data, images=None, filename=None):
         """Save chat with associated images"""
         try:
-            timestamp = datetime.now().isoformat()  # Use ISO format timestamp
-            formatted_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")  # For filename
+            timestamp = datetime.now().isoformat()
+            formatted_timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
             chat_id = f"chat_{formatted_timestamp}"
             
             if filename:
@@ -32,31 +35,60 @@ class UserStorageManager:
             # Add metadata
             chat_metadata = {
                 'id': chat_id,
-                'timestamp': timestamp,  # ISO format for proper datetime parsing
-                'formatted_timestamp': formatted_timestamp,  # For display
+                'timestamp': timestamp,
+                'formatted_timestamp': formatted_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)
+                try:
+                    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)
+                except Exception as e:
+                    st.warning(f"Error loading chat {chat_file.name}: {str(e)}")
+                    continue
             return chats
         except Exception as e:
             st.error(f"Error loading chats: {str(e)}")
@@ -68,10 +100,10 @@ class UserStorageManager:
             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():
@@ -80,7 +112,7 @@ class UserStorageManager:
                     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)}")