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

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of lollms-client might be problematic. Click here for more details.

@@ -168,6 +168,7 @@ def create_dynamic_models(
168
168
 
169
169
  message_metadata = Column(JSON, nullable=True, default=dict)
170
170
  images = Column(JSON, nullable=True, default=list)
171
+ active_images = Column(JSON, nullable=True, default=list) # New: List of booleans for image activation state
171
172
  created_at = Column(DateTime, default=datetime.utcnow)
172
173
 
173
174
  @declared_attr
@@ -217,8 +218,9 @@ class LollmsDataManager:
217
218
  with self.engine.connect() as connection:
218
219
  print("Checking for database schema upgrades...")
219
220
 
221
+ # Discussions table migration
220
222
  cursor = connection.execute(text("PRAGMA table_info(discussions)"))
221
- columns = [row[1] for row in cursor.fetchall()]
223
+ columns = {row[1] for row in cursor.fetchall()}
222
224
 
223
225
  if 'pruning_summary' not in columns:
224
226
  print(" -> Upgrading 'discussions' table: Adding 'pruning_summary' column.")
@@ -247,7 +249,15 @@ class LollmsDataManager:
247
249
  if 'memory' not in columns:
248
250
  print(" -> Upgrading 'discussions' table: Adding 'memory' column.")
249
251
  connection.execute(text("ALTER TABLE discussions ADD COLUMN memory TEXT"))
250
-
252
+
253
+ # Messages table migration
254
+ cursor = connection.execute(text("PRAGMA table_info(messages)"))
255
+ columns = {row[1] for row in cursor.fetchall()}
256
+
257
+ if 'active_images' not in columns:
258
+ print(" -> Upgrading 'messages' table: Adding 'active_images' column.")
259
+ connection.execute(text("ALTER TABLE messages ADD COLUMN active_images TEXT"))
260
+
251
261
  print("Database schema is up to date.")
252
262
  connection.commit()
253
263
 
@@ -356,6 +366,73 @@ class LollmsMessage:
356
366
  def __repr__(self) -> str:
357
367
  """Provides a developer-friendly representation of the message."""
358
368
  return f"<LollmsMessage id={self.id} sender='{self.sender}'>"
369
+
370
+ def get_all_images(self) -> List[Dict[str, Union[str, bool]]]:
371
+ """
372
+ Returns a list of all images associated with this message, including their activation status.
373
+
374
+ Returns:
375
+ A list of dictionaries, where each dictionary represents an image.
376
+ Example: [{"data": "base64_string", "active": True}]
377
+ """
378
+ if not self.images:
379
+ return []
380
+
381
+ # Retrocompatibility: if active_images is not set or mismatched, assume all are active.
382
+ if self.active_images is None or not isinstance(self.active_images, list) or len(self.active_images) != len(self.images):
383
+ active_flags = [True] * len(self.images)
384
+ else:
385
+ active_flags = self.active_images
386
+
387
+ return [
388
+ {"data": img_data, "active": active_flags[i]}
389
+ for i, img_data in enumerate(self.images)
390
+ ]
391
+
392
+ def get_active_images(self) -> List[str]:
393
+ """
394
+ Returns a list of base64 strings for only the active images.
395
+ This is backwards-compatible with messages created before this feature.
396
+ """
397
+ if not self.images:
398
+ return []
399
+
400
+ # Retrocompatibility: if active_images is not set, all images are active.
401
+ if self.active_images is None or not isinstance(self.active_images, list):
402
+ return self.images
403
+
404
+ # Filter images based on the active_images flag list
405
+ return [
406
+ img for i, img in enumerate(self.images)
407
+ if i < len(self.active_images) and self.active_images[i]
408
+ ]
409
+
410
+ def toggle_image_activation(self, index: int, active: Optional[bool] = None):
411
+ """
412
+ Toggles or sets the activation status of an image at a given index.
413
+ This change is committed to the database if the discussion is DB-backed.
414
+
415
+ Args:
416
+ index: The index of the image in the 'images' list.
417
+ active: If provided, sets the status to this boolean. If None, toggles the current status.
418
+ """
419
+ if not self.images or index >= len(self.images):
420
+ raise IndexError("Image index out of range.")
421
+
422
+ # Initialize active_images if it's missing or mismatched
423
+ if self.active_images is None or not isinstance(self.active_images, list) or len(self.active_images) != len(self.images):
424
+ new_active_images = [True] * len(self.images)
425
+ else:
426
+ new_active_images = self.active_images.copy()
427
+
428
+ if active is None:
429
+ new_active_images[index] = not new_active_images[index]
430
+ else:
431
+ new_active_images[index] = bool(active) # Ensure it's a boolean
432
+
433
+ self.active_images = new_active_images
434
+ if self._discussion._is_db_backed:
435
+ self._discussion.commit()
359
436
 
360
437
  def set_metadata_item(self, itemname:str, item_value, discussion):
361
438
  new_metadata = (self.metadata or {}).copy()
@@ -395,7 +472,6 @@ class LollmsDiscussion:
395
472
  object.__setattr__(self, 'autosave', autosave)
396
473
  object.__setattr__(self, 'max_context_size', max_context_size)
397
474
  object.__setattr__(self, 'scratchpad', "")
398
- object.__setattr__(self, 'images', [])
399
475
 
400
476
  # Internal state
401
477
  object.__setattr__(self, '_session', None)
@@ -404,7 +480,6 @@ class LollmsDiscussion:
404
480
  object.__setattr__(self, '_messages_to_delete_from_db', set())
405
481
  object.__setattr__(self, '_is_db_backed', db_manager is not None)
406
482
 
407
- object.__setattr__(self, '_system_prompt', None)
408
483
  if self._is_db_backed:
409
484
  if not db_discussion_obj and not discussion_id:
410
485
  raise ValueError("Either discussion_id or db_discussion_obj must be provided for DB-backed discussions.")
@@ -421,6 +496,30 @@ class LollmsDiscussion:
421
496
  else:
422
497
  self._create_in_memory_proxy(id=discussion_id)
423
498
 
499
+ object.__setattr__(self, '_system_prompt', getattr(self._db_discussion, 'system_prompt', None))
500
+
501
+ # --- THE FIX IS HERE: Load discussion-level images from metadata ---
502
+ metadata = getattr(self._db_discussion, 'discussion_metadata', {}) or {}
503
+ image_data = metadata.get("discussion_images", {})
504
+
505
+ images_list = []
506
+ active_list = []
507
+
508
+ if isinstance(image_data, dict) and 'data' in image_data:
509
+ # New format: {'data': [...], 'active': [...]}
510
+ images_list = image_data.get('data', [])
511
+ active_list = image_data.get('active', [])
512
+ else:
513
+ # Covers case where image_data is None, empty dict, or a legacy list
514
+ # We will rely on migration logic in `get_user_discussion` to fix the DB format
515
+ # For now, we just don't crash and load empty lists.
516
+ images_list = []
517
+ active_list = []
518
+
519
+ object.__setattr__(self, 'images', images_list)
520
+ object.__setattr__(self, 'active_images', active_list)
521
+ # --- END FIX ---
522
+
424
523
  self._rebuild_message_index()
425
524
 
426
525
  @classmethod
@@ -487,18 +586,36 @@ class LollmsDiscussion:
487
586
  return getattr(self._db_discussion, name)
488
587
 
489
588
  def __setattr__(self, name: str, value: Any):
490
- """Proxies attribute setting to the underlying discussion object."""
589
+ """
590
+ Proxies attribute setting to the underlying discussion object, while
591
+ keeping internal state like _system_prompt in sync.
592
+ """
593
+ # A list of attributes that are internal to the LollmsDiscussion wrapper
594
+ # and should not be proxied to the underlying data object.
491
595
  internal_attrs = [
492
596
  'lollmsClient', 'db_manager', 'autosave', 'max_context_size', 'scratchpad',
493
- '_session', '_db_discussion', '_message_index', '_messages_to_delete_from_db', '_is_db_backed'
597
+ 'images', 'active_images',
598
+ '_session', '_db_discussion', '_message_index', '_messages_to_delete_from_db',
599
+ '_is_db_backed', '_system_prompt'
494
600
  ]
601
+
495
602
  if name in internal_attrs:
603
+ # If it's an internal attribute, set it directly on the wrapper object.
496
604
  object.__setattr__(self, name, value)
497
605
  else:
606
+ # If we are setting 'system_prompt', we must update BOTH the internal
607
+ # _system_prompt variable AND the underlying data object.
608
+ if name == 'system_prompt':
609
+ object.__setattr__(self, '_system_prompt', value)
610
+
611
+ # If the attribute is 'metadata', proxy it to the correct column name.
498
612
  if name == 'metadata':
499
613
  setattr(self._db_discussion, 'discussion_metadata', value)
500
614
  else:
615
+ # For all other attributes, proxy them directly to the underlying object.
501
616
  setattr(self._db_discussion, name, value)
617
+
618
+ # Mark the discussion as dirty to trigger a save.
502
619
  self.touch()
503
620
 
504
621
  def _create_in_memory_proxy(self, id: Optional[str] = None):
@@ -527,7 +644,17 @@ class LollmsDiscussion:
527
644
  self._message_index = {msg.id: msg for msg in self._db_discussion.messages}
528
645
 
529
646
  def touch(self):
530
- """Marks the discussion as updated and saves it if autosave is enabled."""
647
+ """Marks the discussion as updated, persists images, and saves if autosave is on."""
648
+ # Persist in-memory discussion images to the metadata field before saving.
649
+ # This works for both DB-backed and in-memory discussions.
650
+ metadata = (getattr(self._db_discussion, 'discussion_metadata', {}) or {}).copy()
651
+ if self.images or "discussion_images" in metadata: # Only update if needed
652
+ metadata["discussion_images"] = {
653
+ "data": self.images,
654
+ "active": self.active_images
655
+ }
656
+ setattr(self._db_discussion, 'discussion_metadata', metadata)
657
+
531
658
  setattr(self._db_discussion, 'updated_at', datetime.utcnow())
532
659
  if self._is_db_backed and self.autosave:
533
660
  self.commit()
@@ -572,6 +699,20 @@ class LollmsDiscussion:
572
699
  msg_id = kwargs.get('id', str(uuid.uuid4()))
573
700
  parent_id = kwargs.get('parent_id', self.active_branch_id)
574
701
 
702
+ # New: Automatically initialize active_images if images are provided
703
+ if 'images' in kwargs and kwargs['images'] and 'active_images' not in kwargs:
704
+ kwargs['active_images'] = [True] * len(kwargs['images'])
705
+
706
+ kwargs.setdefault('images', [])
707
+ kwargs.setdefault('active_images', [])
708
+
709
+ if 'sender_type' not in kwargs:
710
+ if kwargs.get('sender') == 'user':
711
+ kwargs['sender_type'] = 'user'
712
+ else:
713
+ kwargs['sender_type'] = 'assistant'
714
+
715
+
575
716
  message_data = {
576
717
  'id': msg_id,
577
718
  'parent_id': parent_id,
@@ -598,7 +739,7 @@ class LollmsDiscussion:
598
739
  self.active_branch_id = msg_id
599
740
  self.touch()
600
741
  return LollmsMessage(self, new_msg_orm)
601
-
742
+
602
743
  def get_branch(self, leaf_id: Optional[str]) -> List[LollmsMessage]:
603
744
  """Traces a branch of the conversation from a leaf message back to the root.
604
745
 
@@ -620,6 +761,20 @@ class LollmsDiscussion:
620
761
 
621
762
  return [LollmsMessage(self, orm) for orm in reversed(branch_orms)]
622
763
 
764
+ def get_message(self, message_id: str) -> Optional['LollmsMessage']:
765
+ """Retrieves a single message by its ID.
766
+
767
+ Args:
768
+ message_id: The unique ID of the message to retrieve.
769
+
770
+ Returns:
771
+ An LollmsMessage instance if found, otherwise None.
772
+ """
773
+ db_message = self._message_index.get(message_id)
774
+ if db_message:
775
+ return LollmsMessage(self, db_message)
776
+ return None
777
+
623
778
  def get_full_data_zone(self):
624
779
  """Assembles all data zones into a single, formatted string for the prompt."""
625
780
  parts = []
@@ -869,8 +1024,9 @@ class LollmsDiscussion:
869
1024
  if len(self._message_index)>0:
870
1025
  ASCIIColors.warning("No active message to regenerate from.\n")
871
1026
  ASCIIColors.warning(f"Using last available message:{list(self._message_index.keys())[-1]}\n")
872
- else:
1027
+ # Fix for when branch_tip_id is not provided
873
1028
  branch_tip_id = list(self._message_index.keys())[-1]
1029
+ else:
874
1030
  raise ValueError("No active message to regenerate from.")
875
1031
 
876
1032
  last_message_orm = self._message_index[self.active_branch_id]
@@ -886,6 +1042,11 @@ class LollmsDiscussion:
886
1042
  if self._is_db_backed:
887
1043
  self._messages_to_delete_from_db.add(last_message_id)
888
1044
 
1045
+ self.active_branch_id = parent_id
1046
+ # We now pass the parent ID as the tip, because that's what we want to generate from
1047
+ return self.chat(user_message="", add_user_message=False, branch_tip_id=parent_id, **kwargs)
1048
+
1049
+ # If the last message is a user message, we can just call chat on it
889
1050
  return self.chat(user_message="", add_user_message=False, branch_tip_id=branch_tip_id, **kwargs)
890
1051
 
891
1052
  def delete_branch(self, message_id: str):
@@ -958,7 +1119,7 @@ class LollmsDiscussion:
958
1119
 
959
1120
  This method can format the conversation for different backends like OpenAI,
960
1121
  Ollama, or the native `lollms_text` format. It intelligently handles
961
- context limits and non-destructive pruning summaries.
1122
+ context limits, non-destructive pruning summaries, and discussion-level images.
962
1123
 
963
1124
  Args:
964
1125
  format_type: The target format. Can be "lollms_text", "openai_chat",
@@ -977,7 +1138,11 @@ class LollmsDiscussion:
977
1138
  """
978
1139
  branch_tip_id = branch_tip_id or self.active_branch_id
979
1140
  if not branch_tip_id and format_type in ["lollms_text", "openai_chat", "ollama_chat", "markdown"]:
980
- return "" if format_type == "lollms_text" else []
1141
+ if format_type in ["lollms_text", "markdown"]:
1142
+ return ""
1143
+ else:
1144
+ return []
1145
+
981
1146
 
982
1147
  branch = self.get_branch(branch_tip_id)
983
1148
 
@@ -1038,8 +1203,11 @@ class LollmsDiscussion:
1038
1203
  for msg in reversed(messages_to_render):
1039
1204
  sender_str = msg.sender.replace(':', '').replace('!@>', '')
1040
1205
  content = get_full_content(msg)
1041
- if msg.images:
1042
- content += f"\n({len(msg.images)} image(s) attached)"
1206
+
1207
+ active_images = msg.get_active_images()
1208
+ if active_images:
1209
+ content += f"\n({len(active_images)} image(s) attached)"
1210
+
1043
1211
  msg_text = f"!@>{sender_str}:\n{content}\n"
1044
1212
  msg_tokens = self.lollmsClient.count_tokens(msg_text)
1045
1213
 
@@ -1054,11 +1222,49 @@ class LollmsDiscussion:
1054
1222
 
1055
1223
  # --- OPENAI & OLLAMA CHAT FORMATS ---
1056
1224
  messages = []
1057
- if full_system_prompt:
1058
- if format_type == "markdown":
1059
- messages.append(f"system: {full_system_prompt}")
1060
- else:
1061
- messages.append({"role": "system", "content": full_system_prompt})
1225
+ # Get active discussion-level images
1226
+ active_discussion_b64 = [
1227
+ img for i, img in enumerate(self.images or [])
1228
+ if i < len(self.active_images or []) and self.active_images[i]
1229
+ ]
1230
+ # Handle system message, which can now contain text and/or discussion-level images
1231
+ if full_system_prompt or (active_discussion_b64 and format_type in ["openai_chat", "ollama_chat", "markdown"]):
1232
+ discussion_level_images = build_image_dicts(active_discussion_b64)
1233
+
1234
+ if format_type == "openai_chat":
1235
+ content_parts = []
1236
+ if full_system_prompt:
1237
+ content_parts.append({"type": "text", "text": full_system_prompt})
1238
+ for img in discussion_level_images:
1239
+ img_data = img['data']
1240
+ url = f"data:image/jpeg;base64,{img_data}" if img['type'] == 'base64' else img_data
1241
+ content_parts.append({"type": "image_url", "image_url": {"url": url, "detail": "auto"}})
1242
+ if content_parts:
1243
+ messages.append({"role": "system", "content": content_parts})
1244
+
1245
+ elif format_type == "ollama_chat":
1246
+ system_message_dict = {"role": "system", "content": full_system_prompt or ""}
1247
+ base64_images = [img['data'] for img in discussion_level_images if img['type'] == 'base64']
1248
+ if base64_images:
1249
+ system_message_dict["images"] = base64_images
1250
+ messages.append(system_message_dict)
1251
+
1252
+ elif format_type == "markdown":
1253
+ system_md_parts = []
1254
+ if full_system_prompt:
1255
+ system_md_parts.append(f"system: {full_system_prompt}")
1256
+
1257
+ for img in discussion_level_images:
1258
+ img_data = img['data']
1259
+ url = f"![Image](data:image/jpeg;base64,{img_data})" if img['type'] == 'base64' else f"![Image]({img_data})"
1260
+ system_md_parts.append(f"\n{url}\n")
1261
+
1262
+ if system_md_parts:
1263
+ messages.append("".join(system_md_parts))
1264
+
1265
+ else: # Fallback for any other potential format
1266
+ if full_system_prompt:
1267
+ messages.append({"role": "system", "content": full_system_prompt})
1062
1268
 
1063
1269
  for msg in branch:
1064
1270
  if msg.sender_type == 'user':
@@ -1066,9 +1272,9 @@ class LollmsDiscussion:
1066
1272
  else:
1067
1273
  role = participants.get(msg.sender, "assistant")
1068
1274
 
1069
- content, images = get_full_content(msg), msg.images or []
1070
- images = build_image_dicts(images)
1071
-
1275
+ content = get_full_content(msg)
1276
+ active_images_b64 = msg.get_active_images()
1277
+ images = build_image_dicts(active_images_b64)
1072
1278
 
1073
1279
  if format_type == "openai_chat":
1074
1280
  if images:
@@ -1259,6 +1465,7 @@ class LollmsDiscussion:
1259
1465
  This provides a comprehensive snapshot of the context usage. It accurately calculates
1260
1466
  the token count of the combined system context (prompt, all data zones, summary)
1261
1467
  and the message history, reflecting how the `lollms_text` export format works.
1468
+ It also includes the token count for any active images in the message history.
1262
1469
 
1263
1470
  Args:
1264
1471
  branch_tip_id: The ID of the message branch to measure. Defaults to the active branch.
@@ -1281,7 +1488,12 @@ class LollmsDiscussion:
1281
1488
  "message_history": {
1282
1489
  "content": str,
1283
1490
  "tokens": int,
1284
- "message_count": int
1491
+ "message_count": int,
1492
+ "breakdown": {
1493
+ "text_tokens": int,
1494
+ "image_tokens": int,
1495
+ "image_details": [{"message_id": str, "index": int, "tokens": int}]
1496
+ }
1285
1497
  }
1286
1498
  }
1287
1499
  }
