lollms-client 0.29.2__py3-none-any.whl → 0.31.0__py3-none-any.whl

lollms_client/lollms_discussion.py

@@ -168,6 +168,7 @@ def create_dynamic_models(
 
         message_metadata = Column(JSON, nullable=True, default=dict)
         images = Column(JSON, nullable=True, default=list)
+        active_images = Column(JSON, nullable=True, default=list) # New: List of booleans for image activation state
         created_at = Column(DateTime, default=datetime.utcnow)
 
         @declared_attr
@@ -217,8 +218,9 @@ class LollmsDataManager:
             with self.engine.connect() as connection:
                 print("Checking for database schema upgrades...")
                 
+                # Discussions table migration
                 cursor = connection.execute(text("PRAGMA table_info(discussions)"))
-                columns = [row[1] for row in cursor.fetchall()]
+                columns = {row[1] for row in cursor.fetchall()}
                 
                 if 'pruning_summary' not in columns:
                     print("  -> Upgrading 'discussions' table: Adding 'pruning_summary' column.")
@@ -247,7 +249,15 @@ class LollmsDataManager:
                 if 'memory' not in columns:
                     print("  -> Upgrading 'discussions' table: Adding 'memory' column.")
                     connection.execute(text("ALTER TABLE discussions ADD COLUMN memory TEXT"))
-                
+
+                # Messages table migration
+                cursor = connection.execute(text("PRAGMA table_info(messages)"))
+                columns = {row[1] for row in cursor.fetchall()}
+
+                if 'active_images' not in columns:
+                    print("  -> Upgrading 'messages' table: Adding 'active_images' column.")
+                    connection.execute(text("ALTER TABLE messages ADD COLUMN active_images TEXT"))
+
                 print("Database schema is up to date.")
                 connection.commit() 
 
@@ -356,6 +366,73 @@ class LollmsMessage:
     def __repr__(self) -> str:
         """Provides a developer-friendly representation of the message."""
         return f"<LollmsMessage id={self.id} sender='{self.sender}'>"
+
+    def get_all_images(self) -> List[Dict[str, Union[str, bool]]]:
+        """
+        Returns a list of all images associated with this message, including their activation status.
+        
+        Returns:
+            A list of dictionaries, where each dictionary represents an image.
+            Example: [{"data": "base64_string", "active": True}]
+        """
+        if not self.images:
+            return []
+        
+        # Retrocompatibility: if active_images is not set or mismatched, assume all are active.
+        if self.active_images is None or not isinstance(self.active_images, list) or len(self.active_images) != len(self.images):
+            active_flags = [True] * len(self.images)
+        else:
+            active_flags = self.active_images
+
+        return [
+            {"data": img_data, "active": active_flags[i]}
+            for i, img_data in enumerate(self.images)
+        ]
+
+    def get_active_images(self) -> List[str]:
+        """
+        Returns a list of base64 strings for only the active images.
+        This is backwards-compatible with messages created before this feature.
+        """
+        if not self.images:
+            return []
+        
+        # Retrocompatibility: if active_images is not set, all images are active.
+        if self.active_images is None or not isinstance(self.active_images, list):
+            return self.images
+
+        # Filter images based on the active_images flag list
+        return [
+            img for i, img in enumerate(self.images) 
+            if i < len(self.active_images) and self.active_images[i]
+        ]
+
+    def toggle_image_activation(self, index: int, active: Optional[bool] = None):
+        """
+        Toggles or sets the activation status of an image at a given index.
+        This change is committed to the database if the discussion is DB-backed.
+        
+        Args:
+            index: The index of the image in the 'images' list.
+            active: If provided, sets the status to this boolean. If None, toggles the current status.
+        """
+        if not self.images or index >= len(self.images):
+            raise IndexError("Image index out of range.")
+
+        # Initialize active_images if it's missing or mismatched
+        if self.active_images is None or not isinstance(self.active_images, list) or len(self.active_images) != len(self.images):
+            new_active_images = [True] * len(self.images)
+        else:
+            new_active_images = self.active_images.copy()
+
+        if active is None:
+            new_active_images[index] = not new_active_images[index]
+        else:
+            new_active_images[index] = bool(active) # Ensure it's a boolean
+        
+        self.active_images = new_active_images
+        if self._discussion._is_db_backed:
+            self._discussion.commit()
     
     def set_metadata_item(self, itemname:str, item_value, discussion):
         new_metadata = (self.metadata or {}).copy()
@@ -395,7 +472,6 @@ class LollmsDiscussion:
         object.__setattr__(self, 'autosave', autosave)
         object.__setattr__(self, 'max_context_size', max_context_size)
         object.__setattr__(self, 'scratchpad', "")
-        object.__setattr__(self, 'images', [])
         
         # Internal state
         object.__setattr__(self, '_session', None)
@@ -404,7 +480,6 @@ class LollmsDiscussion:
         object.__setattr__(self, '_messages_to_delete_from_db', set())
         object.__setattr__(self, '_is_db_backed', db_manager is not None)
         
-        object.__setattr__(self, '_system_prompt', None)
         if self._is_db_backed:
             if not db_discussion_obj and not discussion_id:
                 raise ValueError("Either discussion_id or db_discussion_obj must be provided for DB-backed discussions.")
@@ -421,6 +496,30 @@ class LollmsDiscussion:
         else:
             self._create_in_memory_proxy(id=discussion_id)
         
+        object.__setattr__(self, '_system_prompt', getattr(self._db_discussion, 'system_prompt', None))
+
+        # --- THE FIX IS HERE: Load discussion-level images from metadata ---
+        metadata = getattr(self._db_discussion, 'discussion_metadata', {}) or {}
+        image_data = metadata.get("discussion_images", {})
+        
+        images_list = []
+        active_list = []
+        
+        if isinstance(image_data, dict) and 'data' in image_data:
+            # New format: {'data': [...], 'active': [...]}
+            images_list = image_data.get('data', [])
+            active_list = image_data.get('active', [])
+        else:
+            # Covers case where image_data is None, empty dict, or a legacy list
+            # We will rely on migration logic in `get_user_discussion` to fix the DB format
+            # For now, we just don't crash and load empty lists.
+            images_list = []
+            active_list = []
+            
+        object.__setattr__(self, 'images', images_list)
+        object.__setattr__(self, 'active_images', active_list)
+        # --- END FIX ---
+
         self._rebuild_message_index()
 
     @classmethod
@@ -487,18 +586,36 @@ class LollmsDiscussion:
         return getattr(self._db_discussion, name)
 
     def __setattr__(self, name: str, value: Any):
-        """Proxies attribute setting to the underlying discussion object."""
+        """
+        Proxies attribute setting to the underlying discussion object, while
+        keeping internal state like _system_prompt in sync.
+        """
+        # A list of attributes that are internal to the LollmsDiscussion wrapper
+        # and should not be proxied to the underlying data object.
         internal_attrs = [
             'lollmsClient', 'db_manager', 'autosave', 'max_context_size', 'scratchpad',
-            '_session', '_db_discussion', '_message_index', '_messages_to_delete_from_db', '_is_db_backed'
+            'images', 'active_images',
+            '_session', '_db_discussion', '_message_index', '_messages_to_delete_from_db',
+            '_is_db_backed', '_system_prompt'
         ]
+        
         if name in internal_attrs:
+            # If it's an internal attribute, set it directly on the wrapper object.
             object.__setattr__(self, name, value)
         else:
+            # If we are setting 'system_prompt', we must update BOTH the internal
+            # _system_prompt variable AND the underlying data object.
+            if name == 'system_prompt':
+                object.__setattr__(self, '_system_prompt', value)
+            
+            # If the attribute is 'metadata', proxy it to the correct column name.
             if name == 'metadata':
                 setattr(self._db_discussion, 'discussion_metadata', value)
             else:
+                # For all other attributes, proxy them directly to the underlying object.
                 setattr(self._db_discussion, name, value)
+            
+            # Mark the discussion as dirty to trigger a save.
             self.touch()
     
     def _create_in_memory_proxy(self, id: Optional[str] = None):
@@ -527,7 +644,17 @@ class LollmsDiscussion:
         self._message_index = {msg.id: msg for msg in self._db_discussion.messages}
 
     def touch(self):
-        """Marks the discussion as updated and saves it if autosave is enabled."""
+        """Marks the discussion as updated, persists images, and saves if autosave is on."""
+        # Persist in-memory discussion images to the metadata field before saving.
+        # This works for both DB-backed and in-memory discussions.
+        metadata = (getattr(self._db_discussion, 'discussion_metadata', {}) or {}).copy()
+        if self.images or "discussion_images" in metadata:  # Only update if needed
+            metadata["discussion_images"] = {
+                "data": self.images,
+                "active": self.active_images
+            }
+            setattr(self._db_discussion, 'discussion_metadata', metadata)
+
         setattr(self._db_discussion, 'updated_at', datetime.utcnow())
         if self._is_db_backed and self.autosave:
             self.commit()
@@ -572,6 +699,20 @@ class LollmsDiscussion:
         msg_id = kwargs.get('id', str(uuid.uuid4()))
         parent_id = kwargs.get('parent_id', self.active_branch_id)
         
+        # New: Automatically initialize active_images if images are provided
+        if 'images' in kwargs and kwargs['images'] and 'active_images' not in kwargs:
+            kwargs['active_images'] = [True] * len(kwargs['images'])
+        
+        kwargs.setdefault('images', [])
+        kwargs.setdefault('active_images', [])
+
+        if 'sender_type' not in kwargs:
+            if kwargs.get('sender') == 'user':
+                kwargs['sender_type'] = 'user'
+            else:
+                kwargs['sender_type'] = 'assistant'
+
+
         message_data = {
             'id': msg_id,
             'parent_id': parent_id,
@@ -598,7 +739,7 @@ class LollmsDiscussion:
         self.active_branch_id = msg_id
         self.touch()
         return LollmsMessage(self, new_msg_orm)
-        
+    
     def get_branch(self, leaf_id: Optional[str]) -> List[LollmsMessage]:
         """Traces a branch of the conversation from a leaf message back to the root.
 
@@ -620,6 +761,20 @@ class LollmsDiscussion:
             
         return [LollmsMessage(self, orm) for orm in reversed(branch_orms)]
 
+    def get_message(self, message_id: str) -> Optional['LollmsMessage']:
+        """Retrieves a single message by its ID.
+
+        Args:
+            message_id: The unique ID of the message to retrieve.
+
+        Returns:
+            An LollmsMessage instance if found, otherwise None.
+        """
+        db_message = self._message_index.get(message_id)
+        if db_message:
+            return LollmsMessage(self, db_message)
+        return None
+
     def get_full_data_zone(self):
         """Assembles all data zones into a single, formatted string for the prompt."""
         parts = []
@@ -869,8 +1024,9 @@ class LollmsDiscussion:
             if len(self._message_index)>0:
                 ASCIIColors.warning("No active message to regenerate from.\n")
                 ASCIIColors.warning(f"Using last available message:{list(self._message_index.keys())[-1]}\n")
-            else:
+                # Fix for when branch_tip_id is not provided
                 branch_tip_id = list(self._message_index.keys())[-1]
+            else:
                 raise ValueError("No active message to regenerate from.")
         
         last_message_orm = self._message_index[self.active_branch_id]
@@ -886,6 +1042,11 @@ class LollmsDiscussion:
             if self._is_db_backed:
                 self._messages_to_delete_from_db.add(last_message_id)
             
+            self.active_branch_id = parent_id
+            # We now pass the parent ID as the tip, because that's what we want to generate from
+            return self.chat(user_message="", add_user_message=False, branch_tip_id=parent_id, **kwargs)
+
+        # If the last message is a user message, we can just call chat on it
         return self.chat(user_message="", add_user_message=False, branch_tip_id=branch_tip_id, **kwargs)
     
     def delete_branch(self, message_id: str):
@@ -958,7 +1119,7 @@ class LollmsDiscussion:
 
         This method can format the conversation for different backends like OpenAI,
         Ollama, or the native `lollms_text` format. It intelligently handles
-        context limits and non-destructive pruning summaries.
+        context limits, non-destructive pruning summaries, and discussion-level images.
 
         Args:
             format_type: The target format. Can be "lollms_text", "openai_chat",
@@ -977,7 +1138,11 @@ class LollmsDiscussion:
         """
         branch_tip_id = branch_tip_id or self.active_branch_id
         if not branch_tip_id and format_type in ["lollms_text", "openai_chat", "ollama_chat", "markdown"]:
-            return "" if format_type == "lollms_text" else []
+            if format_type in ["lollms_text", "markdown"]:
+                return ""
+            else:
+                return []
+
 
         branch = self.get_branch(branch_tip_id)
         
@@ -1038,8 +1203,11 @@ class LollmsDiscussion:
             for msg in reversed(messages_to_render):
                 sender_str = msg.sender.replace(':', '').replace('!@>', '')
                 content = get_full_content(msg)
-                if msg.images:
-                    content += f"\n({len(msg.images)} image(s) attached)"
+                
+                active_images = msg.get_active_images()
+                if active_images:
+                    content += f"\n({len(active_images)} image(s) attached)"
+
                 msg_text = f"!@>{sender_str}:\n{content}\n"
                 msg_tokens = self.lollmsClient.count_tokens(msg_text)
 
@@ -1054,11 +1222,49 @@ class LollmsDiscussion:
 
         # --- OPENAI & OLLAMA CHAT FORMATS ---
         messages = []
-        if full_system_prompt:
-            if format_type == "markdown":
-                messages.append(f"system: {full_system_prompt}")
-            else:
-                messages.append({"role": "system", "content": full_system_prompt})
+        # Get active discussion-level images
+        active_discussion_b64 = [
+            img for i, img in enumerate(self.images or [])
+            if i < len(self.active_images or []) and self.active_images[i]
+        ]
+        # Handle system message, which can now contain text and/or discussion-level images
+        if full_system_prompt or (active_discussion_b64 and format_type in ["openai_chat", "ollama_chat", "markdown"]):
+            discussion_level_images = build_image_dicts(active_discussion_b64)
+
+            if format_type == "openai_chat":
+                content_parts = []
+                if full_system_prompt:
+                    content_parts.append({"type": "text", "text": full_system_prompt})
+                for img in discussion_level_images:
+                    img_data = img['data']
+                    url = f"data:image/jpeg;base64,{img_data}" if img['type'] == 'base64' else img_data
+                    content_parts.append({"type": "image_url", "image_url": {"url": url, "detail": "auto"}})
+                if content_parts:
+                    messages.append({"role": "system", "content": content_parts})
+
+            elif format_type == "ollama_chat":
+                system_message_dict = {"role": "system", "content": full_system_prompt or ""}
+                base64_images = [img['data'] for img in discussion_level_images if img['type'] == 'base64']
+                if base64_images:
+                    system_message_dict["images"] = base64_images
+                messages.append(system_message_dict)
+
+            elif format_type == "markdown":
+                system_md_parts = []
+                if full_system_prompt:
+                    system_md_parts.append(f"system: {full_system_prompt}")
+
+                for img in discussion_level_images:
+                    img_data = img['data']
+                    url = f"![Image](data:image/jpeg;base64,{img_data})" if img['type'] == 'base64' else f"![Image]({img_data})"
+                    system_md_parts.append(f"\n{url}\n")
+                
+                if system_md_parts:
+                    messages.append("".join(system_md_parts))
+
+            else: # Fallback for any other potential format
+                if full_system_prompt:
+                    messages.append({"role": "system", "content": full_system_prompt})
 
         for msg in branch:
             if msg.sender_type == 'user':
@@ -1066,9 +1272,9 @@ class LollmsDiscussion:
             else:
                 role = participants.get(msg.sender, "assistant")
 
-            content, images = get_full_content(msg), msg.images or []
-            images = build_image_dicts(images)        
-
+            content = get_full_content(msg)
+            active_images_b64 = msg.get_active_images()
+            images = build_image_dicts(active_images_b64)
 
             if format_type == "openai_chat":
                 if images:
@@ -1180,6 +1386,7 @@ class LollmsDiscussion:
                 "- Key decisions or conclusions reached.\n"
                 "- Important entities, projects, or topics mentioned that are likely to recur.\n"
                 "Format the output as a concise list of bullet points. Be brief and factual. "
+                "Do not repeat information that is already in the User Data Zone or the Memory"
                 "If no new, significant long-term information is present, output the single word: 'NOTHING'."
             )
             
@@ -1258,6 +1465,7 @@ class LollmsDiscussion:
         This provides a comprehensive snapshot of the context usage. It accurately calculates
         the token count of the combined system context (prompt, all data zones, summary)
         and the message history, reflecting how the `lollms_text` export format works.
+        It also includes the token count for any active images in the message history.
 
         Args:
             branch_tip_id: The ID of the message branch to measure. Defaults to the active branch.
@@ -1272,15 +1480,20 @@ class LollmsDiscussion:
                         "content": str,
                         "tokens": int,
                         "breakdown": {
-                            "system_prompt": str,
-                            "memory": str,
+                            "system_prompt": {"content": str, "tokens": int},
+                            "memory": {"content": str, "tokens": int},
                             ...
                         }
                     },
                     "message_history": {
                         "content": str,
                         "tokens": int,
-                        "message_count": int
+                        "message_count": int,
+                        "breakdown": {
+                            "text_tokens": int,
+                            "image_tokens": int,
+                            "image_details": [{"message_id": str, "index": int, "tokens": int}]
+                        }
                     }
                 }
             }
