Update storage/storage_manager.py

storage/storage_manager.py

@@ -32,6 +32,8 @@ class UserStorageManager:
                 chat_title = generate_chat_name(chat_data, claude_service)
             
             chat_id = filename or f"chat_{formatted_timestamp}"
+            if filename:
+                chat_id = filename.replace('.json', '')
             
             # Save chat data
             chat_file = self.paths["chats"] / f"{chat_id}.json"
@@ -44,12 +46,48 @@ class UserStorageManager:
                 'title': chat_title or f"Trading Analysis {formatted_timestamp}",
                 'data': chat_data
             }
+            
+            # Save chat metadata
+            with open(chat_file, "w") as f:
+                json.dump(chat_metadata, f, indent=4)
+            
+            # 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 metadata 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, indent=4)
+
+            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):
+            # Sort chat files by modification time for most recent first
+            chat_files = sorted(
+                self.paths["chats"].glob("*.json"),
+                key=lambda x: x.stat().st_mtime,
+                reverse=True
+            )
+            for chat_file in chat_files:
                 try:
                     with open(chat_file, "r") as f:
                         chat_data = json.load(f)
@@ -100,7 +138,7 @@ class UserStorageManager:
         try:
             context_file = self.paths["context"] / "context.json"
             with open(context_file, "w") as f:
-                json.dump(context_data, f)
+                json.dump(context_data, f, indent=4)
         except Exception as e:
             st.error(f"Error saving context: {str(e)}")