@@ -1293,8 +1505,10 @@ class LollmsDiscussion:
1293
1505
  "zones": {}
1294
1506
  }
1295
1507
  tokenizer = self.lollmsClient.count_tokens
1508
+ tokenizer_images = self.lollmsClient.count_image_tokens
1296
1509
 
1297
1510
  # --- 1. Assemble and Tokenize the Entire System Context Block ---
1511
+ system_context_tokens = 0
1298
1512
  system_prompt_text = (self._system_prompt or "").strip()
1299
1513
  data_zone_text = self.get_full_data_zone()
1300
1514
  pruning_summary_content = (self.pruning_summary or "").strip()
@@ -1310,7 +1524,7 @@ class LollmsDiscussion:
1310
1524
 
1311
1525
  if full_system_content:
1312
1526
  system_block = f"!@>system:\n{full_system_content}\n"
1313
- system_tokens = tokenizer(system_block)
1527
+ system_context_tokens = tokenizer(system_block)
1314
1528
 
1315
1529
  breakdown = {}
1316
1530
  if system_prompt_text:
@@ -1355,14 +1569,18 @@ class LollmsDiscussion:
1355
1569
 
1356
1570
  result["zones"]["system_context"] = {
1357
1571
  "content": full_system_content,
1358
- "tokens": system_tokens,
1572
+ "tokens": system_context_tokens,
1359
1573
  "breakdown": breakdown
1360
1574
  }
