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.
- {memoryos-0.2.0.dist-info → memoryos-0.2.2.dist-info}/METADATA +67 -26
- memoryos-0.2.2.dist-info/RECORD +169 -0
- memoryos-0.2.2.dist-info/entry_points.txt +3 -0
- memos/__init__.py +1 -1
- memos/api/config.py +562 -0
- memos/api/context/context.py +147 -0
- memos/api/context/dependencies.py +90 -0
- memos/api/exceptions.py +28 -0
- memos/api/mcp_serve.py +502 -0
- memos/api/product_api.py +35 -0
- memos/api/product_models.py +163 -0
- memos/api/routers/__init__.py +1 -0
- memos/api/routers/product_router.py +386 -0
- memos/chunkers/sentence_chunker.py +8 -2
- memos/cli.py +113 -0
- memos/configs/embedder.py +27 -0
- memos/configs/graph_db.py +132 -3
- memos/configs/internet_retriever.py +6 -0
- memos/configs/llm.py +47 -0
- memos/configs/mem_cube.py +1 -1
- memos/configs/mem_os.py +5 -0
- memos/configs/mem_reader.py +9 -0
- memos/configs/mem_scheduler.py +107 -7
- memos/configs/mem_user.py +58 -0
- memos/configs/memory.py +5 -4
- memos/dependency.py +52 -0
- memos/embedders/ark.py +92 -0
- memos/embedders/factory.py +4 -0
- memos/embedders/sentence_transformer.py +8 -2
- memos/embedders/universal_api.py +32 -0
- memos/graph_dbs/base.py +11 -3
- memos/graph_dbs/factory.py +4 -0
- memos/graph_dbs/nebular.py +1364 -0
- memos/graph_dbs/neo4j.py +333 -124
- memos/graph_dbs/neo4j_community.py +300 -0
- memos/llms/base.py +9 -0
- memos/llms/deepseek.py +54 -0
- memos/llms/factory.py +10 -1
- memos/llms/hf.py +170 -13
- memos/llms/hf_singleton.py +114 -0
- memos/llms/ollama.py +4 -0
- memos/llms/openai.py +67 -1
- memos/llms/qwen.py +63 -0
- memos/llms/vllm.py +153 -0
- memos/log.py +1 -1
- memos/mem_cube/general.py +77 -16
- memos/mem_cube/utils.py +109 -0
- memos/mem_os/core.py +251 -51
- memos/mem_os/main.py +94 -12
- memos/mem_os/product.py +1220 -43
- memos/mem_os/utils/default_config.py +352 -0
- memos/mem_os/utils/format_utils.py +1401 -0
- memos/mem_reader/simple_struct.py +18 -10
- memos/mem_scheduler/base_scheduler.py +441 -40
- memos/mem_scheduler/general_scheduler.py +249 -248
- memos/mem_scheduler/modules/base.py +14 -5
- memos/mem_scheduler/modules/dispatcher.py +67 -4
- memos/mem_scheduler/modules/misc.py +104 -0
- memos/mem_scheduler/modules/monitor.py +240 -50
- memos/mem_scheduler/modules/rabbitmq_service.py +319 -0
- memos/mem_scheduler/modules/redis_service.py +32 -22
- memos/mem_scheduler/modules/retriever.py +167 -23
- memos/mem_scheduler/modules/scheduler_logger.py +255 -0
- memos/mem_scheduler/mos_for_test_scheduler.py +140 -0
- memos/mem_scheduler/schemas/__init__.py +0 -0
- memos/mem_scheduler/schemas/general_schemas.py +43 -0
- memos/mem_scheduler/{modules/schemas.py → schemas/message_schemas.py} +63 -61
- memos/mem_scheduler/schemas/monitor_schemas.py +329 -0
- memos/mem_scheduler/utils/__init__.py +0 -0
- memos/mem_scheduler/utils/filter_utils.py +176 -0
- memos/mem_scheduler/utils/misc_utils.py +61 -0
- memos/mem_user/factory.py +94 -0
- memos/mem_user/mysql_persistent_user_manager.py +271 -0
- memos/mem_user/mysql_user_manager.py +500 -0
- memos/mem_user/persistent_factory.py +96 -0
- memos/mem_user/persistent_user_manager.py +260 -0
- memos/mem_user/user_manager.py +4 -4
- memos/memories/activation/item.py +29 -0
- memos/memories/activation/kv.py +10 -3
- memos/memories/activation/vllmkv.py +219 -0
- memos/memories/factory.py +2 -0
- memos/memories/textual/base.py +1 -1
- memos/memories/textual/general.py +43 -97
- memos/memories/textual/item.py +5 -33
- memos/memories/textual/tree.py +22 -12
- memos/memories/textual/tree_text_memory/organize/conflict.py +9 -5
- memos/memories/textual/tree_text_memory/organize/manager.py +26 -18
- memos/memories/textual/tree_text_memory/organize/redundancy.py +25 -44
- memos/memories/textual/tree_text_memory/organize/relation_reason_detector.py +50 -48
- memos/memories/textual/tree_text_memory/organize/reorganizer.py +81 -56
- memos/memories/textual/tree_text_memory/retrieve/internet_retriever.py +6 -3
- memos/memories/textual/tree_text_memory/retrieve/internet_retriever_factory.py +2 -0
- memos/memories/textual/tree_text_memory/retrieve/recall.py +0 -1
- memos/memories/textual/tree_text_memory/retrieve/reranker.py +2 -2
- memos/memories/textual/tree_text_memory/retrieve/retrieval_mid_structs.py +2 -0
- memos/memories/textual/tree_text_memory/retrieve/searcher.py +52 -28
- memos/memories/textual/tree_text_memory/retrieve/task_goal_parser.py +42 -15
- memos/memories/textual/tree_text_memory/retrieve/utils.py +11 -7
- memos/memories/textual/tree_text_memory/retrieve/xinyusearch.py +62 -58
- memos/memos_tools/dinding_report_bot.py +422 -0
- memos/memos_tools/notification_service.py +44 -0
- memos/memos_tools/notification_utils.py +96 -0
- memos/parsers/markitdown.py +8 -2
- memos/settings.py +3 -1
- memos/templates/mem_reader_prompts.py +66 -23
- memos/templates/mem_scheduler_prompts.py +126 -43
- memos/templates/mos_prompts.py +87 -0
- memos/templates/tree_reorganize_prompts.py +85 -30
- memos/vec_dbs/base.py +12 -0
- memos/vec_dbs/qdrant.py +46 -20
- memoryos-0.2.0.dist-info/RECORD +0 -128
- memos/mem_scheduler/utils.py +0 -26
- {memoryos-0.2.0.dist-info → memoryos-0.2.2.dist-info}/LICENSE +0 -0
- {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
|
|
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
|
-
|
|
16
|
-
ANSWER_LABEL = "answer"
|
|
15
|
+
logger = get_logger(__name__)
|
|
17
16
|
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
|
|
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
|
-
|
|
44
|
-
|
|
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.
|
|
42
|
+
default_factory=lambda: datetime.utcnow(), description="submit time for schedule_messages"
|
|
56
43
|
)
|
|
57
44
|
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
|
|
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
|
-
|
|
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.
|
|
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
|
+
}
|