MemoryOS 0.2.0__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 (114) hide show
  1. {memoryos-0.2.0.dist-info → memoryos-0.2.2.dist-info}/METADATA +67 -26
  2. memoryos-0.2.2.dist-info/RECORD +169 -0
  3. memoryos-0.2.2.dist-info/entry_points.txt +3 -0
  4. memos/__init__.py +1 -1
  5. memos/api/config.py +562 -0
  6. memos/api/context/context.py +147 -0
  7. memos/api/context/dependencies.py +90 -0
  8. memos/api/exceptions.py +28 -0
  9. memos/api/mcp_serve.py +502 -0
  10. memos/api/product_api.py +35 -0
  11. memos/api/product_models.py +163 -0
  12. memos/api/routers/__init__.py +1 -0
  13. memos/api/routers/product_router.py +386 -0
  14. memos/chunkers/sentence_chunker.py +8 -2
  15. memos/cli.py +113 -0
  16. memos/configs/embedder.py +27 -0
  17. memos/configs/graph_db.py +132 -3
  18. memos/configs/internet_retriever.py +6 -0
  19. memos/configs/llm.py +47 -0
  20. memos/configs/mem_cube.py +1 -1
  21. memos/configs/mem_os.py +5 -0
  22. memos/configs/mem_reader.py +9 -0
  23. memos/configs/mem_scheduler.py +107 -7
  24. memos/configs/mem_user.py +58 -0
  25. memos/configs/memory.py +5 -4
  26. memos/dependency.py +52 -0
  27. memos/embedders/ark.py +92 -0
  28. memos/embedders/factory.py +4 -0
  29. memos/embedders/sentence_transformer.py +8 -2
  30. memos/embedders/universal_api.py +32 -0
  31. memos/graph_dbs/base.py +11 -3
  32. memos/graph_dbs/factory.py +4 -0
  33. memos/graph_dbs/nebular.py +1364 -0
  34. memos/graph_dbs/neo4j.py +333 -124
  35. memos/graph_dbs/neo4j_community.py +300 -0
  36. memos/llms/base.py +9 -0
  37. memos/llms/deepseek.py +54 -0
  38. memos/llms/factory.py +10 -1
  39. memos/llms/hf.py +170 -13
  40. memos/llms/hf_singleton.py +114 -0
  41. memos/llms/ollama.py +4 -0
  42. memos/llms/openai.py +67 -1
  43. memos/llms/qwen.py +63 -0
  44. memos/llms/vllm.py +153 -0
  45. memos/log.py +1 -1
  46. memos/mem_cube/general.py +77 -16
  47. memos/mem_cube/utils.py +109 -0
  48. memos/mem_os/core.py +251 -51
  49. memos/mem_os/main.py +94 -12
  50. memos/mem_os/product.py +1220 -43
  51. memos/mem_os/utils/default_config.py +352 -0
  52. memos/mem_os/utils/format_utils.py +1401 -0
  53. memos/mem_reader/simple_struct.py +18 -10
  54. memos/mem_scheduler/base_scheduler.py +441 -40
  55. memos/mem_scheduler/general_scheduler.py +249 -248
  56. memos/mem_scheduler/modules/base.py +14 -5
  57. memos/mem_scheduler/modules/dispatcher.py +67 -4
  58. memos/mem_scheduler/modules/misc.py +104 -0
  59. memos/mem_scheduler/modules/monitor.py +240 -50
  60. memos/mem_scheduler/modules/rabbitmq_service.py +319 -0
  61. memos/mem_scheduler/modules/redis_service.py +32 -22
  62. memos/mem_scheduler/modules/retriever.py +167 -23
  63. memos/mem_scheduler/modules/scheduler_logger.py +255 -0
  64. memos/mem_scheduler/mos_for_test_scheduler.py +140 -0
  65. memos/mem_scheduler/schemas/__init__.py +0 -0
  66. memos/mem_scheduler/schemas/general_schemas.py +43 -0
  67. memos/mem_scheduler/{modules/schemas.py → schemas/message_schemas.py} +63 -61
  68. memos/mem_scheduler/schemas/monitor_schemas.py +329 -0
  69. memos/mem_scheduler/utils/__init__.py +0 -0
  70. memos/mem_scheduler/utils/filter_utils.py +176 -0
  71. memos/mem_scheduler/utils/misc_utils.py +61 -0
  72. memos/mem_user/factory.py +94 -0
  73. memos/mem_user/mysql_persistent_user_manager.py +271 -0
  74. memos/mem_user/mysql_user_manager.py +500 -0
  75. memos/mem_user/persistent_factory.py +96 -0
  76. memos/mem_user/persistent_user_manager.py +260 -0
  77. memos/mem_user/user_manager.py +4 -4
  78. memos/memories/activation/item.py +29 -0
  79. memos/memories/activation/kv.py +10 -3
  80. memos/memories/activation/vllmkv.py +219 -0
  81. memos/memories/factory.py +2 -0
  82. memos/memories/textual/base.py +1 -1
  83. memos/memories/textual/general.py +43 -97
  84. memos/memories/textual/item.py +5 -33
  85. memos/memories/textual/tree.py +22 -12
  86. memos/memories/textual/tree_text_memory/organize/conflict.py +9 -5
  87. memos/memories/textual/tree_text_memory/organize/manager.py +26 -18
  88. memos/memories/textual/tree_text_memory/organize/redundancy.py +25 -44
  89. memos/memories/textual/tree_text_memory/organize/relation_reason_detector.py +50 -48
  90. memos/memories/textual/tree_text_memory/organize/reorganizer.py +81 -56
  91. memos/memories/textual/tree_text_memory/retrieve/internet_retriever.py +6 -3
  92. memos/memories/textual/tree_text_memory/retrieve/internet_retriever_factory.py +2 -0
  93. memos/memories/textual/tree_text_memory/retrieve/recall.py +0 -1
  94. memos/memories/textual/tree_text_memory/retrieve/reranker.py +2 -2
  95. memos/memories/textual/tree_text_memory/retrieve/retrieval_mid_structs.py +2 -0
  96. memos/memories/textual/tree_text_memory/retrieve/searcher.py +52 -28
  97. memos/memories/textual/tree_text_memory/retrieve/task_goal_parser.py +42 -15
  98. memos/memories/textual/tree_text_memory/retrieve/utils.py +11 -7
  99. memos/memories/textual/tree_text_memory/retrieve/xinyusearch.py +62 -58
  100. memos/memos_tools/dinding_report_bot.py +422 -0
  101. memos/memos_tools/notification_service.py +44 -0
  102. memos/memos_tools/notification_utils.py +96 -0
  103. memos/parsers/markitdown.py +8 -2
  104. memos/settings.py +3 -1
  105. memos/templates/mem_reader_prompts.py +66 -23
  106. memos/templates/mem_scheduler_prompts.py +126 -43
  107. memos/templates/mos_prompts.py +87 -0
  108. memos/templates/tree_reorganize_prompts.py +85 -30
  109. memos/vec_dbs/base.py +12 -0
  110. memos/vec_dbs/qdrant.py +46 -20
  111. memoryos-0.2.0.dist-info/RECORD +0 -128
  112. memos/mem_scheduler/utils.py +0 -26
  113. {memoryos-0.2.0.dist-info → memoryos-0.2.2.dist-info}/LICENSE +0 -0
  114. {memoryos-0.2.0.dist-info → memoryos-0.2.2.dist-info}/WHEEL +0 -0