1361
1575
 
1362
- # --- 2. Assemble and Tokenize the Message History Block ---
1576
+ # --- 2. Assemble and Tokenize the Message History Block (with images) ---
1363
1577
  branch_tip_id = branch_tip_id or self.active_branch_id
1364
1578
  messages_text = ""
1365
1579
  message_count = 0
1580
+ history_text_tokens = 0
1581
+ total_image_tokens = 0
1582
+ image_details_list = []
1583
+
1366
1584
  if branch_tip_id:
1367
1585
  branch = self.get_branch(branch_tip_id)
1368
1586
  messages_to_render = branch
@@ -1380,27 +1598,124 @@ class LollmsDiscussion:
1380
1598
  for msg in messages_to_render:
1381
1599
  sender_str = msg.sender.replace(':', '').replace('!@>', '')
1382
1600
  content = msg.content.strip()
1383
- if msg.images:
1384
- content += f"\n({len(msg.images)} image(s) attached)"
1601
+
1602
+ active_images = msg.get_active_images()
1603
+ if active_images:
1604
+ content += f"\n({len(active_images)} image(s) attached)"
1605
+ # Count image tokens
1606
+ for i, image_b64 in enumerate(active_images):
1607
+ tokens = tokenizer_images(image_b64)
1608
+ if tokens > 0:
1609
+ total_image_tokens += tokens
1610
+ image_details_list.append({"message_id": msg.id, "index": i, "tokens": tokens})
1611
+
1385
1612
  msg_text = f"!@>{sender_str}:\n{content}\n"
