MemoryOS 0.2.1__py3-none-any.whl → 0.2.2__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 MemoryOS might be problematic. Click here for more details.

Files changed (74) hide show
  1. {memoryos-0.2.1.dist-info → memoryos-0.2.2.dist-info}/METADATA +2 -1
  2. {memoryos-0.2.1.dist-info → memoryos-0.2.2.dist-info}/RECORD +72 -55
  3. memos/__init__.py +1 -1
  4. memos/api/config.py +156 -65
  5. memos/api/context/context.py +147 -0
  6. memos/api/context/dependencies.py +90 -0
  7. memos/api/product_models.py +5 -1
  8. memos/api/routers/product_router.py +54 -26
  9. memos/configs/graph_db.py +49 -1
  10. memos/configs/internet_retriever.py +6 -0
  11. memos/configs/mem_os.py +5 -0
  12. memos/configs/mem_reader.py +9 -0
  13. memos/configs/mem_scheduler.py +18 -4
  14. memos/configs/mem_user.py +58 -0
  15. memos/graph_dbs/base.py +9 -1
  16. memos/graph_dbs/factory.py +2 -0
  17. memos/graph_dbs/nebular.py +1364 -0
  18. memos/graph_dbs/neo4j.py +4 -4
  19. memos/log.py +1 -1
  20. memos/mem_cube/utils.py +13 -6
  21. memos/mem_os/core.py +140 -30
  22. memos/mem_os/main.py +1 -1
  23. memos/mem_os/product.py +266 -152
  24. memos/mem_os/utils/format_utils.py +314 -67
  25. memos/mem_reader/simple_struct.py +13 -5
  26. memos/mem_scheduler/base_scheduler.py +220 -250
  27. memos/mem_scheduler/general_scheduler.py +193 -73
  28. memos/mem_scheduler/modules/base.py +5 -5
  29. memos/mem_scheduler/modules/dispatcher.py +6 -9
  30. memos/mem_scheduler/modules/misc.py +81 -16
  31. memos/mem_scheduler/modules/monitor.py +52 -41
  32. memos/mem_scheduler/modules/rabbitmq_service.py +9 -7
  33. memos/mem_scheduler/modules/retriever.py +108 -191
  34. memos/mem_scheduler/modules/scheduler_logger.py +255 -0
  35. memos/mem_scheduler/mos_for_test_scheduler.py +16 -19
  36. memos/mem_scheduler/schemas/__init__.py +0 -0
  37. memos/mem_scheduler/schemas/general_schemas.py +43 -0
  38. memos/mem_scheduler/schemas/message_schemas.py +148 -0
  39. memos/mem_scheduler/schemas/monitor_schemas.py +329 -0
  40. memos/mem_scheduler/utils/__init__.py +0 -0
  41. memos/mem_scheduler/utils/filter_utils.py +176 -0
  42. memos/mem_scheduler/utils/misc_utils.py +61 -0
  43. memos/mem_user/factory.py +94 -0
  44. memos/mem_user/mysql_persistent_user_manager.py +271 -0
  45. memos/mem_user/mysql_user_manager.py +500 -0
  46. memos/mem_user/persistent_factory.py +96 -0
  47. memos/mem_user/user_manager.py +4 -4
  48. memos/memories/activation/item.py +4 -0
  49. memos/memories/textual/base.py +1 -1
  50. memos/memories/textual/general.py +35 -91
  51. memos/memories/textual/item.py +5 -33
  52. memos/memories/textual/tree.py +13 -7
  53. memos/memories/textual/tree_text_memory/organize/conflict.py +4 -2
  54. memos/memories/textual/tree_text_memory/organize/relation_reason_detector.py +47 -43
  55. memos/memories/textual/tree_text_memory/organize/reorganizer.py +8 -5
  56. memos/memories/textual/tree_text_memory/retrieve/internet_retriever.py +6 -3
  57. memos/memories/textual/tree_text_memory/retrieve/internet_retriever_factory.py +2 -0
  58. memos/memories/textual/tree_text_memory/retrieve/retrieval_mid_structs.py +2 -0
  59. memos/memories/textual/tree_text_memory/retrieve/searcher.py +46 -23
  60. memos/memories/textual/tree_text_memory/retrieve/task_goal_parser.py +42 -15
  61. memos/memories/textual/tree_text_memory/retrieve/utils.py +11 -7
  62. memos/memories/textual/tree_text_memory/retrieve/xinyusearch.py +62 -58
  63. memos/memos_tools/dinding_report_bot.py +422 -0
  64. memos/memos_tools/notification_service.py +44 -0
  65. memos/memos_tools/notification_utils.py +96 -0
  66. memos/settings.py +3 -1
  67. memos/templates/mem_reader_prompts.py +2 -1
  68. memos/templates/mem_scheduler_prompts.py +41 -7
  69. memos/templates/mos_prompts.py +87 -0
  70. memos/mem_scheduler/modules/schemas.py +0 -328
  71. memos/mem_scheduler/utils.py +0 -75
  72. {memoryos-0.2.1.dist-info → memoryos-0.2.2.dist-info}/LICENSE +0 -0
  73. {memoryos-0.2.1.dist-info → memoryos-0.2.2.dist-info}/WHEEL +0 -0
  74. {memoryos-0.2.1.dist-info → memoryos-0.2.2.dist-info}/entry_points.txt +0 -0
memos/mem_os/product.py CHANGED
@@ -16,17 +16,25 @@ from memos.log import get_logger
16
16
  from memos.mem_cube.general import GeneralMemCube