@@ -0,0 +1,255 @@
1
+ from collections.abc import Callable
2
+
3
+ from memos.log import get_logger
4
+ from memos.mem_cube.general import GeneralMemCube
5
+ from memos.mem_scheduler.modules.base import BaseSchedulerModule
6
+ from memos.mem_scheduler.schemas.general_schemas import (
7
+ ACTIVATION_MEMORY_TYPE,
8
+ ADD_LABEL,
9
+ LONG_TERM_MEMORY_TYPE,
10
+ NOT_INITIALIZED,
11
+ PARAMETER_MEMORY_TYPE,
12
+ QUERY_LABEL,
13
+ TEXT_MEMORY_TYPE,
14
+ USER_INPUT_TYPE,
15
+ WORKING_MEMORY_TYPE,
16
+ )
17
+ from memos.mem_scheduler.schemas.message_schemas import (
18
+ ScheduleLogForWebItem,
19
+ ScheduleMessageItem,
20
+ )
21
+ from memos.mem_scheduler.utils.filter_utils import (
22
+ transform_name_to_key,
23
+ )
24
+ from memos.mem_scheduler.utils.misc_utils import log_exceptions
25
+ from memos.memories.textual.tree import TextualMemoryItem, TreeTextMemory
26
+
27
+
28
+ logger = get_logger(__name__)
29
+
30
+
31
+ class SchedulerLoggerModule(BaseSchedulerModule):
32
+ def __init__(self):
33
+ """
34
+ Initialize RabbitMQ connection settings.
35
+ """
36
+ super().__init__()
37
+
38
+ @log_exceptions(logger=logger)
39
+ def create_autofilled_log_item(
40
+ self,
41
+ log_content: str,
42
+ label: str,
43
+ from_memory_type: str,
44
+ to_memory_type: str,
45
+ user_id: str,
46
+ mem_cube_id: str,
47
+ mem_cube: GeneralMemCube,
48
+ ) -> ScheduleLogForWebItem:
49
+ text_mem_base: TreeTextMemory = mem_cube.text_mem
50
+ current_memory_sizes = text_mem_base.get_current_memory_size()
51
+ current_memory_sizes = {
52
+ "long_term_memory_size": current_memory_sizes.get("LongTermMemory", 0),
53
+ "user_memory_size": current_memory_sizes.get("UserMemory", 0),
54
+ "working_memory_size": current_memory_sizes.get("WorkingMemory", 0),
55
+ "transformed_act_memory_size": NOT_INITIALIZED,
56
+ "parameter_memory_size": NOT_INITIALIZED,
57
+ }
58
+ memory_capacities = {
59
+ "long_term_memory_capacity": text_mem_base.memory_manager.memory_size["LongTermMemory"],
60
+ "user_memory_capacity": text_mem_base.memory_manager.memory_size["UserMemory"],
61
+ "working_memory_capacity": text_mem_base.memory_manager.memory_size["WorkingMemory"],
62
+ "transformed_act_memory_capacity": NOT_INITIALIZED,
63
+ "parameter_memory_capacity": NOT_INITIALIZED,
64
+ }
65
+
66
+ if hasattr(self, "monitor"):
67
+ if (
68
+ user_id in self.monitor.activation_memory_monitors
69
+ and mem_cube_id in self.monitor.activation_memory_monitors[user_id]
70
+ ):
71
+ activation_monitor = self.monitor.activation_memory_monitors[user_id][mem_cube_id]
72
+ transformed_act_memory_size = len(activation_monitor.memories)
73
+ logger.info(
74
+ f'activation_memory_monitors currently has "{transformed_act_memory_size}" transformed memory size'
75
+ )
76
+ else:
77
+ transformed_act_memory_size = 0
78
+ logger.info(
79
+ f'activation_memory_monitors is not initialized for user "{user_id}" and mem_cube "{mem_cube_id}'
80
+ )
81
+ current_memory_sizes["transformed_act_memory_size"] = transformed_act_memory_size
82
+ current_memory_sizes["parameter_memory_size"] = 1
83
+
84
+ memory_capacities["transformed_act_memory_capacity"] = (
85
+ self.monitor.activation_mem_monitor_capacity
86
+ )
87
+ memory_capacities["parameter_memory_capacity"] = 1
88
+
89
+ log_message = ScheduleLogForWebItem(
90
+ user_id=user_id,
91
+ mem_cube_id=mem_cube_id,
92
+ label=label,
93
+ from_memory_type=from_memory_type,
94
+ to_memory_type=to_memory_type,
95
+ log_content=log_content,
96
+ current_memory_sizes=current_memory_sizes,
97
+ memory_capacities=memory_capacities,
98
+ )
99
+ return log_message
100
+
101
+ @log_exceptions(logger=logger)
102
+ def log_working_memory_replacement(
103
+ self,
104
+ original_memory: list[TextualMemoryItem],
105
+ new_memory: list[TextualMemoryItem],
106
+ user_id: str,
107
+ mem_cube_id: str,
108
+ mem_cube: GeneralMemCube,
109
+ log_func_callback: Callable[[list[ScheduleLogForWebItem]], None],
110
+ ):
111
+ """Log changes when working memory is replaced."""
112
+ memory_type_map = {
113
+ transform_name_to_key(name=m.memory): m.metadata.memory_type
114
+ for m in original_memory + new_memory
115
+ }
116
+
117
+ original_text_memories = [m.memory for m in original_memory]
118
+ new_text_memories = [m.memory for m in new_memory]
119
+
120
+ # Convert to sets for efficient difference operations
121
+ original_set = set(original_text_memories)
122
+ new_set = set(new_text_memories)
123
+
124
+ # Identify changes
125
+ added_memories = list(new_set - original_set) # Present in new but not original
126
+
127
+ # recording messages
128
+ for memory in added_memories:
129
+ normalized_mem = transform_name_to_key(name=memory)
130
+ if normalized_mem not in memory_type_map:
131
+ logger.error(f"Memory text not found in type mapping: {memory[:50]}...")
132
+ # Get the memory type from the map, default to LONG_TERM_MEMORY_TYPE if not found
133
+ mem_type = memory_type_map.get(normalized_mem, LONG_TERM_MEMORY_TYPE)
134
+
135
+ if mem_type == WORKING_MEMORY_TYPE:
136
+ logger.warning(f"Memory already in working memory: {memory[:50]}...")
137
+ continue
138
+
139
+ log_message = self.create_autofilled_log_item(
140
+ log_content=memory,
141
+ label=QUERY_LABEL,
142
+ from_memory_type=mem_type,
143
+ to_memory_type=WORKING_MEMORY_TYPE,
144
+ user_id=user_id,
145
+ mem_cube_id=mem_cube_id,
146
+ mem_cube=mem_cube,
147
+ )
148
+ log_func_callback([log_message])
149
+ logger.info(
150
+ f"{len(added_memories)} {LONG_TERM_MEMORY_TYPE} memorie(s) "
151
+ f"transformed to {WORKING_MEMORY_TYPE} memories."
152
+ )
153
+
154
+ @log_exceptions(logger=logger)
155
+ def log_activation_memory_update(
156
+ self,
157
+ original_text_memories: list[str],
158
+ new_text_memories: list[str],
159
+ label: str,
160
+ user_id: str,
161
+ mem_cube_id: str,
162
+ mem_cube: GeneralMemCube,
163
+ log_func_callback: Callable[[list[ScheduleLogForWebItem]], None],
164
+ ):
165
+ """Log changes when activation memory is updated."""
166
+ original_set = set(original_text_memories)
167
+ new_set = set(new_text_memories)
168
+
169
+ # Identify changes
170
+ added_memories = list(new_set - original_set) # Present in new but not original
171
+
172
+ # recording messages
173
+ for mem in added_memories:
174
+ log_message_a = self.create_autofilled_log_item(
175
+ log_content=mem,
176
+ label=label,
177
+ from_memory_type=TEXT_MEMORY_TYPE,
178
+ to_memory_type=ACTIVATION_MEMORY_TYPE,
179
+ user_id=user_id,
180
+ mem_cube_id=mem_cube_id,
181
+ mem_cube=mem_cube,
182
+ )
183
+ log_message_b = self.create_autofilled_log_item(
184
+ log_content=mem,
185
+ label=label,
186
+ from_memory_type=ACTIVATION_MEMORY_TYPE,
187
+ to_memory_type=PARAMETER_MEMORY_TYPE,
188
+ user_id=user_id,
189
+ mem_cube_id=mem_cube_id,
190
+ mem_cube=mem_cube,
191
+ )
192
+ log_func_callback([log_message_a, log_message_b])
193
+ logger.info(
194
+ f"{len(added_memories)} {LONG_TERM_MEMORY_TYPE} memorie(s) "
195
+ f"transformed to {WORKING_MEMORY_TYPE} memories."
196
+ )
197
+
198
+ @log_exceptions(logger=logger)
199
+ def log_adding_memory(
200
+ self,
201
+ memory: str,
202
+ memory_type: str,
203
+ user_id: str,
204
+ mem_cube_id: str,
205
+ mem_cube: GeneralMemCube,
206
+ log_func_callback: Callable[[list[ScheduleLogForWebItem]], None],
207
+ ):
208
+ """Log changes when working memory is replaced."""
209
+ log_message = self.create_autofilled_log_item(
210
+ log_content=memory,
211
+ label=ADD_LABEL,
212
+ from_memory_type=USER_INPUT_TYPE,
213
+ to_memory_type=memory_type,
214
+ user_id=user_id,
215
+ mem_cube_id=mem_cube_id,
216
+ mem_cube=mem_cube,
217
+ )
218
+ log_func_callback([log_message])
219
+ logger.info(
220
+ f"{USER_INPUT_TYPE} memory for user {user_id} "
221
+ f"converted to {memory_type} memory in mem_cube {mem_cube_id}: {memory}"
222
+ )
223
+
224
+ @log_exceptions(logger=logger)
225
+ def validate_schedule_message(self, message: ScheduleMessageItem, label: str):
226
+ """Validate if the message matches the expected label.
227
+
228
+ Args:
229
+ message: Incoming message item to validate.
230
+ label: Expected message label (e.g., QUERY_LABEL/ANSWER_LABEL).
231
+
232
+ Returns:
233
+ bool: True if validation passed, False otherwise.
234
+ """
235
+ if message.label != label:
236
+ logger.error(f"Handler validation failed: expected={label}, actual={message.label}")
237
+ return False
238
+ return True
239
+
240
+ @log_exceptions(logger=logger)
241
+ def validate_schedule_messages(self, messages: list[ScheduleMessageItem], label: str):
242
+ """Validate if all messages match the expected label.
243
+
244
+ Args:
245
+ messages: List of message items to validate.
246
+ label: Expected message label (e.g., QUERY_LABEL/ANSWER_LABEL).
247
+
248
+ Returns:
249
+ bool: True if all messages passed validation, False if any failed.
250
+ """
251
+ for message in messages:
252
+ if not self.validate_schedule_message(message, label):
253
+ logger.error("Message batch contains invalid labels, aborting processing")
254
+ return False
255
+ return True
@@ -0,0 +1,140 @@
1
+ from datetime import datetime
2
+
3
+ from memos.configs.mem_os import MOSConfig
4
+ from memos.log import get_logger
5
+ from memos.mem_os.main import MOS
6
+ from memos.mem_scheduler.schemas.general_schemas import (
7
+ ANSWER_LABEL,
8
+ MONITOR_WORKING_MEMORY_TYPE,
9
+ QUERY_LABEL,
10
+ )
11
+ from memos.mem_scheduler.schemas.message_schemas import ScheduleMessageItem
12
+
13
+
14
+ logger = get_logger(__name__)
15
+
16
+
17
+ class MOSForTestScheduler(MOS):
18
+ """This class is only to test abilities of mem scheduler"""
19
+
20
+ def __init__(self, config: MOSConfig):
21
+ super().__init__(config)
22
+
23
+ def _str_memories(self, memories: list[str]) -> str:
24
+ """Format memories for display."""
25
+ if not memories:
26
+ return "No memories."
27
+ return "\n".join(f"{i + 1}. {memory}" for i, memory in enumerate(memories))
28
+
29
+ def chat(self, query: str, user_id: str | None = None) -> str:
30
+ """
31
+ Chat with the MOS.
32
+
33
+ Args:
34
+ query (str): The user's query.
35
+
36
+ Returns:
37
+ str: The response from the MOS.
38
+ """
39
+ target_user_id = user_id if user_id is not None else self.user_id
40
+ accessible_cubes = self.user_manager.get_user_cubes(target_user_id)
41
+ user_cube_ids = [cube.cube_id for cube in accessible_cubes]
42
+ if target_user_id not in self.chat_history_manager:
43
+ self._register_chat_history(target_user_id)
44
+
45
+ chat_history = self.chat_history_manager[target_user_id]
46
+
47
+ topk_for_scheduler = 2
48
+
49
+ if self.config.enable_textual_memory and self.mem_cubes:
50
+ memories_all = []
51
+ for mem_cube_id, mem_cube in self.mem_cubes.items():
52
+ if mem_cube_id not in user_cube_ids:
53
+ continue
54
+ if not mem_cube.text_mem:
55
+ continue
56
+
57
+ message_item = ScheduleMessageItem(
58
+ user_id=target_user_id,
59
+ mem_cube_id=mem_cube_id,
60
+ mem_cube=mem_cube,
61
+ label=QUERY_LABEL,
62
+ content=query,
63
+ timestamp=datetime.now(),
64
+ )
65
+ cur_working_memories = [m.memory for m in mem_cube.text_mem.get_working_memory()]
66
+ print(f"Working memories before schedule: {cur_working_memories}")
67
+
68
+ # --- force to run mem_scheduler ---
69
+ self.mem_scheduler.monitor.query_trigger_interval = 0
70
+ self.mem_scheduler._query_message_consumer(messages=[message_item])
71
+
72
+ # from scheduler
73
+ scheduler_memories = self.mem_scheduler.monitor.get_monitor_memories(
74
+ user_id=target_user_id,
75
+ mem_cube_id=mem_cube_id,
76
+ memory_type=MONITOR_WORKING_MEMORY_TYPE,
77
+ top_k=topk_for_scheduler,
78
+ )
79
+ print(f"Working memories after schedule: {scheduler_memories}")
80
+ memories_all.extend(scheduler_memories)
81
+
82
+ # from mem_cube
83
+ memories = mem_cube.text_mem.search(
84
+ query, top_k=self.config.top_k - topk_for_scheduler
85
+ )
86
+ text_memories = [m.memory for m in memories]
87
+ print(f"Search results with new working memories: {text_memories}")
88
+ memories_all.extend(text_memories)
89
+
90
+ memories_all = list(set(memories_all))
91
+
92
+ logger.info(f"🧠 [Memory] Searched memories:\n{self._str_memories(memories_all)}\n")
93
+ system_prompt = self._build_system_prompt(memories_all)
94
+ else:
95
+ system_prompt = self._build_system_prompt()
96
+ current_messages = [
97
+ {"role": "system", "content": system_prompt},
98
+ *chat_history.chat_history,
99
+ {"role": "user", "content": query},
100
+ ]
101
+ past_key_values = None
102
+
103
+ if self.config.enable_activation_memory:
104
+ assert self.config.chat_model.backend == "huggingface", (
105
+ "Activation memory only used for huggingface backend."
106
+ )
107
+ # TODO this only one cubes
108
+ for mem_cube_id, mem_cube in self.mem_cubes.items():
109
+ if mem_cube_id not in user_cube_ids:
110
+ continue
111
+ if mem_cube.act_mem:
112
+ kv_cache = next(iter(mem_cube.act_mem.get_all()), None)
113
+ past_key_values = (
114
+ kv_cache.memory if (kv_cache and hasattr(kv_cache, "memory")) else None
115
+ )
116
+ break
117
+ # Generate response
118
+ response = self.chat_llm.generate(current_messages, past_key_values=past_key_values)
119
+ else:
120
+ response = self.chat_llm.generate(current_messages)
121
+ logger.info(f"🤖 [Assistant] {response}\n")
122
+ chat_history.chat_history.append({"role": "user", "content": query})
123
+ chat_history.chat_history.append({"role": "assistant", "content": response})
124
+ self.chat_history_manager[user_id] = chat_history
125
+
126
+ # submit message to scheduler
127
+ for accessible_mem_cube in accessible_cubes:
128
+ mem_cube_id = accessible_mem_cube.cube_id
129
+ mem_cube = self.mem_cubes[mem_cube_id]
130
+ if self.enable_mem_scheduler and self.mem_scheduler is not None:
131
+ message_item = ScheduleMessageItem(
132
+ user_id=target_user_id,
133
+ mem_cube_id=mem_cube_id,
134
+ mem_cube=mem_cube,
135
+ label=ANSWER_LABEL,
136
+ content=response,
137
+ timestamp=datetime.now(),
138
+ )
139
+ self.mem_scheduler.submit_messages(messages=[message_item])
140
+ return response
File without changes
@@ -0,0 +1,43 @@
1
+ from pathlib import Path
2
+ from typing import NewType
3
+
4
+
5
+ FILE_PATH = Path(__file__).absolute()
6
+ BASE_DIR = FILE_PATH.parent.parent.parent.parent.parent
7
+
8
+ QUERY_LABEL = "query"
9
+ ANSWER_LABEL = "answer"
10
+ ADD_LABEL = "add"
11
+
12
+ TreeTextMemory_SEARCH_METHOD = "tree_text_memory_search"
13
+ TextMemory_SEARCH_METHOD = "text_memory_search"
14
+ DIRECT_EXCHANGE_TYPE = "direct"
15
+ FANOUT_EXCHANGE_TYPE = "fanout"
16
+ DEFAULT_WORKING_MEM_MONITOR_SIZE_LIMIT = 20
17
+ DEFAULT_ACTIVATION_MEM_MONITOR_SIZE_LIMIT = 5
18
+ DEFAULT_ACT_MEM_DUMP_PATH = f"{BASE_DIR}/outputs/mem_scheduler/mem_cube_scheduler_test.kv_cache"
19
+ DEFAULT_THREAD__POOL_MAX_WORKERS = 5
20
+ DEFAULT_CONSUME_INTERVAL_SECONDS = 3
21
+ NOT_INITIALIZED = -1
22
+
23
+
24
+ # web log
25
+ LONG_TERM_MEMORY_TYPE = "LongTermMemory"
26
+ USER_MEMORY_TYPE = "UserMemory"
27
+ WORKING_MEMORY_TYPE = "WorkingMemory"
28
+ TEXT_MEMORY_TYPE = "TextMemory"
29
+ ACTIVATION_MEMORY_TYPE = "ActivationMemory"
30
+ PARAMETER_MEMORY_TYPE = "ParameterMemory"
31
+ USER_INPUT_TYPE = "UserInput"
32
+ NOT_APPLICABLE_TYPE = "NotApplicable"
33
+
34
+ # monitors
35
+ MONITOR_WORKING_MEMORY_TYPE = "MonitorWorkingMemoryType"
36
+ MONITOR_ACTIVATION_MEMORY_TYPE = "MonitorActivationMemoryType"
37
+ DEFAULT_MAX_QUERY_KEY_WORDS = 1000
38
+ DEFAULT_WEIGHT_VECTOR_FOR_RANKING = [0.9, 0.05, 0.05]
39
+
40
+
41
+ # new types
42
+ UserID = NewType("UserID", str)
43
+ MemCubeID = NewType("CubeID", str)
@@ -1,47 +1,34 @@
1
1
  from datetime import datetime