1386
1613
  message_parts.append(msg_text)
1387
1614
 
1388
1615
  messages_text = "".join(message_parts)
1389
1616
  message_count = len(messages_to_render)
1390
1617
 
1391
- if messages_text:
1392
- tokens = tokenizer(messages_text)
1618
+ if messages_text or total_image_tokens > 0:
1619
+ history_text_tokens = tokenizer(messages_text)
1393
1620
  result["zones"]["message_history"] = {
1394
1621
  "content": messages_text,
1395
- "tokens": tokens,
1396
- "message_count": message_count
1622
+ "tokens": history_text_tokens + total_image_tokens,
1623
+ "message_count": message_count,
1624
+ "breakdown": {
1625
+ "text_tokens": history_text_tokens,
1626
+ "image_tokens": total_image_tokens,
1627
+ "image_details": image_details_list
1628
+ }
1629
+ }
1630
+
1631
+ # Calculate discussion-level image tokens separately and add them to the total.
1632
+ active_discussion_b64 = [
1633
+ img for i, img in enumerate(self.images or [])
1634
+ if i < len(self.active_images or []) and self.active_images[i]
1635
+ ]
1636
+ discussion_image_tokens = sum(tokenizer_images(img) for img in active_discussion_b64)
1637
+
1638
+ # Add a new zone for discussion images for clarity
1639
+ if discussion_image_tokens > 0:
1640
+ result["zones"]["discussion_images"] = {
1641
+ "tokens": discussion_image_tokens,
1642
+ "image_count": len(active_discussion_b64)
1397
1643
  }