@@ -1291,56 +1504,87 @@ class LollmsDiscussion:
             "current_tokens": 0,
             "zones": {}
         }
+        tokenizer = self.lollmsClient.count_tokens
+        tokenizer_images = self.lollmsClient.count_image_tokens
 
         # --- 1. Assemble and Tokenize the Entire System Context Block ---
+        system_context_tokens = 0
         system_prompt_text = (self._system_prompt or "").strip()
-        data_zone_text = self.get_full_data_zone() # This already formats all zones correctly
-        
-        pruning_summary_text = ""
-        if self.pruning_summary and self.pruning_point_id:
-            pruning_summary_text = f"--- Conversation Summary ---\n{self.pruning_summary.strip()}"
+        data_zone_text = self.get_full_data_zone()
+        pruning_summary_content = (self.pruning_summary or "").strip()
+
+        pruning_summary_block = ""
+        if pruning_summary_content and self.pruning_point_id:
+            pruning_summary_block = f"--- Conversation Summary ---\n{pruning_summary_content}"
 
-        # Combine all parts that go into the system block, separated by newlines
         full_system_content_parts = [
-            part for part in [system_prompt_text, data_zone_text, pruning_summary_text] if part
+            part for part in [system_prompt_text, data_zone_text, pruning_summary_block] if part
         ]
         full_system_content = "\n\n".join(full_system_content_parts).strip()
 
         if full_system_content:
-            # Create the final system block as it would be exported
             system_block = f"!@>system:\n{full_system_content}\n"
-            system_tokens = self.lollmsClient.count_tokens(system_block)
+            system_context_tokens = tokenizer(system_block)
             
-            # Create the breakdown for user visibility
             breakdown = {}
             if system_prompt_text:
-                breakdown["system_prompt"] = system_prompt_text
-            if self.memory and self.memory.strip():
-                breakdown["memory"] = self.memory.strip()
-            if self.user_data_zone and self.user_data_zone.strip():
-                breakdown["user_data_zone"] = self.user_data_zone.strip()
-            if self.discussion_data_zone and self.discussion_data_zone.strip():
-                breakdown["discussion_data_zone"] = self.discussion_data_zone.strip()
-            if self.personality_data_zone and self.personality_data_zone.strip():
-                breakdown["personality_data_zone"] = self.personality_data_zone.strip()
-            if self.pruning_summary and self.pruning_summary.strip():
-                breakdown["pruning_summary"] = self.pruning_summary.strip()
+                breakdown["system_prompt"] = {
+                    "content": system_prompt_text,
+                    "tokens": tokenizer(system_prompt_text)
+                }
+
+            memory_text = (self.memory or "").strip()
+            if memory_text:
+                breakdown["memory"] = {
+                    "content": memory_text,
+                    "tokens": tokenizer(memory_text)
+                }
+
+            user_data_text = (self.user_data_zone or "").strip()
+            if user_data_text:
+                breakdown["user_data_zone"] = {
+                    "content": user_data_text,
+                    "tokens": tokenizer(user_data_text)
+                }
+            
+            discussion_data_text = (self.discussion_data_zone or "").strip()
+            if discussion_data_text:
+                breakdown["discussion_data_zone"] = {
+                    "content": discussion_data_text,
+                    "tokens": tokenizer(discussion_data_text)
+                }
+
+            personality_data_text = (self.personality_data_zone or "").strip()
+            if personality_data_text:
+                breakdown["personality_data_zone"] = {
+                    "content": personality_data_text,
+                    "tokens": tokenizer(personality_data_text)
+                }
+
+            if pruning_summary_content:
+                breakdown["pruning_summary"] = {
+                    "content": pruning_summary_content,
+                    "tokens": tokenizer(pruning_summary_content)
+                }
 
             result["zones"]["system_context"] = {
                 "content": full_system_content,
-                "tokens": system_tokens,
+                "tokens": system_context_tokens,
                 "breakdown": breakdown
             }
 