2
- from pathlib import Path
3
- from typing import ClassVar, TypeVar
2
+ from typing import Any
4
3
  from uuid import uuid4
5
4
 
6
- from pydantic import BaseModel, Field
5
+ from pydantic import BaseModel, ConfigDict, Field, field_serializer
7
6
  from typing_extensions import TypedDict
8
7
 
8
+ from memos.log import get_logger
9
9
  from memos.mem_cube.general import GeneralMemCube
10
+ from memos.mem_scheduler.modules.misc import DictConversionMixin
10
11
 
12
+ from .general_schemas import NOT_INITIALIZED
11
13
 
12
- FILE_PATH = Path(__file__).absolute()
13
- BASE_DIR = FILE_PATH.parent.parent.parent.parent.parent
14
14
 
15
- QUERY_LABEL = "query"
16
- ANSWER_LABEL = "answer"
15
+ logger = get_logger(__name__)
17
16
 
18
- TreeTextMemory_SEARCH_METHOD = "tree_text_memory_search"
19
- TextMemory_SEARCH_METHOD = "text_memory_search"
20
- DEFAULT_ACTIVATION_MEM_SIZE = 5
21
- DEFAULT_ACT_MEM_DUMP_PATH = f"{BASE_DIR}/outputs/mem_scheduler/mem_cube_scheduler_test.kv_cache"
22
- DEFAULT_THREAD__POOL_MAX_WORKERS = 5
23
- DEFAULT_CONSUME_INTERVAL_SECONDS = 3
24
- NOT_INITIALIZED = -1
25
- BaseModelType = TypeVar("T", bound="BaseModel")
26
-
27
-
28
- class DictConversionMixin:
29
- def to_dict(self) -> dict:
30
- """Convert the instance to a dictionary."""
31
- return {
32
- **self.dict(),
33
- "timestamp": self.timestamp.isoformat() if hasattr(self, "timestamp") else None,
34
- }
35
-
36
- @classmethod
37
- def from_dict(cls: type[BaseModelType], data: dict) -> BaseModelType:
38
- """Create an instance from a dictionary."""
39
- if "timestamp" in data:
40
- data["timestamp"] = datetime.fromisoformat(data["timestamp"])
41
- return cls(**data)
17
+ DEFAULT_MEMORY_SIZES = {
18
+ "long_term_memory_size": NOT_INITIALIZED,
19
+ "user_memory_size": NOT_INITIALIZED,
20
+ "working_memory_size": NOT_INITIALIZED,
21
+ "transformed_act_memory_size": NOT_INITIALIZED,
22
+ "parameter_memory_size": NOT_INITIALIZED,
23
+ }
42
24
 