1398
1644
 
1645
+
1399
1646
  # --- 3. Finalize the Total Count ---
1400
- result["current_tokens"] = self.count_discussion_tokens("lollms_text", branch_tip_id)
1647
+ # Sum up the tokens from all calculated zones
1648
+ total_tokens = 0
1649
+ for zone in result["zones"].values():
1650
+ total_tokens += zone.get("tokens", 0)
1651
+
1652
+ result["current_tokens"] = total_tokens
1401
1653
 
1402
1654
  return result
1403
1655
 
1656
+ def get_all_images(self, branch_tip_id: Optional[str] = None) -> List[Dict[str, Any]]:
1657
+ """
1658
+ Retrieves all images from all messages in the specified or active branch.
1659
+
1660
+ Each image is returned as a dictionary containing its data, its active status,
1661
+ the ID of the message it belongs to, and its index within that message.
1662
+
1663
+ Args:
1664
+ branch_tip_id: The ID of the leaf message of the desired branch.
1665
+ If None, the active branch's leaf is used.
1666
+
1667
+ Returns:
1668
+ A list of dictionaries, each representing an image in the discussion branch.
1669
+ Example: [{"message_id": "...", "index": 0, "data": "...", "active": True}]
1670
+ """
1671
+ all_discussion_images = []
1672
+ branch = self.get_branch(branch_tip_id or self.active_branch_id)
1673
+
1674
+ if not branch:
1675
+ return []
1676
+
1677
+ for message in branch:
1678
+ message_images = message.get_all_images() # This returns [{"data":..., "active":...}]
1679
+ for i, img_info in enumerate(message_images):
1680
+ all_discussion_images.append({
1681
+ "message_id": message.id,
1682
+ "index": i,
1683
+ "data": img_info["data"],
1684
+ "active": img_info["active"]
1685
+ })
1686
+
1687
+ return all_discussion_images
1688
+
1689
+ def get_active_images(self, branch_tip_id: Optional[str] = None) -> List[str]:
1690
+ """
1691
+ Retrieves all *active* images from the discussion and from all messages
1692
+ in the specified or active branch.
1693
+
1694
+ This method aggregates the active images from the discussion level and
1695
+ from each message's `get_active_images` call into a single flat list.
1696
+
1697
+ Args:
1698
+ branch_tip_id: The ID of the leaf message of the desired branch.
1699
+ If None, the active branch's leaf is used.
1700
+
1701
+ Returns:
1702
+ A flat list of base64-encoded strings for all active images.
1703
+ """
1704
+ # Start with active discussion-level images
1705
+ active_discussion_images = [
1706
+ img for i, img in enumerate(self.images or [])
1707
+ if i < len(self.active_images or []) and self.active_images[i]
1708
+ ]
1709
+
1710
+ branch = self.get_branch(branch_tip_id or self.active_branch_id)
1711
+ if not branch:
1712
+ return active_discussion_images
1713
+
1714
+ for message in branch:
1715
+ active_discussion_images.extend(message.get_active_images())
1716
+
1717
+ return active_discussion_images
1718
+
1404
1719
  def switch_to_branch(self, branch_id):
