Update agents/memory.py

agents/memory.py

@@ -1,27 +1,34 @@
 class Memory:
+    """
+    Stores conversation history and image context for one user session.
+    """
 
     def __init__(self):
 
-        # Conversation history
+        # Conversation
         self.chat_history = []
 
-        # Uploaded files
+        # Uploaded image paths
         self.uploaded_files = []
 
-        # Current image (original or latest edited)
-        self.last_used_image = None
+        # Current working image
+        self.current_image = None
 
-        # Full analysis dictionary
-        self.image_context = None
+        # Latest edited image
+        self.last_used_image = None
 
-        # Current analysis
+        # Full image analysis
         self.analysis = None
 
-    # --------------------------
-    # Chat History
-    # --------------------------
+        # Natural language description
+        self.image_context = ""
+
+    # =====================================================
+    # CHAT
+    # =====================================================
 
     def add_message(self, role, content):
+
         self.chat_history.append({
             "role": role,
             "content": content
@@ -30,44 +37,81 @@ class Memory:
     def get_history(self):
         return self.chat_history
 
-    # --------------------------
-    # Uploaded Files
-    # --------------------------
+    # =====================================================
+    # FILES
+    # =====================================================
 
     def add_files(self, files):
-        self.uploaded_files.extend(files)
 
-    # --------------------------
-    # Image Context
-    # --------------------------
+        self.uploaded_files = files
 
-    def set_image_context(self, analysis):
-        """
-        Store the complete image analysis.
-        """
-        self.image_context = analysis
-        self.analysis = analysis
+        if files:
 
-    def get_image_context(self):
-        return self.image_context
+            self.current_image = files[0]
+            self.last_used_image = files[0]
+
+    def get_files(self):
+
+        return self.uploaded_files
 
-    # --------------------------
-    # Current Image
-    # --------------------------
+    # =====================================================
+    # IMAGE
+    # =====================================================
+
+    def set_current_image(self, image_path):
+
+        self.current_image = image_path
+
+    def get_current_image(self):
+
+        return self.current_image
 
     def set_last_image(self, image_path):
+
         self.last_used_image = image_path
+        self.current_image = image_path
 
     def get_last_image(self):
+
         return self.last_used_image
 
-    # --------------------------
-    # Clear Memory
-    # --------------------------
+    # =====================================================
+    # ANALYSIS
+    # =====================================================
+
+    def set_analysis(self, analysis):
+
+        self.analysis = analysis
+
+        if isinstance(analysis, dict):
+
+            self.image_context = analysis.get(
+                "description",
+                ""
+            )
+
+    def get_analysis(self):
+
+        return self.analysis
+
+    def get_image_context(self):
+
+        return self.image_context
+
+    # =====================================================
+    # CLEAR
+    # =====================================================
 
     def clear(self):
-        self.chat_history.clear()
-        self.uploaded_files.clear()
+
+        self.chat_history = []
+
+        self.uploaded_files = []
+
+        self.current_image = None
+
         self.last_used_image = None
-        self.image_context = None
-        self.analysis = None
+
+        self.analysis = None
+
+        self.image_context = ""