43
- class Config:
44
- json_encoders: ClassVar[dict[type, object]] = {datetime: lambda v: v.isoformat()}
25
+ DEFAULT_MEMORY_CAPACITIES = {
26
+ "long_term_memory_capacity": 10000,
27
+ "user_memory_capacity": 10000,
28
+ "working_memory_capacity": 20,
29
+ "transformed_act_memory_capacity": NOT_INITIALIZED,
30
+ "parameter_memory_capacity": NOT_INITIALIZED,
31
+ }
45
32
 
46
33
 
47
34
  class ScheduleMessageItem(BaseModel, DictConversionMixin):
@@ -52,15 +39,36 @@ class ScheduleMessageItem(BaseModel, DictConversionMixin):
52
39
  mem_cube: GeneralMemCube | str = Field(..., description="memcube for schedule")
53
40
  content: str = Field(..., description="Content of the schedule message")
54
41
  timestamp: datetime = Field(
55
- default_factory=datetime.now, description="submit time for schedule_messages"
42
+ default_factory=lambda: datetime.utcnow(), description="submit time for schedule_messages"
56
43
  )
57
44
 
58
- class Config:
59
- arbitrary_types_allowed = True
60
- json_encoders: ClassVar[dict[type, object]] = {
61
- datetime: lambda v: v.isoformat(),
62
- GeneralMemCube: lambda v: f"<GeneralMemCube:{id(v)}>",
63
- }
45
+ # Pydantic V2 model configuration
46
+ model_config = ConfigDict(
47
+ # Allows arbitrary Python types as model fields without validation
48
+ # Required when using custom types like GeneralMemCube that aren't Pydantic models
49
+ arbitrary_types_allowed=True,
50
+ # Additional metadata for JSON Schema generation
51
+ json_schema_extra={
52
+ # Example payload demonstrating the expected structure and sample values
53
+ # Used for API documentation, testing, and developer reference
54
+ "example": {
55
+ "item_id": "123e4567-e89b-12d3-a456-426614174000", # Sample UUID
56
+ "user_id": "user123", # Example user identifier
57
+ "mem_cube_id": "cube456", # Sample memory cube ID
58
+ "label": "sample_label", # Demonstration label value
59
+ "mem_cube": "obj of GeneralMemCube", # Added mem_cube example
60
+ "content": "sample content", # Example message content
61
+ "timestamp": "2024-07-22T12:00:00Z", # Added timestamp example
62
+ }
63
+ },
64
+ )
65
+
66
+ @field_serializer("mem_cube")
67
+ def serialize_mem_cube(self, cube: GeneralMemCube | str, _info) -> str:
68
+ """Custom serializer for GeneralMemCube objects to string representation"""
69
+ if isinstance(cube, str):
70
+ return cube
71
+ return f"<GeneralMemCube:{id(cube)}>"
64
72
 