-        # --- 2. Assemble and Tokenize the Message History Block ---
+        # --- 2. Assemble and Tokenize the Message History Block (with images) ---
         branch_tip_id = branch_tip_id or self.active_branch_id
         messages_text = ""
         message_count = 0
+        history_text_tokens = 0
+        total_image_tokens = 0
+        image_details_list = []
+        
         if branch_tip_id:
             branch = self.get_branch(branch_tip_id)
             messages_to_render = branch
             
-            # Adjust for pruning to get the active set of messages
             if self.pruning_summary and self.pruning_point_id:
                 pruning_index = -1
                 for i, msg in enumerate(branch):
@@ -1354,28 +1598,124 @@ class LollmsDiscussion:
             for msg in messages_to_render:
                 sender_str = msg.sender.replace(':', '').replace('!@>', '')
                 content = msg.content.strip()
-                if msg.images:
-                    content += f"\n({len(msg.images)} image(s) attached)"
+
+                active_images = msg.get_active_images()
+                if active_images:
+                    content += f"\n({len(active_images)} image(s) attached)"
+                    # Count image tokens
+                    for i, image_b64 in enumerate(active_images):
+                        tokens = tokenizer_images(image_b64)
+                        if tokens > 0:
+                            total_image_tokens += tokens
+                            image_details_list.append({"message_id": msg.id, "index": i, "tokens": tokens})
+
                 msg_text = f"!@>{sender_str}:\n{content}\n"
                 message_parts.append(msg_text)
             
             messages_text = "".join(message_parts)
             message_count = len(messages_to_render)
 
