Create storage/storage_manager.py

storage/storage_manager.py

@@ -0,0 +1,97 @@
+# storage/storage_manager.py
+
+import json
+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
+        self.ensure_directories()
+
+    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):
+        """Save chat with associated images"""
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+        chat_id = f"chat_{timestamp}"
+        
+        # Save chat data
+        chat_file = self.paths["chats"] / f"{chat_id}.json"
+        with open(chat_file, "w") as f:
+            json.dump(chat_data, f)
+        
+        # Save associated images
+        if images:
+            image_dir = self.paths["images"] / chat_id
+            image_dir.mkdir(exist_ok=True)
+            
+            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)
+            
+            # Update chat data with image references
+            chat_data["image_paths"] = [str(p) for p in image_dir.glob("*.png")]
+            with open(chat_file, "w") as f:
+                json.dump(chat_data, f)
+
+        return chat_id
+
+    def get_all_chats(self):
+        """Get all user's chats with metadata"""
+        chats = []
+        for chat_file in self.paths["chats"].glob("*.json"):
+            with open(chat_file, "r") as f:
+                chat_data = json.load(f)
+                chat_data["id"] = chat_file.stem
+                chats.append(chat_data)
+        
+        # Sort by timestamp
+        return sorted(chats, key=lambda x: x.get("timestamp", ""), reverse=True)
+
+    def get_chat(self, chat_id):
+        """Get specific chat with images"""
+        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
+
+    def save_context(self, context_data):
+        """Save conversation context"""
+        context_file = self.paths["context"] / "context.json"
+        with open(context_file, "w") as f:
+            json.dump(context_data, f)
+
+    def get_context(self):
+        """Get saved conversation context"""
+        context_file = self.paths["context"] / "context.json"
+        if context_file.exists():
+            with open(context_file, "r") as f:
+                return json.load(f)
+        return None
+
+    def clear_context(self):
+        """Clear conversation context"""
+        context_file = self.paths["context"] / "context.json"
+        if context_file.exists():
+            context_file.unlink()