65
73
  def to_dict(self) -> dict:
66
74
  """Convert model to dictionary suitable for Redis Stream"""
@@ -68,7 +76,6 @@ class ScheduleMessageItem(BaseModel, DictConversionMixin):
68
76
  "item_id": self.item_id,
69
77
  "user_id": self.user_id,
70
78
  "cube_id": self.mem_cube_id,
71
- "message_id": self.message_id,
72
79
  "label": self.label,
73
80
  "cube": "Not Applicable", # Custom cube serialization
74
81
  "content": self.content,
@@ -82,7 +89,6 @@ class ScheduleMessageItem(BaseModel, DictConversionMixin):
82
89
  item_id=data.get("item_id", str(uuid4())),
83
90
  user_id=data["user_id"],
84
91
  cube_id=data["cube_id"],
85
- message_id=data.get("message_id", str(uuid4())),
86
92
  label=data["label"],
87
93
  cube="Not Applicable", # Custom cube deserialization
88
94
  content=data["content"],
@@ -104,23 +110,6 @@ class MemoryCapacities(TypedDict):
104
110
  transformed_act_memory_capacity: int
105
111
 
106
112
 
107
- DEFAULT_MEMORY_SIZES = {
108
- "long_term_memory_size": NOT_INITIALIZED,
109
- "user_memory_size": NOT_INITIALIZED,
110
- "working_memory_size": NOT_INITIALIZED,
111
- "transformed_act_memory_size": NOT_INITIALIZED,
112
- "parameter_memory_size": NOT_INITIALIZED,
113
- }
114
-
115
- DEFAULT_MEMORY_CAPACITIES = {
116
- "long_term_memory_capacity": 10000,
117
- "user_memory_capacity": 10000,
118
- "working_memory_capacity": 20,
119
- "transformed_act_memory_capacity": NOT_INITIALIZED,
120
- "parameter_memory_capacity": NOT_INITIALIZED,
121
- }
122
-
123
-
124
113
  class ScheduleLogForWebItem(BaseModel, DictConversionMixin):