-        if messages_text:
-            tokens = self.lollmsClient.count_tokens(messages_text)
+        if messages_text or total_image_tokens > 0:
+            history_text_tokens = tokenizer(messages_text)
             result["zones"]["message_history"] = {
                 "content": messages_text,
-                "tokens": tokens,
-                "message_count": message_count
+                "tokens": history_text_tokens + total_image_tokens,
+                "message_count": message_count,
+                "breakdown": {
+                    "text_tokens": history_text_tokens,
+                    "image_tokens": total_image_tokens,
+                    "image_details": image_details_list
+                }
+            }
+        
+        # Calculate discussion-level image tokens separately and add them to the total.
+        active_discussion_b64 = [
+            img for i, img in enumerate(self.images or [])
+            if i < len(self.active_images or []) and self.active_images[i]
+        ]
+        discussion_image_tokens = sum(tokenizer_images(img) for img in active_discussion_b64)
+        
+        # Add a new zone for discussion images for clarity
+        if discussion_image_tokens > 0:
+            result["zones"]["discussion_images"] = {
+                "tokens": discussion_image_tokens,
+                "image_count": len(active_discussion_b64)
             }
 
+
         # --- 3. Finalize the Total Count ---
-        # This remains the most accurate way to get the final count, as it uses the
-        # exact same export logic as the chat method.
-        result["current_tokens"] = self.count_discussion_tokens("lollms_text", branch_tip_id)
+        # Sum up the tokens from all calculated zones
+        total_tokens = 0
+        for zone in result["zones"].values():
+            total_tokens += zone.get("tokens", 0)
+        
+        result["current_tokens"] = total_tokens
 
         return result