17
17
  from memos.mem_os.core import MOSCore
18
18
  from memos.mem_os.utils.format_utils import (
19
+ clean_json_response,
19
20
  convert_graph_to_tree_forworkmem,
21
+ ensure_unique_tree_ids,
20
22
  filter_nodes_by_tree_ids,
21
23
  remove_embedding_recursive,
22
24
  sort_children_by_memory_type,
25
+ split_continuous_references,
23
26
  )
24
- from memos.mem_scheduler.modules.schemas import ANSWER_LABEL, QUERY_LABEL, ScheduleMessageItem
25
- from memos.mem_user.persistent_user_manager import PersistentUserManager
27
+ from memos.mem_scheduler.schemas.general_schemas import (
28
+ ANSWER_LABEL,
29
+ QUERY_LABEL,
30
+ )
31
+ from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
32
+ from memos.mem_user.persistent_factory import PersistentUserManagerFactory
26
33
  from memos.mem_user.user_manager import UserRole
27
34
  from memos.memories.textual.item import (
28
35
  TextualMemoryItem,
29
36
  )
37
+ from memos.templates.mos_prompts import MEMOS_PRODUCT_BASE_PROMPT, MEMOS_PRODUCT_ENHANCE_PROMPT
30
38
  from memos.types import MessageList
31
39
 
32
40
 
@@ -46,8 +54,10 @@ class MOSProduct(MOSCore):
46
54
  def __init__(
47
55
  self,
48
56
  default_config: MOSConfig | None = None,
49
- max_user_instances: int = 100,
57
+ max_user_instances: int = 1,
50
58
  default_cube_config: GeneralMemCubeConfig | None = None,
59
+ online_bot=None,
60
+ error_bot=None,
51
61
  ):
52
62
  """
53
63
  Initialize MOSProduct with an optional default configuration.
@@ -56,6 +66,8 @@ class MOSProduct(MOSCore):
56
66
  default_config (MOSConfig | None): Default configuration for new users
57
67
  max_user_instances (int): Maximum number of user instances to keep in memory
58
68
  default_cube_config (GeneralMemCubeConfig | None): Default cube configuration for loading cubes
69
+ online_bot: DingDing online_bot function or None if disabled
70
+ error_bot: DingDing error_bot function or None if disabled
59
71
  """
60
72
  # Initialize with a root config for shared resources
61
73
  if default_config is None:
@@ -75,21 +87,28 @@ class MOSProduct(MOSCore):
75
87
  root_config.user_id = "root"
76
88
  root_config.session_id = "root_session"
77
89
 
78
- # Initialize parent MOSCore with root config
79
- super().__init__(root_config)
90
+ # Create persistent user manager BEFORE calling parent constructor
91
+ persistent_user_manager_client = PersistentUserManagerFactory.from_config(
92
+ config_factory=root_config.user_manager
93
+ )
94
+
95
+ # Initialize parent MOSCore with root config and persistent user manager
96
+ super().__init__(root_config, user_manager=persistent_user_manager_client)
80
97
 
81
98
  # Product-specific attributes
82
99
  self.default_config = default_config
83
100
  self.default_cube_config = default_cube_config
84
101
  self.max_user_instances = max_user_instances
102
+ self.online_bot = online_bot
103
+ self.error_bot = error_bot
85
104
 
86
105
  # User-specific data structures
87
106
  self.user_configs: dict[str, MOSConfig] = {}
88
107
  self.user_cube_access: dict[str, set[str]] = {} # user_id -> set of cube_ids
89
108
  self.user_chat_histories: dict[str, dict] = {}
90
109
 
91
- # Use PersistentUserManager for user management
92
- self.global_user_manager = PersistentUserManager(user_id="root")
110
+ # Note: self.user_manager is now the persistent user manager from parent class
111
+ # No need for separate global_user_manager as they are the same instance
93
112
 
94
113
  # Initialize tiktoken for streaming
95
114
  try:
@@ -116,10 +135,10 @@ class MOSProduct(MOSCore):
116
135
  """
117
136
  try:
118
137
  # Get all user configurations from persistent storage
119
- user_configs = self.global_user_manager.list_user_configs()
138
+ user_configs = self.user_manager.list_user_configs()
120
139
 
121
140
  # Get the raw database records for sorting by updated_at
122
- session = self.global_user_manager._get_session()
141
+ session = self.user_manager._get_session()
123
142
  try:
124
143
  from memos.mem_user.persistent_user_manager import UserConfig
125
144
 
@@ -165,7 +184,7 @@ class MOSProduct(MOSCore):
165
184
  """
166
185
  try:
167
186
  # Get user's accessible cubes from persistent storage
168
- accessible_cubes = self.global_user_manager.get_user_cubes(user_id)
187
+ accessible_cubes = self.user_manager.get_user_cubes(user_id)
169
188
 
170
189
  for cube in accessible_cubes:
171
190
  if cube.cube_id not in self.mem_cubes:
@@ -188,11 +207,12 @@ class MOSProduct(MOSCore):
188
207
  )
189
208
  except Exception as e:
190
209
  logger.error(
191
- f"Failed to pre-load cube {cube.cube_id} for user {user_id}: {e}"
210
+ f"Failed to pre-load cube {cube.cube_id} for user {user_id}: {e}",
211
+ exc_info=True,
192
212
  )
193
213
 
194
214
  except Exception as e:
195
- logger.error(f"Error pre-loading cubes for user {user_id}: {e}")
215
+ logger.error(f"Error pre-loading cubes for user {user_id}: {e}", exc_info=True)
196
216
 
197
217
  def _load_user_cubes(
198
218
  self, user_id: str, default_cube_config: GeneralMemCubeConfig | None = None
@@ -204,7 +224,7 @@ class MOSProduct(MOSCore):
204
224
  default_cube_config (GeneralMemCubeConfig | None, optional): Default cube configuration. Defaults to None.
205
225
  """
206
226
  # Get user's accessible cubes from persistent storage
207
- accessible_cubes = self.global_user_manager.get_user_cubes(user_id)
227
+ accessible_cubes = self.user_manager.get_user_cubes(user_id)
208
228
 
209
229
  for cube in accessible_cubes[:1]:
210
230
  if cube.cube_id not in self.mem_cubes:
@@ -238,7 +258,7 @@ class MOSProduct(MOSCore):
238
258
  return
239
259
 
240
260
  # Try to get config from persistent storage first
241
- stored_config = self.global_user_manager.get_user_config(user_id)
261
+ stored_config = self.user_manager.get_user_config(user_id)
242
262
  if stored_config:
243
263
  self.user_configs[user_id] = stored_config
244
264
  self._load_user_cube_access(user_id)
@@ -268,7 +288,7 @@ class MOSProduct(MOSCore):
268
288
  """Load user's cube access permissions."""
269
289
  try:
270
290
  # Get user's accessible cubes from persistent storage
271
- accessible_cubes = self.global_user_manager.get_user_cube_access(user_id)
291
+ accessible_cubes = self.user_manager.get_user_cube_access(user_id)
272
292
  self.user_cube_access[user_id] = set(accessible_cubes)
273
293
  except Exception as e:
274
294
  logger.warning(f"Failed to load cube access for user {user_id}: {e}")
@@ -304,7 +324,7 @@ class MOSProduct(MOSCore):
304
324
  user_config.session_id = f"{user_id}_session"
305
325
 
306
326
  # Save configuration to persistent storage
307
- self.global_user_manager.save_user_config(user_id, user_config)
327
+ self.user_manager.save_user_config(user_id, user_config)
308
328
 
309
329
  return user_config
310
330
 
@@ -316,7 +336,7 @@ class MOSProduct(MOSCore):
316
336
  return self.user_configs[user_id]
317
337
 
318
338
  # Try to get config from persistent storage first
319
- stored_config = self.global_user_manager.get_user_config(user_id)
339
+ stored_config = self.user_manager.get_user_config(user_id)
320
340
  if stored_config:
321
341
  return self._create_user_config(user_id, stored_config)
322
342
 
@@ -327,7 +347,9 @@ class MOSProduct(MOSCore):
327
347
 
328
348
  return self._create_user_config(user_id, user_config)
329
349
 
330
- def _build_system_prompt(self, user_id: str, memories_all: list[TextualMemoryItem]) -> str:
350
+ def _build_system_prompt(
351
+ self, memories_all: list[TextualMemoryItem], base_prompt: str | None = None
352
+ ) -> str:
331
353
  """
332
354
  Build custom system prompt for the user with memory references.
333
355
 
@@ -340,28 +362,49 @@ class MOSProduct(MOSCore):
340
362
  """
341
363
 
342
364
  # Build base prompt
343
- base_prompt = (
344
- "You are a knowledgeable and helpful AI assistant with access to user memories. "
345
- "When responding to user queries, you should reference relevant memories using the provided memory IDs. "
346
- "Use the reference format: [1-n:memoriesID] "
347
- "where refid is a sequential number starting from 1 and increments for each reference in your response, "
348
- "and memoriesID is the specific memory ID provided in the available memories list. "
349
- "For example: [1:abc123], [2:def456], [3:ghi789], [4:jkl101], [5:mno112] "
350
- "Only reference memories that are directly relevant to the user's question. "
351
- "Make your responses natural and conversational while incorporating memory references when appropriate."
352
- )
353
-
354
365
  # Add memory context if available
355
366
  if memories_all:
356
367
  memory_context = "\n\n## Available ID Memories:\n"
357
368
  for i, memory in enumerate(memories_all, 1):
358
369
  # Format: [memory_id]: memory_content
359
370
  memory_id = f"{memory.id.split('-')[0]}" if hasattr(memory, "id") else f"mem_{i}"
360
- memory_content = memory.memory if hasattr(memory, "memory") else str(memory)
371
+ memory_content = memory.memory[:500] if hasattr(memory, "memory") else str(memory)
372
+ memory_content = memory_content.replace("\n", " ")
361
373
  memory_context += f"{memory_id}: {memory_content}\n"
362
- return base_prompt + memory_context
374
+ return MEMOS_PRODUCT_BASE_PROMPT + memory_context
375
+
376
+ return MEMOS_PRODUCT_BASE_PROMPT
363
377
 
364
- return base_prompt
378
+ def _build_enhance_system_prompt(
379
+ self, user_id: str, memories_all: list[TextualMemoryItem]
380
+ ) -> str:
381
+ """
382
+ Build enhance prompt for the user with memory references.
383
+ """
384
+ if memories_all:
385
+ personal_memory_context = "\n\n## Available ID and PersonalMemory Memories:\n"
386
+ outer_memory_context = "\n\n## Available ID and OuterMemory Memories:\n"
387
+ for i, memory in enumerate(memories_all, 1):
388
+ # Format: [memory_id]: memory_content
389
+ if memory.metadata.memory_type != "OuterMemory":
390
+ memory_id = (
391
+ f"{memory.id.split('-')[0]}" if hasattr(memory, "id") else f"mem_{i}"
392
+ )
393
+ memory_content = (
394
+ memory.memory[:500] if hasattr(memory, "memory") else str(memory)
395
+ )
396
+ personal_memory_context += f"{memory_id}: {memory_content}\n"
397
+ else:
398
+ memory_id = (
399
+ f"{memory.id.split('-')[0]}" if hasattr(memory, "id") else f"mem_{i}"
400
+ )
401
+ memory_content = (
402
+ memory.memory[:500] if hasattr(memory, "memory") else str(memory)
403
+ )
404
+ memory_content = memory_content.replace("\n", " ")
405
+ outer_memory_context += f"{memory_id}: {memory_content}\n"
406
+ return MEMOS_PRODUCT_ENHANCE_PROMPT + personal_memory_context + outer_memory_context
407
+ return MEMOS_PRODUCT_ENHANCE_PROMPT
365
408
 
366
409
  def _process_streaming_references_complete(self, text_buffer: str) -> tuple[str, str]:
367
410
  """
@@ -386,9 +429,13 @@ class MOSProduct(MOSCore):
386
429
  last_match = complete_matches[-1]
387
430
  end_pos = last_match.end()
388
431
 
389
- # Return text up to the end of the last complete tag
432
+ # Get text up to the end of the last complete tag
390
433
  processed_text = text_buffer[:end_pos]
391
434
  remaining_buffer = text_buffer[end_pos:]
435
+
436
+ # Apply reference splitting to the processed text
437
+ processed_text = split_continuous_references(processed_text)
438
+
392
439
  return processed_text, remaining_buffer
393
440
 
394
441
  # Check for incomplete reference tags
@@ -407,37 +454,56 @@ class MOSProduct(MOSCore):
407
454
  return "", text_buffer
408
455
  else:
409
456
  # Incomplete opening pattern, return text before it
410
- return text_buffer[:opening_start], text_buffer[opening_start:]
457
+ processed_text = text_buffer[:opening_start]
458
+ # Apply reference splitting to the processed text
459
+ processed_text = split_continuous_references(processed_text)
460
+ return processed_text, text_buffer[opening_start:]
411
461
 
412
462
  # Check for partial opening pattern (starts with [ but not complete)
413
463
  if "[" in text_buffer:
414
464
  ref_start = text_buffer.find("[")
415
- return text_buffer[:ref_start], text_buffer[ref_start:]
465
+ processed_text = text_buffer[:ref_start]
466
+ # Apply reference splitting to the processed text
467
+ processed_text = split_continuous_references(processed_text)
468
+ return processed_text, text_buffer[ref_start:]
416
469
 
417
- # No reference tags found, return all text
418
- return text_buffer, ""
470
+ # No reference tags found, apply reference splitting and return all text
471
+ processed_text = split_continuous_references(text_buffer)
472
+ return processed_text, ""
419
473
 
420
- def _extract_references_from_response(self, response: str) -> list[dict]:
474
+ def _extract_references_from_response(self, response: str) -> tuple[str, list[dict]]:
421
475
  """
422
- Extract reference information from the response.
476
+ Extract reference information from the response and return clean text.
423
477
 
424
478
  Args:
425
479
  response (str): The complete response text.
426
480
 
427
481
  Returns:
428
- list[dict]: List of reference information.
482
+ tuple[str, list[dict]]: A tuple containing:
483
+ - clean_text: Text with reference markers removed
484
+ - references: List of reference information
429
485
  """
430
486
  import re
431
487
 
432
- references = []
433
- # Pattern to match [refid:memoriesID]
434
- pattern = r"\[(\d+):([^\]]+)\]"
488
+ try:
489
+ references = []
490
+ # Pattern to match [refid:memoriesID]
491
+ pattern = r"\[(\d+):([^\]]+)\]"
492
+
493
+ matches = re.findall(pattern, response)
494
+ for ref_number, memory_id in matches:
495
+ references.append({"memory_id": memory_id, "reference_number": int(ref_number)})
496
+
497
+ # Remove all reference markers from the text to get clean text
498
+ clean_text = re.sub(pattern, "", response)
435
499
 
436
- matches = re.findall(pattern, response)
437
- for ref_number, memory_id in matches:
438
- references.append({"memory_id": memory_id, "reference_number": int(ref_number)})
500
+ # Clean up any extra whitespace that might be left after removing markers
501
+ clean_text = re.sub(r"\s+", " ", clean_text).strip()
439
502
 
440
- return references
503
+ return clean_text, references
504
+ except Exception as e:
505
+ logger.error(f"Error extracting references from response: {e}", exc_info=True)
506
+ return response, []
441
507
 
442
508
  def _chunk_response_with_tiktoken(
443
509
  self, response: str, chunk_size: int = 5
@@ -492,6 +558,14 @@ class MOSProduct(MOSCore):
492
558
  )
493
559
  self.mem_scheduler.submit_messages(messages=[message_item])
494
560
 
561
+ def _filter_memories_by_threshold(
562
+ self, memories: list[TextualMemoryItem], threshold: float = 0.20
563
+ ) -> list[TextualMemoryItem]:
564
+ """
565
+ Filter memories by threshold.
566
+ """
567
+ return [memory for memory in memories if memory.metadata.relativity >= threshold]
568
+
495
569
  def register_mem_cube(
496
570
  self,
497
571
  mem_cube_name_or_path_or_object: str | GeneralMemCube,
@@ -581,9 +655,7 @@ class MOSProduct(MOSCore):
581
655
  user_name = user_id
582
656
 
583
657
  # Create user with configuration using persistent user manager
584
- self.global_user_manager.create_user_with_config(
585
- user_id, user_config, UserRole.USER, user_id
586
- )
658
+ self.user_manager.create_user_with_config(user_id, user_config, UserRole.USER, user_id)
587
659
 
588
660
  # Create user configuration
589
661
  user_config = self._create_user_config(user_id, user_config)
@@ -599,7 +671,7 @@ class MOSProduct(MOSCore):
599
671
  try:
600
672
  default_mem_cube.dump(mem_cube_name_or_path)
601
673
  except Exception as e:
602
- print(e)
674
+ logger.error(f"Failed to dump default cube: {e}")
603
675
 
604
676
  # Register the default cube with MOS
605
677
  self.register_mem_cube(
@@ -656,7 +728,7 @@ class MOSProduct(MOSCore):
656
728
  you should generate some suggestion query, the query should be user what to query,
657
729
  user recently memories is:
658
730
  {memories}
659
- please generate 3 suggestion query in English,
731
+ if the user recently memories is empty, please generate 3 suggestion query in English,
660
732
  output should be a json format, the key is "query", the value is a list of suggestion query.
661
733
 
662
734
  example:
@@ -664,76 +736,27 @@ class MOSProduct(MOSCore):
664
736
  "query": ["query1", "query2", "query3"]
665
737
  }}
666
738
  """
667
- text_mem_result = super().search("my recently memories", user_id=user_id, top_k=10)[
739
+ text_mem_result = super().search("my recently memories", user_id=user_id, top_k=3)[
668
740
  "text_mem"
669
741
  ]
670
742
  if text_mem_result:
671
- memories = "\n".join([m.memory for m in text_mem_result[0]["memories"]])
743
+ memories = "\n".join([m.memory[:200] for m in text_mem_result[0]["memories"]])
672
744
  else:
673
745
  memories = ""
674
746
  message_list = [{"role": "system", "content": suggestion_prompt.format(memories=memories)}]
675
747
  response = self.chat_llm.generate(message_list)
676
- response_json = json.loads(response)
677
-
748
+ clean_response = clean_json_response(response)
749
+ response_json = json.loads(clean_response)
678
750
  return response_json["query"]
679
751
 
680
- def chat(
681
- self,
682
- query: str,
683
- user_id: str,
684
- cube_id: str | None = None,
685
- history: MessageList | None = None,
686
- ) -> Generator[str, None, None]:
687
- """Chat with LLM SSE Type.
688
- Args:
689
- query (str): Query string.
690
- user_id (str): User ID.
691
- cube_id (str, optional): Custom cube ID for user.
692
- history (list[dict], optional): Chat history.
693
-
694
- Returns:
695
- Generator[str, None, None]: The response string generator.
696
- """
697
- # Use MOSCore's built-in validation
698
- if cube_id:
699
- self._validate_cube_access(user_id, cube_id)
700
- else:
701
- self._validate_user_exists(user_id)
702
-
703
- # Load user cubes if not already loaded
704
- self._load_user_cubes(user_id, self.default_cube_config)
705
- time_start = time.time()
706
- memories_list = super().search(query, user_id)["text_mem"]
707
- # Get response from parent MOSCore (returns string, not generator)
708
- response = super().chat(query, user_id)
709
- time_end = time.time()
710
-
711
- # Use tiktoken for proper token-based chunking
712
- for chunk in self._chunk_response_with_tiktoken(response, chunk_size=5):
713
- chunk_data = f"data: {json.dumps({'type': 'text', 'content': chunk})}\n\n"
714
- yield chunk_data
715
-
716
- # Prepare reference data
717
- reference = []
718
- for memories in memories_list:
719
- memories_json = memories.model_dump()
720
- memories_json["metadata"]["ref_id"] = f"[{memories.id.split('-')[0]}]"
721
- memories_json["metadata"]["embedding"] = []
722
- memories_json["metadata"]["sources"] = []
723
- reference.append(memories_json)
724
-
725
- yield f"data: {json.dumps({'type': 'reference', 'content': reference})}\n\n"
726
- total_time = round(float(time_end - time_start), 1)
727
-
728
- yield f"data: {json.dumps({'type': 'time', 'content': {'total_time': total_time, 'speed_improvement': '23%'}})}\n\n"
729
- yield f"data: {json.dumps({'type': 'end'})}\n\n"
730
-
731
752
  def chat_with_references(
732
753
  self,
733
754
  query: str,
734
755
  user_id: str,
735
756
  cube_id: str | None = None,
736
757
  history: MessageList | None = None,
758
+ top_k: int = 10,
759
+ internet_search: bool = False,
737
760
  ) -> Generator[str, None, None]:
738
761
  """
739
762
  Chat with LLM with memory references and streaming output.
@@ -749,30 +772,46 @@ class MOSProduct(MOSCore):
749
772
  """
750
773
 
751
774
  self._load_user_cubes(user_id, self.default_cube_config)
752
-
753
775
  time_start = time.time()
754
776
  memories_list = []
777
+ yield f"data: {json.dumps({'type': 'status', 'data': '0'})}\n\n"
755
778
  memories_result = super().search(
756
- query, user_id, install_cube_ids=[cube_id] if cube_id else None, top_k=10
779
+ query,
780
+ user_id,
781
+ install_cube_ids=[cube_id] if cube_id else None,
782
+ top_k=top_k,
783
+ mode="fine",
784
+ internet_search=internet_search,
757
785
  )["text_mem"]
786
+ yield f"data: {json.dumps({'type': 'status', 'data': '1'})}\n\n"
787
+ search_time_end = time.time()
788
+ logger.info(
789
+ f"time chat: search text_mem time user_id: {user_id} time is: {search_time_end - time_start}"
790
+ )
791
+ self._send_message_to_scheduler(
792
+ user_id=user_id, mem_cube_id=cube_id, query=query, label=QUERY_LABEL
793
+ )
758
794
  if memories_result:
759
795
  memories_list = memories_result[0]["memories"]
760
-
761
- # Build custom system prompt with relevant memories
762
- system_prompt = self._build_system_prompt(user_id, memories_list)
763
-
796
+ memories_list = self._filter_memories_by_threshold(memories_list)
797
+ # Build custom system prompt with relevant memories)
798
+ system_prompt = self._build_enhance_system_prompt(user_id, memories_list)
764
799
  # Get chat history
765
- target_user_id = user_id if user_id is not None else self.user_id
766
- if target_user_id not in self.chat_history_manager:
767
- self._register_chat_history(target_user_id)
800
+ if user_id not in self.chat_history_manager:
801
+ self._register_chat_history(user_id)
768
802
 
769
- chat_history = self.chat_history_manager[target_user_id]
803
+ chat_history = self.chat_history_manager[user_id]
804
+ if history:
805
+ chat_history.chat_history = history[-10:]
770
806
  current_messages = [
771
807
  {"role": "system", "content": system_prompt},
772
808
  *chat_history.chat_history,
773
809
  {"role": "user", "content": query},
774
810
  ]
775
-
811
+ logger.info(
812
+ f"user_id: {user_id}, cube_id: {cube_id}, current_system_prompt: {system_prompt}"
813
+ )
814
+ yield f"data: {json.dumps({'type': 'status', 'data': '2'})}\n\n"
776
815
  # Generate response with custom prompt
777
816
  past_key_values = None
778
817
  response_stream = None
@@ -802,13 +841,16 @@ class MOSProduct(MOSCore):
802
841
  response_stream = self.chat_llm.generate(current_messages)
803
842
 
804
843
  time_end = time.time()
805
-
844
+ chat_time_end = time.time()
845
+ logger.info(
846
+ f"time chat: chat time user_id: {user_id} time is: {chat_time_end - search_time_end}"
847
+ )
806
848
  # Simulate streaming output with proper reference handling using tiktoken
807
849
 
808
850
  # Initialize buffer for streaming
809
851
  buffer = ""
810
852
  full_response = ""
811
-
853
+ token_count = 0
812
854
  # Use tiktoken for proper token-based chunking
813
855
  if self.config.chat_model.backend not in ["huggingface", "vllm"]:
814
856
  # For non-huggingface backends, we need to collect the full response first
@@ -821,6 +863,7 @@ class MOSProduct(MOSCore):
821
863
  for chunk in response_stream:
822
864
  if chunk in ["<think>", "</think>"]:
823
865
  continue
866
+ token_count += 1
824
867
  buffer += chunk
825
868
  full_response += chunk
826
869
 
@@ -847,22 +890,60 @@ class MOSProduct(MOSCore):
847
890
  memories_json["metadata"]["embedding"] = []
848
891
  memories_json["metadata"]["sources"] = []
849
892
  memories_json["metadata"]["memory"] = memories.memory
893
+ memories_json["metadata"]["id"] = memories.id
850
894
  reference.append({"metadata": memories_json["metadata"]})
851
895
 
852
896
  yield f"data: {json.dumps({'type': 'reference', 'data': reference})}\n\n"
897
+ # set kvcache improve speed
898
+ speed_improvement = round(float((len(system_prompt) / 2) * 0.0048 + 44.5), 1)
853
899
  total_time = round(float(time_end - time_start), 1)
854
- yield f"data: {json.dumps({'type': 'time', 'data': {'total_time': total_time, 'speed_improvement': '23%'}})}\n\n"
855
- chat_history.chat_history.append({"role": "user", "content": query})
856
- chat_history.chat_history.append({"role": "assistant", "content": full_response})
857
- self._send_message_to_scheduler(
858
- user_id=user_id, mem_cube_id=cube_id, query=query, label=QUERY_LABEL
859
- )
860
- self._send_message_to_scheduler(
861
- user_id=user_id, mem_cube_id=cube_id, query=full_response, label=ANSWER_LABEL
862
- )
863
- self.chat_history_manager[user_id] = chat_history
864
900
 
901
+ yield f"data: {json.dumps({'type': 'time', 'data': {'total_time': total_time, 'speed_improvement': f'{speed_improvement}%'}})}\n\n"
865
902
  yield f"data: {json.dumps({'type': 'end'})}\n\n"
903
+
904
+ logger.info(f"user_id: {user_id}, cube_id: {cube_id}, current_messages: {current_messages}")
905
+ logger.info(f"user_id: {user_id}, cube_id: {cube_id}, full_response: {full_response}")
906
+
907
+ clean_response, extracted_references = self._extract_references_from_response(full_response)
908
+ logger.info(f"Extracted {len(extracted_references)} references from response")
909
+
910
+ # Send chat report if online_bot is available
911
+ try:
912
+ from memos.memos_tools.notification_utils import send_online_bot_notification
913
+
914
+ # Prepare data for online_bot
915
+ chat_data = {
916
+ "query": query,
917
+ "user_id": user_id,
918
+ "cube_id": cube_id,
919
+ "system_prompt": system_prompt,
920
+ "full_response": full_response,
921
+ }
922
+
923
+ system_data = {
924
+ "references": extracted_references,
925
+ "time_start": time_start,
926
+ "time_end": time_end,
927
+ "speed_improvement": speed_improvement,
928
+ }
929
+
930
+ emoji_config = {"chat": "💬", "system_info": "📊"}
931
+
932
+ send_online_bot_notification(
933
+ online_bot=self.online_bot,
934
+ header_name="MemOS Chat Report",
935
+ sub_title_name="chat_with_references",
936
+ title_color="#00956D",
937
+ other_data1=chat_data,
938
+ other_data2=system_data,
939
+ emoji=emoji_config,
940
+ )
941
+ except Exception as e:
942
+ logger.warning(f"Failed to send chat notification: {e}")
943
+
944
+ self._send_message_to_scheduler(
945
+ user_id=user_id, mem_cube_id=cube_id, query=clean_response, label=ANSWER_LABEL
946
+ )
866
947
  self.add(
867
948
  user_id=user_id,
868
949
  messages=[
@@ -873,18 +954,12 @@ class MOSProduct(MOSCore):
873
954
  },
874
955
  {
875
956
  "role": "assistant",
876
- "content": full_response,
957
+ "content": clean_response, # Store clean text without reference markers
877
958
  "chat_time": str(datetime.now().strftime("%Y-%m-%d %H:%M:%S")),
878
959
  },
879
960
  ],
880
961
  mem_cube_id=cube_id,
881
962
  )
882
- # Keep chat history under 30 messages by removing oldest conversation pair
883
- if len(self.chat_history_manager[user_id].chat_history) > 10:
884
- self.chat_history_manager[user_id].chat_history.pop(0) # Remove oldest user message
885
- self.chat_history_manager[user_id].chat_history.pop(
886
- 0
887
- ) # Remove oldest assistant response
888
963
 
889
964
  def get_all(
890
965
  self,
@@ -905,9 +980,14 @@ class MOSProduct(MOSCore):
905
980
 
906
981
  # Load user cubes if not already loaded
907
982
  self._load_user_cubes(user_id, self.default_cube_config)
983
+ time_start = time.time()
908
984
  memory_list = super().get_all(
909
985
  mem_cube_id=mem_cube_ids[0] if mem_cube_ids else None, user_id=user_id
910
986
  )[memory_type]
987
+ get_all_time_end = time.time()
988
+ logger.info(
989
+ f"time get_all: get_all time user_id: {user_id} time is: {get_all_time_end - time_start}"
990
+ )
911
991
  reformat_memory_list = []
912
992
  if memory_type == "text_mem":
913
993
  for memory in memory_list:
@@ -918,8 +998,10 @@ class MOSProduct(MOSCore):
918
998
  "UserMemory": 0.40,
919
999
  }
920
1000
  tree_result, node_type_count = convert_graph_to_tree_forworkmem(
921
- memories, target_node_count=150, type_ratios=custom_type_ratios
1001
+ memories, target_node_count=200, type_ratios=custom_type_ratios
922
1002
  )
1003
+ # Ensure all node IDs are unique in the tree structure
1004
+ tree_result = ensure_unique_tree_ids(tree_result)
923
1005
  memories_filtered = filter_nodes_by_tree_ids(tree_result, memories)
924
1006
  children = tree_result["children"]
925
1007
  children_sort = sort_children_by_memory_type(children)
@@ -963,6 +1045,10 @@ class MOSProduct(MOSCore):
963
1045
  "memories": act_mem_params[0].model_dump(),
964
1046
  }