1405
1720
  self.active_branch_id = branch_id
1406
1721
 
@@ -1438,3 +1753,84 @@ class LollmsDiscussion:
1438
1753
  new_metadata[itemname] = item_value
1439
1754
  self.metadata = new_metadata
1440
1755
  self.commit()
1756
+
1757
+ def add_discussion_image(self, image_b64: str):
1758
+ """
1759
+ Adds an image at the discussion level and marks it as active.
1760
+
1761
+ This image is not tied to a specific message but is considered part of the
1762
+ overall discussion context. It will be included with the system prompt
1763
+ when exporting the conversation for multi-modal models. The change is
1764
+ persisted to the database on the next commit.
1765
+
1766
+ Args:
1767
+ image_b64: A base64-encoded string of the image to add.
1768
+ """
1769
+ self.images.append(image_b64)
1770
+ self.active_images.append(True)
1771
+ self.touch()
1772
+
1773
+ def get_discussion_images(self) -> List[Dict[str, Union[str, bool]]]:
1774
+ """
1775
+ Returns a list of all images attached to the discussion, including their
1776
+ activation status.
1777
+
1778
+ Returns:
1779
+ A list of dictionaries, where each dictionary represents an image.
1780
+ Example: [{"data": "base64_string", "active": True}]
1781
+ """
1782
+ if not self.images:
1783
+ return []
1784
+
1785
+ # Ensure active_images list is in sync, default to True if not
1786
+ if len(self.active_images) != len(self.images):
1787
+ active_flags = [True] * len(self.images)
1788
+ else:
1789
+ active_flags = self.active_images
1790
+
1791
+ return [
1792
+ {"data": img_data, "active": active_flags[i]}
1793
+ for i, img_data in enumerate(self.images)
1794
+ ]
1795
+
1796
+ def toggle_discussion_image_activation(self, index: int, active: Optional[bool] = None):
1797
+ """
1798
+ Toggles or sets the activation status of a discussion-level image at a given index.
1799
+ The change is persisted to the database on the next commit.
1800
+
1801
+ Args:
1802
+ index: The index of the image in the discussion's 'images' list.
1803
+ active: If provided, sets the status to this boolean. If None, toggles the current status.
1804
+ """
1805
+ if not self.images or index >= len(self.images):
1806
+ raise IndexError("Discussion image index out of range.")
1807
+
1808
+ # Ensure active_images list is in sync before modification
1809
+ if len(self.active_images) != len(self.images):
1810
+ self.active_images = [True] * len(self.images)
1811
+
1812
+ if active is None:
1813
+ self.active_images[index] = not self.active_images[index]
1814
+ else:
1815
+ self.active_images[index] = bool(active)
1816
+
1817
+ self.touch()
1818
+
1819
+ def remove_discussion_image(self, index: int):
1820
+ """
1821
+ Removes a discussion-level image at a given index.
1822
+ The change is persisted to the database on the next commit.
1823
+
1824
+ Args:
1825
+ index: The index of the image in the discussion's 'images' list.
1826
+ """
1827
+ if not self.images or index >= len(self.images):
1828
+ raise IndexError("Discussion image index out of range.")
1829
+
1830
+ # Ensure active_images list is in sync before modification
1831
+ if len(self.active_images) != len(self.images):
1832
+ self.active_images = [True] * len(self.images)
1833
+
1834
+ del self.images[index]
1835
+ del self.active_images[index]
1836
+ self.touch()