+    
+    def get_all_images(self, branch_tip_id: Optional[str] = None) -> List[Dict[str, Any]]:
+        """
+        Retrieves all images from all messages in the specified or active branch.
+
+        Each image is returned as a dictionary containing its data, its active status,
+        the ID of the message it belongs to, and its index within that message.
+
+        Args:
+            branch_tip_id: The ID of the leaf message of the desired branch.
+                           If None, the active branch's leaf is used.
+
+        Returns:
+            A list of dictionaries, each representing an image in the discussion branch.
+            Example: [{"message_id": "...", "index": 0, "data": "...", "active": True}]
+        """
+        all_discussion_images = []
+        branch = self.get_branch(branch_tip_id or self.active_branch_id)
+
+        if not branch:
+            return []
+
+        for message in branch:
+            message_images = message.get_all_images() # This returns [{"data":..., "active":...}]
+            for i, img_info in enumerate(message_images):
+                all_discussion_images.append({
+                    "message_id": message.id,
+                    "index": i,
+                    "data": img_info["data"],
+                    "active": img_info["active"]
+                })
+        
+        return all_discussion_images
+
+    def get_active_images(self, branch_tip_id: Optional[str] = None) -> List[str]:
+        """
+        Retrieves all *active* images from the discussion and from all messages
+        in the specified or active branch.
+
+        This method aggregates the active images from the discussion level and
+        from each message's `get_active_images` call into a single flat list.
+
+        Args:
+            branch_tip_id: The ID of the leaf message of the desired branch.
+                           If None, the active branch's leaf is used.
+
+        Returns:
+            A flat list of base64-encoded strings for all active images.
+        """
+        # Start with active discussion-level images
+        active_discussion_images = [
+            img for i, img in enumerate(self.images or [])
+            if i < len(self.active_images or []) and self.active_images[i]
+        ]
+        
+        branch = self.get_branch(branch_tip_id or self.active_branch_id)
+        if not branch:
+            return active_discussion_images
+
+        for message in branch:
+            active_discussion_images.extend(message.get_active_images())
+
+        return active_discussion_images
+
     def switch_to_branch(self, branch_id):
         self.active_branch_id = branch_id
 