965
1047
  )
1048
+ make_format_time_end = time.time()
1049
+ logger.info(
1050
+ f"time get_all: make_format time user_id: {user_id} time is: {make_format_time_end - get_all_time_end}"
1051
+ )
966
1052
  return reformat_memory_list
967
1053
 
968
1054
  def _get_subgraph(
@@ -985,6 +1071,7 @@ class MOSProduct(MOSCore):
985
1071
  user_id: str,
986
1072
  query: str,
987
1073
  mem_cube_ids: list[str] | None = None,
1074
+ top_k: int = 20,
988
1075
  ) -> list[dict[str, Any]]:
989
1076
  """Get all memory items for a user.
990
1077
 
@@ -1000,7 +1087,7 @@ class MOSProduct(MOSCore):
1000
1087
  # Load user cubes if not already loaded
1001
1088
  self._load_user_cubes(user_id, self.default_cube_config)
1002
1089
  memory_list = self._get_subgraph(
1003
- query=query, mem_cube_id=mem_cube_ids[0], user_id=user_id, top_k=20
1090
+ query=query, mem_cube_id=mem_cube_ids[0], user_id=user_id, top_k=top_k
1004
1091
  )["text_mem"]
1005
1092
  reformat_memory_list = []
1006
1093
  for memory in memory_list:
@@ -1009,6 +1096,8 @@ class MOSProduct(MOSCore):
1009
1096
  tree_result, node_type_count = convert_graph_to_tree_forworkmem(
1010
1097
  memories, target_node_count=150, type_ratios=custom_type_ratios
1011
1098
  )
1099
+ # Ensure all node IDs are unique in the tree structure
1100
+ tree_result = ensure_unique_tree_ids(tree_result)
1012
1101
  memories_filtered = filter_nodes_by_tree_ids(tree_result, memories)
1013
1102
  children = tree_result["children"]
1014
1103
  children_sort = sort_children_by_memory_type(children)
@@ -1025,15 +1114,27 @@ class MOSProduct(MOSCore):
1025
1114
  return reformat_memory_list
1026
1115
 
1027
1116
  def search(
1028
- self, query: str, user_id: str, install_cube_ids: list[str] | None = None, top_k: int = 20
1117
+ self,
1118
+ query: str,
1119
+ user_id: str,
1120
+ install_cube_ids: list[str] | None = None,
1121
+ top_k: int = 10,
1122
+ mode: Literal["fast", "fine"] = "fast",
1029
1123
  ):
1030
1124
  """Search memories for a specific user."""
1031
- # Validate user access
1032
- self._validate_user_access(user_id)
1033
1125
 
1034
1126
  # Load user cubes if not already loaded
1127
+ time_start = time.time()
1035
1128
  self._load_user_cubes(user_id, self.default_cube_config)
1036
- search_result = super().search(query, user_id, install_cube_ids, top_k)
1129
+ load_user_cubes_time_end = time.time()
1130
+ logger.info(
1131
+ f"time search: load_user_cubes time user_id: {user_id} time is: {load_user_cubes_time_end - time_start}"
1132
+ )
1133
+ search_result = super().search(query, user_id, install_cube_ids, top_k, mode=mode)
1134
+ search_time_end = time.time()
1135
+ logger.info(
1136
+ f"time search: search text_mem time user_id: {user_id} time is: {search_time_end - load_user_cubes_time_end}"
1137
+ )
1037
1138
  text_memory_list = search_result["text_mem"]
1038
1139
  reformat_memory_list = []
1039
1140
  for memory in text_memory_list:
@@ -1049,7 +1150,10 @@ class MOSProduct(MOSCore):
1049
1150
  memories_list.append(memories)
1050
1151
  reformat_memory_list.append({"cube_id": memory["cube_id"], "memories": memories_list})
1051
1152
  search_result["text_mem"] = reformat_memory_list
1052
-
1153
+ time_end = time.time()
1154
+ logger.info(
1155
+ f"time search: total time for user_id: {user_id} time is: {time_end - time_start}"
1156
+ )
1053
1157
  return search_result
1054
1158
 
1055
1159
  def add(
@@ -1059,24 +1163,34 @@ class MOSProduct(MOSCore):
1059
1163
  memory_content: str | None = None,
1060
1164
  doc_path: str | None = None,
1061
1165
  mem_cube_id: str | None = None,
1166
+ source: str | None = None,
1167
+ user_profile: bool = False,
1062
1168
  ):
1063
1169
  """Add memory for a specific user."""
1064
- # Use MOSCore's built-in user/cube validation
1065
- if mem_cube_id:
1066
- self._validate_cube_access(user_id, mem_cube_id)
1067
- else:
1068
- self._validate_user_exists(user_id)
1069
1170
 
1070
1171
  # Load user cubes if not already loaded
1071
1172
  self._load_user_cubes(user_id, self.default_cube_config)
1072
1173
 
1073
1174
  result = super().add(messages, memory_content, doc_path, mem_cube_id, user_id)
1175
+ if user_profile:
1176
+ try:
1177
+ user_interests = memory_content.split("'userInterests': '")[1].split("', '")[0]
1178
+ user_interests = user_interests.replace(",", " ")
1179
+ user_profile_memories = self.mem_cubes[
1180
+ mem_cube_id
1181
+ ].text_mem.internet_retriever.retrieve_from_internet(query=user_interests, top_k=5)
1182
+ for memory in user_profile_memories:
1183
+ self.mem_cubes[mem_cube_id].text_mem.add(memory)
1184
+ except Exception as e:
1185
+ logger.error(
1186
+ f"Failed to retrieve user profile: {e}, memory_content: {memory_content}"
1187
+ )
1074
1188
 
1075
1189
  return result
1076
1190
 
1077
1191
  def list_users(self) -> list:
1078
1192
  """List all registered users."""
1079
- return self.global_user_manager.list_users()
1193
+ return self.user_manager.list_users()
1080
1194
 
1081
1195
  def get_user_info(self, user_id: str) -> dict:
1082
1196
  """Get user information including accessible cubes."""
@@ -1116,7 +1230,7 @@ class MOSProduct(MOSCore):
1116
1230
  """
1117
1231
  try:
1118
1232
  # Save to persistent storage
1119
- success = self.global_user_manager.save_user_config(user_id, config)
1233
+ success = self.user_manager.save_user_config(user_id, config)
1120
1234
  if success:
1121
1235
  # Update in-memory config
1122
1236
  self.user_configs[user_id] = config
@@ -1136,7 +1250,7 @@ class MOSProduct(MOSCore):
1136
1250
  Returns:
1137
1251
  MOSConfig | None: The user's configuration or None if not found.
1138
1252
  """
1139
- return self.global_user_manager.get_user_config(user_id)
1253
+ return self.user_manager.get_user_config(user_id)
1140
1254
 
1141
1255
  def get_active_user_count(self) -> int:
1142
1256
  """Get the number of active user configurations in memory."""