125
114
  item_id: str = Field(
126
115
  description="Unique identifier for the log entry", default_factory=lambda: str(uuid4())
@@ -130,7 +119,8 @@ class ScheduleLogForWebItem(BaseModel, DictConversionMixin):
130
119
  ..., description="Identifier for the memcube associated with this log entry"
131
120
  )
132
121
  label: str = Field(..., description="Label categorizing the type of log")
133
- log_title: str = Field(..., description="Title or brief summary of the log content")
122
+ from_memory_type: str = Field(..., description="Source memory type")
123
+ to_memory_type: str = Field(..., description="Destination memory type")
134
124
  log_content: str = Field(..., description="Detailed content of the log entry")
135
125
  current_memory_sizes: MemorySizes = Field(
136
126
  default_factory=lambda: dict(DEFAULT_MEMORY_SIZES),
@@ -141,6 +131,18 @@ class ScheduleLogForWebItem(BaseModel, DictConversionMixin):
141
131
  description="Maximum capacities of memory partitions",
142
132
  )
143
133
  timestamp: datetime = Field(
144
- default_factory=datetime.now,
134
+ default_factory=lambda: datetime.utcnow(),
145
135
  description="Timestamp indicating when the log entry was created",
146
136
  )
137
+
138
+ def debug_info(self) -> dict[str, Any]:
139
+ """Return structured debug information for logging purposes."""
140
+ return {
141
+ "log_id": self.item_id,
142
+ "user_id": self.user_id,
143
+ "mem_cube_id": self.mem_cube_id,
144
+ "operation": f"{self.from_memory_type} → {self.to_memory_type}",
145
+ "label": self.label,
146
+ "content_length": len(self.log_content),
147
+ "timestamp": self.timestamp.isoformat(),
148
+ }