@@ -1412,4 +1752,85 @@ class LollmsDiscussion:
         new_metadata = (self.metadata or {}).copy()
         new_metadata[itemname] = item_value
         self.metadata = new_metadata
-        self.commit()
+        self.commit()
+
+    def add_discussion_image(self, image_b64: str):
+        """
+        Adds an image at the discussion level and marks it as active.
+
+        This image is not tied to a specific message but is considered part of the
+        overall discussion context. It will be included with the system prompt
+        when exporting the conversation for multi-modal models. The change is
+        persisted to the database on the next commit.
+
+        Args:
+            image_b64: A base64-encoded string of the image to add.
+        """
+        self.images.append(image_b64)
+        self.active_images.append(True)
+        self.touch()
+
+    def get_discussion_images(self) -> List[Dict[str, Union[str, bool]]]:
+        """
+        Returns a list of all images attached to the discussion, including their
+        activation status.
+
+        Returns:
+            A list of dictionaries, where each dictionary represents an image.
+            Example: [{"data": "base64_string", "active": True}]
+        """
+        if not self.images:
+            return []
+
+        # Ensure active_images list is in sync, default to True if not
+        if len(self.active_images) != len(self.images):
+            active_flags = [True] * len(self.images)
+        else:
+            active_flags = self.active_images
+
+        return [
+            {"data": img_data, "active": active_flags[i]}
+            for i, img_data in enumerate(self.images)
+        ]
+
+    def toggle_discussion_image_activation(self, index: int, active: Optional[bool] = None):
+        """
+        Toggles or sets the activation status of a discussion-level image at a given index.
+        The change is persisted to the database on the next commit.
+
+        Args:
+            index: The index of the image in the discussion's 'images' list.
+            active: If provided, sets the status to this boolean. If None, toggles the current status.
+        """
+        if not self.images or index >= len(self.images):
+            raise IndexError("Discussion image index out of range.")
+
+        # Ensure active_images list is in sync before modification
+        if len(self.active_images) != len(self.images):
+            self.active_images = [True] * len(self.images)
+
+        if active is None:
+            self.active_images[index] = not self.active_images[index]
+        else:
+            self.active_images[index] = bool(active)
+        
+        self.touch()
+
+    def remove_discussion_image(self, index: int):
+        """
+        Removes a discussion-level image at a given index.
+        The change is persisted to the database on the next commit.
+
+        Args:
+            index: The index of the image in the discussion's 'images' list.
+        """
+        if not self.images or index >= len(self.images):
+            raise IndexError("Discussion image index out of range.")
+
+        # Ensure active_images list is in sync before modification
+        if len(self.active_images) != len(self.images):
+            self.active_images = [True] * len(self.images)
+
+        del self.images[index]
+        del self.active_images[index]
+        self.touch()